
'''simple docstring'''
import gc
import unittest

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

from diffusers import PriorTransformer, UnCLIPPipeline, UnCLIPScheduler, UNetaDConditionModel, UNetaDModel
from diffusers.pipelines.unclip.text_proj import UnCLIPTextProjModel
from diffusers.utils import load_numpy, nightly, slow, torch_device
from diffusers.utils.testing_utils import enable_full_determinism, require_torch_gpu, skip_mps

from ..pipeline_params import TEXT_TO_IMAGE_BATCH_PARAMS, TEXT_TO_IMAGE_PARAMS
from ..test_pipelines_common import PipelineTesterMixin, assert_mean_pixel_difference


enable_full_determinism()






class 				UpperCAmelCase_			(		snake_case__      ,  unittest.TestCase   ):
 UpperCAmelCase_						      =  UnCLIPPipeline
 UpperCAmelCase_						      =  TEXT_TO_IMAGE_PARAMS - {
     """negative_prompt""",
     """height""",
     """width""",
     """negative_prompt_embeds""",
     """guidance_scale""",
     """prompt_embeds""",
     """cross_attention_kwargs""",
 }
 UpperCAmelCase_						      =  TEXT_TO_IMAGE_BATCH_PARAMS
 UpperCAmelCase_						      =  [
     """generator""",
     """return_dict""",
     """prior_num_inference_steps""",
     """decoder_num_inference_steps""",
     """super_res_num_inference_steps""",
 ]
 UpperCAmelCase_						      =  False





 @property
 def     snake_case__						(      self):
     return 32





 @property
 def     snake_case__						(      self):
     return 32





 @property
 def     snake_case__						(      self):
     return self.time_input_dim





 @property
 def     snake_case__						(      self):
     return self.time_input_dim * 4





 @property
 def     snake_case__						(      self):
     return 1_00





 @property
 def     snake_case__						(      self):
     snake_case_	:      List[Any]            =  CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip")
     return tokenizer





 @property
 def     snake_case__						(      self):
     torch.manual_seed(0)
     snake_case_	:      Dict            =  CLIPTextConfig(
         bos_token_id=0	,    eos_token_id=2	,    hidden_size=self.text_embedder_hidden_size	,    projection_dim=self.text_embedder_hidden_size	,    intermediate_size=37	,    layer_norm_eps=1E-05	,    num_attention_heads=4	,    num_hidden_layers=5	,    pad_token_id=1	,    vocab_size=10_00	,    )
     return CLIPTextModelWithProjection(lowercase_)





 @property
 def     snake_case__						(      self):
     torch.manual_seed(0)

     snake_case_	:      str            =  {
         "num_attention_heads": 2,
         "attention_head_dim": 12,
         "embedding_dim": self.text_embedder_hidden_size,
         "num_layers": 1,
     }

     snake_case_	:      List[str]            =  PriorTransformer(**lowercase_)
     return model





 @property
 def     snake_case__						(      self):
     torch.manual_seed(0)

     snake_case_	:      Tuple            =  {
         "clip_embeddings_dim": self.text_embedder_hidden_size,
         "time_embed_dim": self.time_embed_dim,
         "cross_attention_dim": self.cross_attention_dim,
     }

     snake_case_	:      Union[str, Any]            =  UnCLIPTextProjModel(**lowercase_)
     return model





 @property
 def     snake_case__						(      self):
     torch.manual_seed(0)

     snake_case_	:      Union[str, Any]            =  {
         "sample_size": 32,
         # RGB in channels
         "in_channels": 3,
         # Out channels is double in channels because predicts mean and variance
         "out_channels": 6,
         "down_block_types": ("ResnetDownsampleBlock2D", "SimpleCrossAttnDownBlock2D"),
         "up_block_types": ("SimpleCrossAttnUpBlock2D", "ResnetUpsampleBlock2D"),
         "mid_block_type": "UNetMidBlock2DSimpleCrossAttn",
         "block_out_channels": (self.block_out_channels_a, self.block_out_channels_a * 2),
         "layers_per_block": 1,
         "cross_attention_dim": self.cross_attention_dim,
         "attention_head_dim": 4,
         "resnet_time_scale_shift": "scale_shift",
         "class_embed_type": "identity",
     }

     snake_case_	:      List[str]            =  UNetaDConditionModel(**lowercase_)
     return model





 @property
 def     snake_case__						(      self):
     return {
         "sample_size": 64,
         "layers_per_block": 1,
         "down_block_types": ("ResnetDownsampleBlock2D", "ResnetDownsampleBlock2D"),
         "up_block_types": ("ResnetUpsampleBlock2D", "ResnetUpsampleBlock2D"),
         "block_out_channels": (self.block_out_channels_a, self.block_out_channels_a * 2),
         "in_channels": 6,
         "out_channels": 3,
     }





 @property
 def     snake_case__						(      self):
     torch.manual_seed(0)

     snake_case_	:      int            =  UNetaDModel(**self.dummy_super_res_kwargs)
     return model





 @property
 def     snake_case__						(      self):
     # seeded differently to get different unet than `self.dummy_super_res_first`
     torch.manual_seed(1)

     snake_case_	:      List[Any]            =  UNetaDModel(**self.dummy_super_res_kwargs)
     return model





 def     snake_case__						(      self):
     snake_case_	:      Dict            =  self.dummy_prior
     snake_case_	:      Optional[int]            =  self.dummy_decoder
     snake_case_	:      str            =  self.dummy_text_proj
     snake_case_	:      Union[str, Any]            =  self.dummy_text_encoder
     snake_case_	:      Tuple            =  self.dummy_tokenizer
     snake_case_	:      Optional[Any]            =  self.dummy_super_res_first
     snake_case_	:      Tuple            =  self.dummy_super_res_last

     snake_case_	:      Any            =  UnCLIPScheduler(
         variance_type="fixed_small_log"	,    prediction_type="sample"	,    num_train_timesteps=10_00	,    clip_sample_range=5.0	,    )

     snake_case_	:      List[Any]            =  UnCLIPScheduler(
         variance_type="learned_range"	,    prediction_type="epsilon"	,    num_train_timesteps=10_00	,    )

     snake_case_	:      Union[str, Any]            =  UnCLIPScheduler(
         variance_type="fixed_small_log"	,    prediction_type="epsilon"	,    num_train_timesteps=10_00	,    )

     snake_case_	:      Any            =  {
         "prior": prior,
         "decoder": decoder,
         "text_proj": text_proj,
         "text_encoder": text_encoder,
         "tokenizer": tokenizer,
         "super_res_first": super_res_first,
         "super_res_last": super_res_last,
         "prior_scheduler": prior_scheduler,
         "decoder_scheduler": decoder_scheduler,
         "super_res_scheduler": super_res_scheduler,
     }

     return components





 def     snake_case__						(      self	,    lowercase_	,    lowercase_=0):
     if str(lowercase_).startswith("mps"):
         snake_case_	:      List[Any]            =  torch.manual_seed(lowercase_)
     else:
         snake_case_	:      str            =  torch.Generator(device=lowercase_).manual_seed(lowercase_)
     snake_case_	:      List[str]            =  {
         "prompt": "horse",
         "generator": generator,
         "prior_num_inference_steps": 2,
         "decoder_num_inference_steps": 2,
         "super_res_num_inference_steps": 2,
         "output_type": "numpy",
     }
     return inputs





 def     snake_case__						(      self):
     snake_case_	:      str            =  "cpu"

     snake_case_	:      int            =  self.get_dummy_components()

     snake_case_	:      List[str]            =  self.pipeline_class(**lowercase_)
     snake_case_	:      List[Any]            =  pipe.to(lowercase_)

     pipe.set_progress_bar_config(disable=lowercase_)

     snake_case_	:      Dict            =  pipe(**self.get_dummy_inputs(lowercase_))
     snake_case_	:      int            =  output.images

     snake_case_	:      Optional[int]            =  pipe(
         **self.get_dummy_inputs(lowercase_)	,    return_dict=lowercase_	,    )[0]

     snake_case_	:      Any            =  image[0, -3:, -3:, -1]
     snake_case_	:      Optional[Any]            =  image_from_tuple[0, -3:, -3:, -1]

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

     snake_case_	:      Union[str, Any]            =  np.array(
         [
             0.9_997,
             0.9_988,
             0.0_028,
             0.9_997,
             0.9_984,
             0.9_965,
             0.0_029,
             0.9_986,
             0.0_025,
         ])

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





 def     snake_case__						(      self):
     snake_case_	:      int            =  torch.device("cpu")






     class 				UpperCAmelCase_			:
      UpperCAmelCase_						      =  1


     snake_case_	:      Optional[int]            =  self.get_dummy_components()

     snake_case_	:      Union[str, Any]            =  self.pipeline_class(**lowercase_)
     snake_case_	:      Optional[int]            =  pipe.to(lowercase_)

     snake_case_	:      Tuple            =  components["prior"]
     snake_case_	:      Any            =  components["decoder"]
     snake_case_	:      List[Any]            =  components["super_res_first"]
     snake_case_	:      Tuple            =  components["tokenizer"]
     snake_case_	:      Optional[Any]            =  components["text_encoder"]

     snake_case_	:      Dict            =  torch.Generator(device=lowercase_).manual_seed(0)
     snake_case_	:      Tuple            =  prior.dtype
     snake_case_	:      Union[str, Any]            =  1

     snake_case_	:      int            =  (batch_size, prior.config.embedding_dim)
     snake_case_	:      List[Any]            =  pipe.prepare_latents(
         lowercase_	,    dtype=lowercase_	,    device=lowercase_	,    generator=lowercase_	,    latents=lowercase_	,    scheduler=DummyScheduler())
     snake_case_	:      List[str]            =  (batch_size, decoder.config.in_channels, decoder.config.sample_size, decoder.config.sample_size)
     snake_case_	:      Any            =  pipe.prepare_latents(
         lowercase_	,    dtype=lowercase_	,    device=lowercase_	,    generator=lowercase_	,    latents=lowercase_	,    scheduler=DummyScheduler())

     snake_case_	:      int            =  (
         batch_size,
         super_res_first.config.in_channels // 2,
         super_res_first.config.sample_size,
         super_res_first.config.sample_size,
     )
     snake_case_	:      Optional[Any]            =  pipe.prepare_latents(
         lowercase_	,    dtype=lowercase_	,    device=lowercase_	,    generator=lowercase_	,    latents=lowercase_	,    scheduler=DummyScheduler())

     pipe.set_progress_bar_config(disable=lowercase_)

     snake_case_	:      List[Any]            =  "this is a prompt example"

     snake_case_	:      Union[str, Any]            =  torch.Generator(device=lowercase_).manual_seed(0)
     snake_case_	:      Optional[int]            =  pipe(
         [prompt]	,    generator=lowercase_	,    prior_num_inference_steps=2	,    decoder_num_inference_steps=2	,    super_res_num_inference_steps=2	,    prior_latents=lowercase_	,    decoder_latents=lowercase_	,    super_res_latents=lowercase_	,    output_type="np"	,    )
     snake_case_	:      Optional[int]            =  output.images

     snake_case_	:      Union[str, Any]            =  tokenizer(
         lowercase_	,    padding="max_length"	,    max_length=tokenizer.model_max_length	,    return_tensors="pt"	,    )
     snake_case_	:      List[str]            =  text_encoder(text_inputs.input_ids)
     snake_case_	:      Any            =  text_inputs.attention_mask

     snake_case_	:      List[str]            =  torch.Generator(device=lowercase_).manual_seed(0)
     snake_case_	:      List[str]            =  pipe(
         generator=lowercase_	,    prior_num_inference_steps=2	,    decoder_num_inference_steps=2	,    super_res_num_inference_steps=2	,    prior_latents=lowercase_	,    decoder_latents=lowercase_	,    super_res_latents=lowercase_	,    text_model_output=lowercase_	,    text_attention_mask=lowercase_	,    output_type="np"	,    )[0]

     # make sure passing text embeddings manually is identical
     assert np.abs(image - image_from_text).max() < 1E-4





 @skip_mps
 def     snake_case__						(      self):
     snake_case_	:      str            =  torch_device == "cpu"

     self._test_attention_slicing_forward_pass(test_max_difference=lowercase_	,    expected_max_diff=0.01)





 @skip_mps
 def     snake_case__						(      self):
     snake_case_	:      Dict            =  torch_device == "cpu"
     snake_case_	:      int            =  True
     snake_case_	:      Union[str, Any]            =  [
         "prior_num_inference_steps",
         "decoder_num_inference_steps",
         "super_res_num_inference_steps",
     ]

     self._test_inference_batch_single_identical(
         test_max_difference=lowercase_	,    relax_max_difference=lowercase_	,    additional_params_copy_to_batched_inputs=lowercase_	,    )





 def     snake_case__						(      self):
     snake_case_	:      Dict            =  [
         "prior_num_inference_steps",
         "decoder_num_inference_steps",
         "super_res_num_inference_steps",
     ]

     if torch_device == "mps":
         # TODO: MPS errors with larger batch sizes
         snake_case_	:      Dict            =  [2, 3]
         self._test_inference_batch_consistent(
             batch_sizes=lowercase_	,    additional_params_copy_to_batched_inputs=lowercase_	,    )
     else:
         self._test_inference_batch_consistent(
             additional_params_copy_to_batched_inputs=lowercase_)





 @skip_mps
 def     snake_case__						(      self):
     return super().test_dict_tuple_outputs_equivalent()





 @skip_mps
 def     snake_case__						(      self):
     return super().test_save_load_local()





 @skip_mps
 def     snake_case__						(      self):
     return super().test_save_load_optional_components()






