
"""simple docstring"""






import gc
import unittest

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

from diffusers import (
    AutoencoderKL,
    DDIMParallelScheduler,
    DDPMParallelScheduler,
    StableDiffusionParadigmsPipeline,
    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 							snake_case       (     SCREAMING_SNAKE_CASE_							,       SCREAMING_SNAKE_CASE_							,       unittest.TestCase		):
   a_					:		List[Any]     					=     StableDiffusionParadigmsPipeline
   a_					:		str     					=     TEXT_TO_IMAGE_PARAMS
   a_					:		List[Any]     					=     TEXT_TO_IMAGE_BATCH_PARAMS
   a_					:		Optional[Any]     					=     TEXT_TO_IMAGE_IMAGE_PARAMS
   a_					:		Any     					=     TEXT_TO_IMAGE_IMAGE_PARAMS







   def    UpperCAmelCase__	( self)   ->Tuple:
      torch.manual_seed(0)
      a_						   =     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       ,							attention_head_dim=(2, 4)       ,							use_linear_projection=__UpperCAmelCase       ,							)
      a_						   =     DDIMParallelScheduler(
          beta_start=0.00_085       ,							beta_end=0.012       ,							beta_schedule="scaled_linear"       ,							clip_sample=__UpperCAmelCase       ,							set_alpha_to_one=__UpperCAmelCase       ,							)
      torch.manual_seed(0)
      a_						   =     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       ,							sample_size=1_28       ,							)
      torch.manual_seed(0)
      a_						   =     CLIPTextConfig(
          bos_token_id=0       ,							eos_token_id=2       ,							hidden_size=32       ,							intermediate_size=37       ,							layer_norm_eps=1E-05       ,							num_attention_heads=4       ,							num_hidden_layers=5       ,							pad_token_id=1       ,							vocab_size=10_00       ,							hidden_act="gelu"       ,							projection_dim=5_12       ,							)
      a_						   =     CLIPTextModel(__UpperCAmelCase)
      a_						   =     CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip")

      a_						   =     {
          "unet": unet,
          "scheduler": scheduler,
          "vae": vae,
          "text_encoder": text_encoder,
          "tokenizer": tokenizer,
          "safety_checker": None,
          "feature_extractor": None,
      }
      return components







   def    UpperCAmelCase__	( self       ,							__UpperCAmelCase       ,							__UpperCAmelCase=0)   ->Dict:
      if str(__UpperCAmelCase).startswith("mps"):
         a_						   =     torch.manual_seed(__UpperCAmelCase)
      else:
         a_						   =     torch.Generator(device=__UpperCAmelCase).manual_seed(__UpperCAmelCase)
      a_						   =     {
          "prompt": "a photograph of an astronaut riding a horse",
          "generator": generator,
          "num_inference_steps": 10,
          "guidance_scale": 6.0,
          "output_type": "numpy",
          "parallel": 3,
          "debug": True,
      }
      return inputs







   def    UpperCAmelCase__	( self)   ->List[Any]:
      a_						   =     "cpu"  # ensure determinism for the device-dependent torch.Generator
      a_						   =     self.get_dummy_components()
      a_						   =     StableDiffusionParadigmsPipeline(**__UpperCAmelCase)
      a_						   =     sd_pipe.to(__UpperCAmelCase)
      sd_pipe.set_progress_bar_config(disable=__UpperCAmelCase)

      a_						   =     self.get_dummy_inputs(__UpperCAmelCase)
      a_						   =     sd_pipe(**__UpperCAmelCase).images
      a_						   =     image[0, -3:, -3:, -1]
      assert image.shape == (1, 64, 64, 3)

      a_						   =     np.array([0.4_773, 0.5_417, 0.4_723, 0.4_925, 0.5_631, 0.4_752, 0.5_240, 0.4_935, 0.5_023])

      assert np.abs(image_slice.flatten() - expected_slice).max() < 1E-2







   def    UpperCAmelCase__	( self)   ->Dict:
      a_						   =     "cpu"  # ensure determinism for the device-dependent torch.Generator
      a_						   =     self.get_dummy_components()
      torch.manual_seed(0)
      a_						   =     DDPMParallelScheduler()
      torch.manual_seed(0)
      a_						   =     StableDiffusionParadigmsPipeline(**__UpperCAmelCase)
      a_						   =     sd_pipe.to(__UpperCAmelCase)
      sd_pipe.set_progress_bar_config(disable=__UpperCAmelCase)

      a_						   =     self.get_dummy_inputs(__UpperCAmelCase)
      a_						   =     sd_pipe(**__UpperCAmelCase).images
      a_						   =     image[0, -3:, -3:, -1]
      assert image.shape == (1, 64, 64, 3)

      a_						   =     np.array([0.3_573, 0.4_420, 0.4_960, 0.4_799, 0.3_796, 0.3_879, 0.4_819, 0.4_365, 0.4_468])

      assert np.abs(image_slice.flatten() - expected_slice).max() < 1E-2







   def    UpperCAmelCase__	( self)   ->Any:
      super().test_inference_batch_consistent(batch_sizes=[1, 2])







   def    UpperCAmelCase__	( self)   ->Tuple:
      super().test_inference_batch_single_identical(batch_size=2       ,							expected_max_diff=3E-3)







   def    UpperCAmelCase__	( self)   ->int:
      a_						   =     "cpu"  # ensure determinism for the device-dependent torch.Generator
      a_						   =     self.get_dummy_components()
      a_						   =     StableDiffusionParadigmsPipeline(**__UpperCAmelCase)
      a_						   =     sd_pipe.to(__UpperCAmelCase)
      sd_pipe.set_progress_bar_config(disable=__UpperCAmelCase)

      a_						   =     self.get_dummy_inputs(__UpperCAmelCase)
      a_						   =     "french fries"
      a_						   =     sd_pipe(**__UpperCAmelCase       ,							negative_prompt=__UpperCAmelCase)
      a_						   =     output.images
      a_						   =     image[0, -3:, -3:, -1]

      assert image.shape == (1, 64, 64, 3)

      a_						   =     np.array([0.4_771, 0.5_420, 0.4_683, 0.4_918, 0.5_636, 0.4_725, 0.5_230, 0.4_923, 0.5_015])

      assert np.abs(image_slice.flatten() - expected_slice).max() < 1E-2






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







   def    UpperCAmelCase__	( self)   ->Optional[Any]:
      super().tearDown()
      gc.collect()
      torch.cuda.empty_cache()







   def    UpperCAmelCase__	( self       ,							__UpperCAmelCase=0)   ->Tuple:
      a_						   =     torch.Generator(device=__UpperCAmelCase).manual_seed(__UpperCAmelCase)
      a_						   =     {
          "prompt": "a photograph of an astronaut riding a horse",
          "generator": generator,
          "num_inference_steps": 10,
          "guidance_scale": 7.5,
          "output_type": "numpy",
          "parallel": 3,
          "debug": True,
      }
      return inputs







   def    UpperCAmelCase__	( self)   ->Any:
      a_						   =     "stabilityai/stable-diffusion-2-base"
      a_						   =     DDIMParallelScheduler.from_pretrained(__UpperCAmelCase       ,							subfolder="scheduler")
      a_						   =     StableDiffusionParadigmsPipeline.from_pretrained(__UpperCAmelCase       ,							scheduler=__UpperCAmelCase       ,							safety_checker=__UpperCAmelCase)
      pipe.to(__UpperCAmelCase)
      pipe.set_progress_bar_config(disable=__UpperCAmelCase)
      pipe.enable_attention_slicing()

      a_						   =     self.get_inputs()
      a_						   =     pipe(**__UpperCAmelCase).images
      a_						   =     image[0, -3:, -3:, -1].flatten()

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

      a_						   =     np.array([0.9_622, 0.9_602, 0.9_748, 0.9_591, 0.9_630, 0.9_691, 0.9_661, 0.9_631, 0.9_741])

      assert np.abs(expected_slice - image_slice).max() < 1E-2

