





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       (						a__	,    unittest.TestCase  ):
       _UpperCamelCase					:       Optional[Any]       = UnCLIPPipeline
       _UpperCamelCase					:       Optional[Any]       = TEXT_TO_IMAGE_PARAMS - {
           '''negative_prompt''',
           '''height''',
           '''width''',
           '''negative_prompt_embeds''',
           '''guidance_scale''',
           '''prompt_embeds''',
           '''cross_attention_kwargs''',
       }
       _UpperCamelCase					:       Optional[int]       = TEXT_TO_IMAGE_BATCH_PARAMS
       _UpperCamelCase					:       List[str]       = [
           '''generator''',
           '''return_dict''',
           '''prior_num_inference_steps''',
           '''decoder_num_inference_steps''',
           '''super_res_num_inference_steps''',
       ]
       _UpperCamelCase					:       int       = False

       @property
       def lowercase			(		self:			Optional[Any]		)							->			List[str]:




              """simple docstring"""


              return 32

       @property
       def lowercase			(		self:			List[Any]		)							->			Optional[int]:




              """simple docstring"""


              return 32

       @property
       def lowercase			(		self:			List[Any]		)							->			Tuple:




              """simple docstring"""


              return self.time_input_dim

       @property
       def lowercase			(		self:			str		)							->			Dict:




              """simple docstring"""


              return self.time_input_dim * 4

       @property
       def lowercase			(		self:			Optional[int]		)							->			Any:




              """simple docstring"""


              return 100

       @property
       def lowercase			(		self:			Union[str, Any]		)							->			List[Any]:




              """simple docstring"""


              UpperCamelCase_ 							=	CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip"		)
              return tokenizer

       @property
       def lowercase			(		self:			Any		)							->			Tuple:




              """simple docstring"""


              torch.manual_seed(0		)
              UpperCamelCase_ 							=	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=1000       , )
              return CLIPTextModelWithProjection(_SCREAMING_SNAKE_CASE		)

       @property
       def lowercase			(		self:			str		)							->			List[str]:




              """simple docstring"""


              torch.manual_seed(0		)

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

              UpperCamelCase_ 							=	PriorTransformer(**_SCREAMING_SNAKE_CASE		)
              return model

       @property
       def lowercase			(		self:			str		)							->			Optional[Any]:




              """simple docstring"""


              torch.manual_seed(0		)

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

              UpperCamelCase_ 							=	UnCLIPTextProjModel(**_SCREAMING_SNAKE_CASE		)
              return model

       @property
       def lowercase			(		self:			Optional[int]		)							->			List[str]:




              """simple docstring"""


              torch.manual_seed(0		)

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

              UpperCamelCase_ 							=	UNetaDConditionModel(**_SCREAMING_SNAKE_CASE		)
              return model

       @property
       def lowercase			(		self:			Any		)							->			int:




              """simple docstring"""


              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 lowercase			(		self:			Dict		)							->			Dict:




              """simple docstring"""


              torch.manual_seed(0		)

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

       @property
       def lowercase			(		self:			Any		)							->			Any:




              """simple docstring"""


              torch.manual_seed(1		)

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

       def lowercase			(		self:			List[str]		)							->			Optional[Any]:




              """simple docstring"""


              UpperCamelCase_ 							=	self.dummy_prior
              UpperCamelCase_ 							=	self.dummy_decoder
              UpperCamelCase_ 							=	self.dummy_text_proj
              UpperCamelCase_ 							=	self.dummy_text_encoder
              UpperCamelCase_ 							=	self.dummy_tokenizer
              UpperCamelCase_ 							=	self.dummy_super_res_first
              UpperCamelCase_ 							=	self.dummy_super_res_last

              UpperCamelCase_ 							=	UnCLIPScheduler(
                  variance_type="fixed_small_log"       , prediction_type="sample"       , num_train_timesteps=1000       , clip_sample_range=5.0       , )

              UpperCamelCase_ 							=	UnCLIPScheduler(
                  variance_type="learned_range"       , prediction_type="epsilon"       , num_train_timesteps=1000       , )

              UpperCamelCase_ 							=	UnCLIPScheduler(
                  variance_type="fixed_small_log"       , prediction_type="epsilon"       , num_train_timesteps=1000       , )

              UpperCamelCase_ 							=	{
                  "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 lowercase			(		self:			str       , _SCREAMING_SNAKE_CASE:			List[Any]       , _SCREAMING_SNAKE_CASE:			Dict=0		)							->			List[str]:




              """simple docstring"""


              if str(_SCREAMING_SNAKE_CASE		).startswith("mps"		):
                     UpperCamelCase_ 							=	torch.manual_seed(_SCREAMING_SNAKE_CASE		)
              else:
                     UpperCamelCase_ 							=	torch.Generator(device=_SCREAMING_SNAKE_CASE		).manual_seed(_SCREAMING_SNAKE_CASE		)
              UpperCamelCase_ 							=	{
                  "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 lowercase			(		self:			List[Any]		)							->			Any:




              """simple docstring"""


              UpperCamelCase_ 							=	"cpu"

              UpperCamelCase_ 							=	self.get_dummy_components()

              UpperCamelCase_ 							=	self.pipeline_class(**_SCREAMING_SNAKE_CASE		)
              UpperCamelCase_ 							=	pipe.to(_SCREAMING_SNAKE_CASE		)

              pipe.set_progress_bar_config(disable=_SCREAMING_SNAKE_CASE		)

              UpperCamelCase_ 							=	pipe(**self.get_dummy_inputs(_SCREAMING_SNAKE_CASE		)		)
              UpperCamelCase_ 							=	output.images

              UpperCamelCase_ 							=	pipe(
                  **self.get_dummy_inputs(_SCREAMING_SNAKE_CASE		)       , return_dict=_SCREAMING_SNAKE_CASE       , )[0]

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

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

              UpperCamelCase_ 							=	np.array(
                  [
                      0.99_97,
                      0.99_88,
                      0.00_28,
                      0.99_97,
                      0.99_84,
                      0.99_65,
                      0.00_29,
                      0.99_86,
                      0.00_25,
                  ]		)

              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 lowercase			(		self:			str		)							->			Optional[Any]:




              """simple docstring"""


              UpperCamelCase_ 							=	torch.device("cpu"		)



              class 				_UpperCamelCase       :
                     _UpperCamelCase					:       Optional[Any]       = 1

              UpperCamelCase_ 							=	self.get_dummy_components()

              UpperCamelCase_ 							=	self.pipeline_class(**_SCREAMING_SNAKE_CASE		)
              UpperCamelCase_ 							=	pipe.to(_SCREAMING_SNAKE_CASE		)

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

              UpperCamelCase_ 							=	torch.Generator(device=_SCREAMING_SNAKE_CASE		).manual_seed(0		)
              UpperCamelCase_ 							=	prior.dtype
              UpperCamelCase_ 							=	1

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

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

              pipe.set_progress_bar_config(disable=_SCREAMING_SNAKE_CASE		)

              UpperCamelCase_ 							=	"this is a prompt example"

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

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

              UpperCamelCase_ 							=	torch.Generator(device=_SCREAMING_SNAKE_CASE		).manual_seed(0		)
              UpperCamelCase_ 							=	pipe(
                  generator=_SCREAMING_SNAKE_CASE       , prior_num_inference_steps=2       , decoder_num_inference_steps=2       , super_res_num_inference_steps=2       , prior_latents=_SCREAMING_SNAKE_CASE       , decoder_latents=_SCREAMING_SNAKE_CASE       , super_res_latents=_SCREAMING_SNAKE_CASE       , text_model_output=_SCREAMING_SNAKE_CASE       , text_attention_mask=_SCREAMING_SNAKE_CASE       , 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 lowercase			(		self:			List[Any]		)							->			List[Any]:




              """simple docstring"""


              UpperCamelCase_ 							=	torch_device == "cpu"

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

       @skip_mps
       def lowercase			(		self:			Dict		)							->			List[str]:




              """simple docstring"""


              UpperCamelCase_ 							=	torch_device == "cpu"
              UpperCamelCase_ 							=	True
              UpperCamelCase_ 							=	[
                  "prior_num_inference_steps",
                  "decoder_num_inference_steps",
                  "super_res_num_inference_steps",
              ]

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

       def lowercase			(		self:			List[str]		)							->			Dict:




              """simple docstring"""


              UpperCamelCase_ 							=	[
                  "prior_num_inference_steps",
                  "decoder_num_inference_steps",
                  "super_res_num_inference_steps",
              ]

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

       @skip_mps
       def lowercase			(		self:			List[str]		)							->			str:




              """simple docstring"""


              return super().test_dict_tuple_outputs_equivalent()

       @skip_mps
       def lowercase			(		self:			Optional[int]		)							->			List[str]:




              """simple docstring"""


              return super().test_save_load_local()



       @skip_mps
       def lowercase			(		self:			Optional[Any]		)							->			str:




              """simple docstring"""


              return super().test_save_load_optional_components()



@nightly
class 				_UpperCamelCase       (						unittest.TestCase  ):

       def lowercase			(		self:			Dict		)							->			Tuple:




              """simple docstring"""


              super().tearDown()
              gc.collect()
              torch.cuda.empty_cache()



       def lowercase			(		self:			List[str]		)							->			List[str]:




              """simple docstring"""


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

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

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

              UpperCamelCase_ 							=	output.images[0]

              assert image.shape == (256, 256, 3)
              assert np.abs(expected_image - image		).max() < 1e-1



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

       def lowercase			(		self:			List[Any]		)							->			List[str]:




              """simple docstring"""


              super().tearDown()
              gc.collect()
              torch.cuda.empty_cache()

       def lowercase			(		self:			str		)							->			Dict:




              """simple docstring"""


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

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

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

              UpperCamelCase_ 							=	output.images[0]

              assert image.shape == (256, 256, 3)

              assert_mean_pixel_difference(_SCREAMING_SNAKE_CASE       , _SCREAMING_SNAKE_CASE		)



       def lowercase			(		self:			Tuple		)							->			Any:




              """simple docstring"""


              torch.cuda.empty_cache()
              torch.cuda.reset_max_memory_allocated()
              torch.cuda.reset_peak_memory_stats()

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

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

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





'''simple docstring'''







import itertools







def _UpperCamelCase			(	__A = 1000000       )				->						int:


      '''simple docstring'''



      UpperCamelCase__      =							[1]

      for i in itertools.count(len(__A       )       ):
            UpperCamelCase__      =							0
            for j in itertools.count(1       ):
                  UpperCamelCase__      =							-1 if j % 2 == 0 else +1
                  UpperCamelCase__      =							(j * j * 3 - j) // 2
                  if index > i:
                        break
                  item += partitions[i - index] * sign
                  item %= number
                  index += j
                  if index > i:
                        break
                  item += partitions[i - index] * sign
                  item %= number

            if item == 0:
                  return i
            partitions.append(__A       )

      return 0


if __name__ == "__main__":
  import doctest

  doctest.testmod()

  print(F"""{solution() = }""")





