
'''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 		snake_case_       (						A__		,	unittest.TestCase						):



   """simple docstring"""





   __lowerCAmelCase	:							Optional[int]								=UnCLIPPipeline
   __lowerCAmelCase	:							Optional[Any]								=TEXT_TO_IMAGE_PARAMS - {
       '''negative_prompt''',
       '''height''',
       '''width''',
       '''negative_prompt_embeds''',
       '''guidance_scale''',
       '''prompt_embeds''',
       '''cross_attention_kwargs''',
   }
   __lowerCAmelCase	:							List[str]								=TEXT_TO_IMAGE_BATCH_PARAMS
   __lowerCAmelCase	:							Any								=[
       '''generator''',
       '''return_dict''',
       '''prior_num_inference_steps''',
       '''decoder_num_inference_steps''',
       '''super_res_num_inference_steps''',
   ]
   __lowerCAmelCase	:							str								=False







   @property
   def     __UpperCAmelCase      ( self):
          return 32







   @property
   def     __UpperCAmelCase      ( self):
          return 32







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







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







   @property
   def     __UpperCAmelCase      ( self):
          return 1_00







   @property
   def     __UpperCAmelCase      ( self):
          lowerCamelCase__				   =      CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip")
          return tokenizer







   @property
   def     __UpperCAmelCase      ( self):
          torch.manual_seed(0)
          lowerCamelCase__				   =      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-0_5			,   num_attention_heads=4			,   num_hidden_layers=5			,   pad_token_id=1			,   vocab_size=10_00			,   )
          return CLIPTextModelWithProjection(UpperCamelCase)







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

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

          lowerCamelCase__				   =      PriorTransformer(**UpperCamelCase)
          return model







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

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

          lowerCamelCase__				   =      UnCLIPTextProjModel(**UpperCamelCase)
          return model







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

          lowerCamelCase__				   =      {
              "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",
          }

          lowerCamelCase__				   =      UNetaDConditionModel(**UpperCamelCase)
          return model







   @property
   def     __UpperCAmelCase      ( 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     __UpperCAmelCase      ( self):
          torch.manual_seed(0)

          lowerCamelCase__				   =      UNetaDModel(**self.dummy_super_res_kwargs)
          return model







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

          lowerCamelCase__				   =      UNetaDModel(**self.dummy_super_res_kwargs)
          return model







   def     __UpperCAmelCase      ( self):
          lowerCamelCase__				   =      self.dummy_prior
          lowerCamelCase__				   =      self.dummy_decoder
          lowerCamelCase__				   =      self.dummy_text_proj
          lowerCamelCase__				   =      self.dummy_text_encoder
          lowerCamelCase__				   =      self.dummy_tokenizer
          lowerCamelCase__				   =      self.dummy_super_res_first
          lowerCamelCase__				   =      self.dummy_super_res_last

          lowerCamelCase__				   =      UnCLIPScheduler(
              variance_type="fixed_small_log"			,   prediction_type="sample"			,   num_train_timesteps=10_00			,   clip_sample_range=5.0			,   )

          lowerCamelCase__				   =      UnCLIPScheduler(
              variance_type="learned_range"			,   prediction_type="epsilon"			,   num_train_timesteps=10_00			,   )

          lowerCamelCase__				   =      UnCLIPScheduler(
              variance_type="fixed_small_log"			,   prediction_type="epsilon"			,   num_train_timesteps=10_00			,   )

          lowerCamelCase__				   =      {
              "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     __UpperCAmelCase      ( self			,   UpperCamelCase			,   UpperCamelCase=0):
          if str(UpperCamelCase).startswith("mps"):
                 lowerCamelCase__				   =      torch.manual_seed(UpperCamelCase)
          else:
                 lowerCamelCase__				   =      torch.Generator(device=UpperCamelCase).manual_seed(UpperCamelCase)
          lowerCamelCase__				   =      {
              "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     __UpperCAmelCase      ( self):
          lowerCamelCase__				   =      "cpu"

          lowerCamelCase__				   =      self.get_dummy_components()

          lowerCamelCase__				   =      self.pipeline_class(**UpperCamelCase)
          lowerCamelCase__				   =      pipe.to(UpperCamelCase)

          pipe.set_progress_bar_config(disable=UpperCamelCase)

          lowerCamelCase__				   =      pipe(**self.get_dummy_inputs(UpperCamelCase))
          lowerCamelCase__				   =      output.images

          lowerCamelCase__				   =      pipe(
              **self.get_dummy_inputs(UpperCamelCase)			,   return_dict=UpperCamelCase			,   )[0]

          lowerCamelCase__				   =      image[0, -3:, -3:, -1]
          lowerCamelCase__				   =      image_from_tuple[0, -3:, -3:, -1]

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

          lowerCamelCase__				   =      np.array(
              [
                  0.9_9_9_7,
                  0.9_9_8_8,
                  0.0_0_2_8,
                  0.9_9_9_7,
                  0.9_9_8_4,
                  0.9_9_6_5,
                  0.0_0_2_9,
                  0.9_9_8_6,
                  0.0_0_2_5,
              ])

          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     __UpperCAmelCase      ( self):
          lowerCamelCase__				   =      torch.device("cpu")







          class 		snake_case_       :



             """simple docstring"""





             __lowerCAmelCase	:							Optional[Any]								=1






          lowerCamelCase__				   =      self.get_dummy_components()

          lowerCamelCase__				   =      self.pipeline_class(**UpperCamelCase)
          lowerCamelCase__				   =      pipe.to(UpperCamelCase)

          lowerCamelCase__				   =      components["prior"]
          lowerCamelCase__				   =      components["decoder"]
          lowerCamelCase__				   =      components["super_res_first"]
          lowerCamelCase__				   =      components["tokenizer"]
          lowerCamelCase__				   =      components["text_encoder"]

          lowerCamelCase__				   =      torch.Generator(device=UpperCamelCase).manual_seed(0)
          lowerCamelCase__				   =      prior.dtype
          lowerCamelCase__				   =      1

          lowerCamelCase__				   =      (batch_size, prior.config.embedding_dim)
          lowerCamelCase__				   =      pipe.prepare_latents(
              UpperCamelCase			,   dtype=UpperCamelCase			,   device=UpperCamelCase			,   generator=UpperCamelCase			,   latents=UpperCamelCase			,   scheduler=DummyScheduler())
          lowerCamelCase__				   =      (batch_size, decoder.config.in_channels, decoder.config.sample_size, decoder.config.sample_size)
          lowerCamelCase__				   =      pipe.prepare_latents(
              UpperCamelCase			,   dtype=UpperCamelCase			,   device=UpperCamelCase			,   generator=UpperCamelCase			,   latents=UpperCamelCase			,   scheduler=DummyScheduler())

          lowerCamelCase__				   =      (
              batch_size,
              super_res_first.config.in_channels // 2,
              super_res_first.config.sample_size,
              super_res_first.config.sample_size,
          )
          lowerCamelCase__				   =      pipe.prepare_latents(
              UpperCamelCase			,   dtype=UpperCamelCase			,   device=UpperCamelCase			,   generator=UpperCamelCase			,   latents=UpperCamelCase			,   scheduler=DummyScheduler())

          pipe.set_progress_bar_config(disable=UpperCamelCase)

          lowerCamelCase__				   =      "this is a prompt example"

          lowerCamelCase__				   =      torch.Generator(device=UpperCamelCase).manual_seed(0)
          lowerCamelCase__				   =      pipe(
              [prompt]			,   generator=UpperCamelCase			,   prior_num_inference_steps=2			,   decoder_num_inference_steps=2			,   super_res_num_inference_steps=2			,   prior_latents=UpperCamelCase			,   decoder_latents=UpperCamelCase			,   super_res_latents=UpperCamelCase			,   output_type="np"			,   )
          lowerCamelCase__				   =      output.images

          lowerCamelCase__				   =      tokenizer(
              UpperCamelCase			,   padding="max_length"			,   max_length=tokenizer.model_max_length			,   return_tensors="pt"			,   )
          lowerCamelCase__				   =      text_encoder(text_inputs.input_ids)
          lowerCamelCase__				   =      text_inputs.attention_mask

          lowerCamelCase__				   =      torch.Generator(device=UpperCamelCase).manual_seed(0)
          lowerCamelCase__				   =      pipe(
              generator=UpperCamelCase			,   prior_num_inference_steps=2			,   decoder_num_inference_steps=2			,   super_res_num_inference_steps=2			,   prior_latents=UpperCamelCase			,   decoder_latents=UpperCamelCase			,   super_res_latents=UpperCamelCase			,   text_model_output=UpperCamelCase			,   text_attention_mask=UpperCamelCase			,   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     __UpperCAmelCase      ( self):
          lowerCamelCase__				   =      torch_device == "cpu"

          self._test_attention_slicing_forward_pass(test_max_difference=UpperCamelCase			,   expected_max_diff=0.0_1)







   @skip_mps
   def     __UpperCAmelCase      ( self):
          lowerCamelCase__				   =      torch_device == "cpu"
          lowerCamelCase__				   =      True
          lowerCamelCase__				   =      [
              "prior_num_inference_steps",
              "decoder_num_inference_steps",
              "super_res_num_inference_steps",
          ]

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







   def     __UpperCAmelCase      ( self):
          lowerCamelCase__				   =      [
              "prior_num_inference_steps",
              "decoder_num_inference_steps",
              "super_res_num_inference_steps",
          ]

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







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







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







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







@nightly
class 		snake_case_       (						unittest.TestCase						):



   """simple docstring"""





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







   def     __UpperCAmelCase      ( self):
          lowerCamelCase__				   =      load_numpy(
              "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main"
              "/unclip/karlo_v1_alpha_horse_cpu.npy")

          lowerCamelCase__				   =      UnCLIPPipeline.from_pretrained("kakaobrain/karlo-v1-alpha")
          pipeline.set_progress_bar_config(disable=UpperCamelCase)

          lowerCamelCase__				   =      torch.manual_seed(0)
          lowerCamelCase__				   =      pipeline(
              "horse"			,   num_images_per_prompt=1			,   generator=UpperCamelCase			,   output_type="np"			,   )

          lowerCamelCase__				   =      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 		snake_case_       (						unittest.TestCase						):



   """simple docstring"""





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







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

          lowerCamelCase__				   =      UnCLIPPipeline.from_pretrained("kakaobrain/karlo-v1-alpha"			,   torch_dtype=torch.floataa)
          lowerCamelCase__				   =      pipeline.to(UpperCamelCase)
          pipeline.set_progress_bar_config(disable=UpperCamelCase)

          lowerCamelCase__				   =      torch.Generator(device="cpu").manual_seed(0)
          lowerCamelCase__				   =      pipeline(
              "horse"			,   generator=UpperCamelCase			,   output_type="np"			,   )

          lowerCamelCase__				   =      output.images[0]

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

          assert_mean_pixel_difference(UpperCamelCase			,   UpperCamelCase)







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

          lowerCamelCase__				   =      UnCLIPPipeline.from_pretrained("kakaobrain/karlo-v1-alpha"			,   torch_dtype=torch.floataa)
          lowerCamelCase__				   =      pipe.to(UpperCamelCase)
          pipe.set_progress_bar_config(disable=UpperCamelCase)
          pipe.enable_attention_slicing()
          pipe.enable_sequential_cpu_offload()

          lowerCamelCase__				   =      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"			,   )

          lowerCamelCase__				   =      torch.cuda.max_memory_allocated()
          # make sure that less than 7 GB is allocated
          assert mem_bytes < 7 * 10**9



'''simple docstring'''







def   lowerCAmelCase(  a__    :  int = 1000      ):







       '''simple docstring'''






       lowerCamelCase__				   =      0
       for i in range(a__      ):
              if i % 3 == 0 or i % 5 == 0:
                     result += i
       return result


if __name__ == "__main__":
    print(f'{solution() = }')