"""simple docstring"""






import copy

from transformers.configuration_utils import PretrainedConfig
from transformers.utils import logging


UpperCamelCase_						   =   logging.get_logger(__name__)





class 							snake_case       (     SCREAMING_SNAKE_CASE_		):
   a_					:		Tuple     					=     """hybrid-clip"""
   a_					:		Optional[Any]     					=     True







   def __init__( self       ,							__UpperCAmelCase=5_12       ,							**__UpperCAmelCase)   ->Dict:
      super().__init__(**__UpperCAmelCase)

      if "text_config" not in kwargs:
         raise ValueError("`text_config` can not be `None`.")

      if "vision_config" not in kwargs:
         raise ValueError("`vision_config` can not be `None`.")

      a_						   =     kwargs.pop("text_config")
      a_						   =     kwargs.pop("vision_config")

      a_						   =     text_config.pop("model_type")
      a_						   =     vision_config.pop("model_type")

      from transformers import AutoConfig

      a_						   =     AutoConfig.for_model(__UpperCAmelCase       ,							**__UpperCAmelCase)

      if vision_model_type == "clip":
         a_						   =     AutoConfig.for_model(__UpperCAmelCase       ,							**__UpperCAmelCase).vision_config
      elif vision_model_type == "clip_vision_model":
         from transformers import CLIPVisionConfig

         a_						   =     CLIPVisionConfig(**__UpperCAmelCase)
      else:
         a_						   =     AutoConfig.for_model(__UpperCAmelCase       ,							**__UpperCAmelCase)

      a_						   =     projection_dim
      a_						   =     1.0







   @classmethod
   def    UpperCAmelCase__	( cls       ,							__UpperCAmelCase       ,							__UpperCAmelCase       ,							**__UpperCAmelCase)   ->Optional[int]:
      return cls(text_config=text_config.to_dict()       ,							vision_config=vision_config.to_dict()       ,							**__UpperCAmelCase)







   def    UpperCAmelCase__	( self)   ->str:
      a_						   =     copy.deepcopy(self.__dict__)
      a_						   =     self.text_config.to_dict()
      a_						   =     self.vision_config.to_dict()
      a_						   =     self.__class__.model_type
      return output