




import gc
import unittest

import numpy as np
import torch
from transformers import CLIPTextConfig, CLIPTextModel, CLIPTokenizer

from diffusers import (
    AutoencoderKL,
    DDIMScheduler,
    StableDiffusionSAGPipeline,
    UNetaDConditionModel,
)
from diffusers.utils import slow, torch_device
from diffusers.utils.testing_utils import enable_full_determinism, require_torch_gpu

from ..pipeline_params import TEXT_TO_IMAGE_BATCH_PARAMS, TEXT_TO_IMAGE_IMAGE_PARAMS, TEXT_TO_IMAGE_PARAMS
from ..test_pipelines_common import PipelineLatentTesterMixin, PipelineTesterMixin


enable_full_determinism()




class    __lowerCAmelCase       (							_UpperCamelCase      , _UpperCamelCase      , unittest.TestCase):
       _a 					=						StableDiffusionSAGPipeline
       _a 					=						TEXT_TO_IMAGE_PARAMS
       _a 					=						TEXT_TO_IMAGE_BATCH_PARAMS
       _a 					=						TEXT_TO_IMAGE_IMAGE_PARAMS
       _a 					=						TEXT_TO_IMAGE_IMAGE_PARAMS
       _a 					=						False






       def 					SCREAMING_SNAKE_CASE    (   self:					str				):
         torch.manual_seed(0				)
         lowercase    :List[Any]      	=						UNetaDConditionModel(
             block_out_channels=(32, 64)       ,					layers_per_block=2       ,					sample_size=32       ,					in_channels=4       ,					out_channels=4       ,					down_block_types=("DownBlock2D", "CrossAttnDownBlock2D")       ,					up_block_types=("CrossAttnUpBlock2D", "UpBlock2D")       ,					cross_attention_dim=32       ,					)
         lowercase    :Any      	=						DDIMScheduler(
             beta_start=0.0_00_85       ,					beta_end=0.0_12       ,					beta_schedule="scaled_linear"       ,					clip_sample=_lowerCAmelCase       ,					set_alpha_to_one=_lowerCAmelCase       ,					)
         torch.manual_seed(0				)
         lowercase    :Any      	=						AutoencoderKL(
             block_out_channels=[32, 64]       ,					in_channels=3       ,					out_channels=3       ,					down_block_types=["DownEncoderBlock2D", "DownEncoderBlock2D"]       ,					up_block_types=["UpDecoderBlock2D", "UpDecoderBlock2D"]       ,					latent_channels=4       ,					)
         torch.manual_seed(0				)
         lowercase    :List[Any]      	=						CLIPTextConfig(
             bos_token_id=0       ,					eos_token_id=2       ,					hidden_size=32       ,					intermediate_size=37       ,					layer_norm_eps=1e-0_5       ,					num_attention_heads=4       ,					num_hidden_layers=5       ,					pad_token_id=1       ,					vocab_size=10_00       ,					)
         lowercase    :int      	=						CLIPTextModel(_lowerCAmelCase				)
         lowercase    :Tuple      	=						CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip"				)

         lowercase    :Optional[int]      	=						{
             "unet": unet,
             "scheduler": scheduler,
             "vae": vae,
             "text_encoder": text_encoder,
             "tokenizer": tokenizer,
             "safety_checker": None,
             "feature_extractor": None,
         }
         return components






       def 					SCREAMING_SNAKE_CASE    (   self:					Optional[Any]       ,					_lowerCAmelCase:					Tuple       ,					_lowerCAmelCase:					List[Any]=0				):
         if str(_lowerCAmelCase				).startswith("mps"				):
           lowercase    :int      	=						torch.manual_seed(_lowerCAmelCase				)
         else:
           lowercase    :List[str]      	=						torch.Generator(device=_lowerCAmelCase				).manual_seed(_lowerCAmelCase				)
         lowercase    :Union[str, Any]      	=						{
             "prompt": ".",
             "generator": generator,
             "num_inference_steps": 2,
             "guidance_scale": 1.0,
             "sag_scale": 1.0,
             "output_type": "numpy",
         }
         return inputs






       def 					SCREAMING_SNAKE_CASE    (   self:					Optional[Any]				):
         super().test_inference_batch_single_identical(expected_max_diff=3e-3				)