@nightly
class 				UpperCAmelCase_			(		unittest.TestCase   ):





 def     snake_case__						(      self):
     # clean up the VRAM after each test
     super().tearDown()
     gc.collect()
     torch.cuda.empty_cache()





 def     snake_case__						(      self):
     snake_case_	:      List[str]            =  load_numpy(
         "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main"
         "/unclip/karlo_v1_alpha_horse_cpu.npy")

     snake_case_	:      List[str]            =  UnCLIPPipeline.from_pretrained("kakaobrain/karlo-v1-alpha")
     pipeline.set_progress_bar_config(disable=lowercase_)

     snake_case_	:      Any            =  torch.manual_seed(0)
     snake_case_	:      Any            =  pipeline(
         "horse"	,    num_images_per_prompt=1	,    generator=lowercase_	,    output_type="np"	,    )

     snake_case_	:      Dict            =  output.images[0]

     assert image.shape == (2_56, 2_56, 3)
     assert np.abs(expected_image - image).max() < 1E-1






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





 def     snake_case__						(      self):
     # clean up the VRAM after each test
     super().tearDown()
     gc.collect()
     torch.cuda.empty_cache()





 def     snake_case__						(      self):
     snake_case_	:      Tuple            =  load_numpy(
         "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main"
         "/unclip/karlo_v1_alpha_horse_fp16.npy")

     snake_case_	:      List[str]            =  UnCLIPPipeline.from_pretrained("kakaobrain/karlo-v1-alpha"	,    torch_dtype=torch.floataa)
     snake_case_	:      str            =  pipeline.to(lowercase_)
     pipeline.set_progress_bar_config(disable=lowercase_)

     snake_case_	:      Union[str, Any]            =  torch.Generator(device="cpu").manual_seed(0)
     snake_case_	:      int            =  pipeline(
         "horse"	,    generator=lowercase_	,    output_type="np"	,    )

     snake_case_	:      str            =  output.images[0]

     assert image.shape == (2_56, 2_56, 3)

     assert_mean_pixel_difference(lowercase_	,    lowercase_)





 def     snake_case__						(      self):
     torch.cuda.empty_cache()
     torch.cuda.reset_max_memory_allocated()
     torch.cuda.reset_peak_memory_stats()

     snake_case_	:      int            =  UnCLIPPipeline.from_pretrained("kakaobrain/karlo-v1-alpha"	,    torch_dtype=torch.floataa)
     snake_case_	:      Any            =  pipe.to(lowercase_)
     pipe.set_progress_bar_config(disable=lowercase_)
     pipe.enable_attention_slicing()
     pipe.enable_sequential_cpu_offload()

     snake_case_	:      Tuple            =  pipe(
         "horse"	,    num_images_per_prompt=1	,    prior_num_inference_steps=2	,    decoder_num_inference_steps=2	,    super_res_num_inference_steps=2	,    output_type="np"	,    )

     snake_case_	:      Union[str, Any]            =  torch.cuda.max_memory_allocated()
     # make sure that less than 7 GB is allocated
     assert mem_bytes < 7 * 10**9