@slow
@require_torch_gpu
class    __lowerCAmelCase       (							unittest.TestCase):






       def 					SCREAMING_SNAKE_CASE    (   self:					Optional[int]				):
         # clean up the VRAM after each test
         super().tearDown()
         gc.collect()
         torch.cuda.empty_cache()






       def 					SCREAMING_SNAKE_CASE    (   self:					str				):
         lowercase    :str      	=						StableDiffusionSAGPipeline.from_pretrained("CompVis/stable-diffusion-v1-4"				)
         lowercase    :int      	=						sag_pipe.to(_lowerCAmelCase				)
         sag_pipe.set_progress_bar_config(disable=_lowerCAmelCase				)

         lowercase    :Tuple      	=						"."
         lowercase    :Tuple      	=						torch.manual_seed(0				)
         lowercase    :Any      	=						sag_pipe(
             [prompt]       ,					generator=_lowerCAmelCase       ,					guidance_scale=7.5       ,					sag_scale=1.0       ,					num_inference_steps=20       ,					output_type="np"				)

         lowercase    :str      	=						output.images

         lowercase    :Dict      	=						image[0, -3:, -3:, -1]

         assert image.shape == (1, 5_12, 5_12, 3)
         lowercase    :List[Any]      	=						np.array([0.15_68, 0.17_38, 0.16_95, 0.16_93, 0.15_07, 0.17_05, 0.15_47, 0.17_51, 0.19_49]				)

         assert np.abs(image_slice.flatten() - expected_slice				).max() < 5e-2






       def 					SCREAMING_SNAKE_CASE    (   self:					str				):
         lowercase    :Dict      	=						StableDiffusionSAGPipeline.from_pretrained("stabilityai/stable-diffusion-2-1-base"				)
         lowercase    :Tuple      	=						sag_pipe.to(_lowerCAmelCase				)
         sag_pipe.set_progress_bar_config(disable=_lowerCAmelCase				)

         lowercase    :Optional[int]      	=						"."
         lowercase    :int      	=						torch.manual_seed(0				)
         lowercase    :Optional[Any]      	=						sag_pipe(
             [prompt]       ,					generator=_lowerCAmelCase       ,					guidance_scale=7.5       ,					sag_scale=1.0       ,					num_inference_steps=20       ,					output_type="np"				)

         lowercase    :List[str]      	=						output.images

         lowercase    :Union[str, Any]      	=						image[0, -3:, -3:, -1]

         assert image.shape == (1, 5_12, 5_12, 3)
         lowercase    :List[str]      	=						np.array([0.34_59, 0.28_76, 0.25_37, 0.30_02, 0.26_71, 0.21_60, 0.30_26, 0.22_62, 0.23_71]				)

         assert np.abs(image_slice.flatten() - expected_slice				).max() < 5e-2






       def 					SCREAMING_SNAKE_CASE    (   self:					Dict				):
         lowercase    :Optional[int]      	=						StableDiffusionSAGPipeline.from_pretrained("stabilityai/stable-diffusion-2-1-base"				)
         lowercase    :List[Any]      	=						sag_pipe.to(_lowerCAmelCase				)
         sag_pipe.set_progress_bar_config(disable=_lowerCAmelCase				)

         lowercase    :Tuple      	=						"."
         lowercase    :List[str]      	=						torch.manual_seed(0				)
         lowercase    :Tuple      	=						sag_pipe(
             [prompt]       ,					width=7_68       ,					height=5_12       ,					generator=_lowerCAmelCase       ,					guidance_scale=7.5       ,					sag_scale=1.0       ,					num_inference_steps=20       ,					output_type="np"       ,					)

         lowercase    :int      	=						output.images

         assert image.shape == (1, 5_12, 7_68, 3)









from typing import TYPE_CHECKING

from ...utils import OptionalDependencyNotAvailable, _LazyModule, is_tokenizers_available, is_torch_available


_UpperCAmelCase	:						Optional[int]			=     {
    "configuration_biogpt": ["BIOGPT_PRETRAINED_CONFIG_ARCHIVE_MAP", "BioGptConfig"],
    "tokenization_biogpt": ["BioGptTokenizer"],
}

try:
     if not is_torch_available():
          raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
     pass
else:
     _UpperCAmelCase	:						int			=     [
         "BIOGPT_PRETRAINED_MODEL_ARCHIVE_LIST",
         "BioGptForCausalLM",
         "BioGptForTokenClassification",
         "BioGptForSequenceClassification",
         "BioGptModel",
         "BioGptPreTrainedModel",
     ]


if TYPE_CHECKING:
     from .configuration_biogpt import BIOGPT_PRETRAINED_CONFIG_ARCHIVE_MAP, BioGptConfig
     from .tokenization_biogpt import BioGptTokenizer

     try:
          if not is_torch_available():
               raise OptionalDependencyNotAvailable()
     except OptionalDependencyNotAvailable:
          pass
     else:
          from .modeling_biogpt import (
              BIOGPT_PRETRAINED_MODEL_ARCHIVE_LIST,
              BioGptForCausalLM,
              BioGptForSequenceClassification,
              BioGptForTokenClassification,
              BioGptModel,
              BioGptPreTrainedModel,
          )


else:
     import sys

     _UpperCAmelCase	:						Optional[Any]			=     _LazyModule(__name__, globals()["__file__"], _import_structure, module_spec=__spec__)