'''simple docstring'''
import importlib.metadata
import sys

from transformers.testing_utils import TestCasePlus
from transformers.utils.versions import require_version, require_version_core


a_     =  importlib.metadata.version("numpy")
a_     =  ".".join([str(x) for x in sys.version_info[:3]])






class 				UpperCAmelCase_			(		snake_case__   ):





 def     snake_case__						(      self):
     # lt + different version strings
     require_version_core("numpy<1000.4.5")
     require_version_core("numpy<1000.4")
     require_version_core("numpy<1000")

     # le
     require_version_core("numpy<=1000.4.5")
     require_version_core(F'numpy<={numpy_ver}')

     # eq
     require_version_core(F'numpy=={numpy_ver}')

     # ne
     require_version_core("numpy!=1000.4.5")

     # ge
     require_version_core("numpy>=1.0")
     require_version_core("numpy>=1.0.0")
     require_version_core(F'numpy>={numpy_ver}')

     # gt
     require_version_core("numpy>1.0.0")

     # mix
     require_version_core("numpy>1.0.0,<1000")

     # requirement w/o version
     require_version_core("numpy")

     # unmet requirements due to version conflict
     for req in ["numpy==1.0.0", "numpy>=1000.0.0", F'numpy<{numpy_ver}']:
         try:
             require_version_core(lowercase_)
         except ImportError as e:
             self.assertIn(F'{req} is required'	,    str(lowercase_))
             self.assertIn("but found"	,    str(lowercase_))

        # unmet requirements due to missing module
     for req in ["numpipypie>1", "numpipypie2"]:
         try:
             require_version_core(lowercase_)
         except importlib.metadata.PackageNotFoundError as e:
             self.assertIn(F'The \'{req}\' distribution was not found and is required by this application'	,    str(lowercase_))
             self.assertIn("Try: pip install transformers -U"	,    str(lowercase_))

        # bogus requirements formats:
        # 1. whole thing
     for req in ["numpy??1.0.0", "numpy1.0.0"]:
         try:
             require_version_core(lowercase_)
         except ValueError as e:
             self.assertIn("requirement needs to be in the pip package format"	,    str(lowercase_))
        # 2. only operators
     for req in ["numpy=1.0.0", "numpy == 1.00", "numpy<>1.0.0", "numpy><1.00", "numpy>>1.0.0"]:
         try:
             require_version_core(lowercase_)
         except ValueError as e:
             self.assertIn("need one of "	,    str(lowercase_))





 def     snake_case__						(      self):
     # matching requirement
     require_version("python>=3.6.0")

     # not matching requirements
     for req in ["python>9.9.9", "python<3.0.0"]:
         try:
             require_version_core(lowercase_)
         except ImportError as e:
             self.assertIn(F'{req} is required'	,    str(lowercase_))
             self.assertIn(F'but found python=={python_ver}'	,    str(lowercase_))




