

'''simple docstring'''



import gc
import unittest

import torch
from transformers import CLIPTextConfig, CLIPTextModel, CLIPTextModelWithProjection, CLIPTokenizer

from diffusers import (
    AutoencoderKL,
    DDIMScheduler,
    DDPMScheduler,
    PriorTransformer,
    StableUnCLIPPipeline,
    UNetaDConditionModel,
)
from diffusers.pipelines.stable_diffusion.stable_unclip_image_normalizer import StableUnCLIPImageNormalizer
from diffusers.utils.testing_utils import enable_full_determinism, load_numpy, require_torch_gpu, slow, torch_device

from ..pipeline_params import TEXT_TO_IMAGE_BATCH_PARAMS, TEXT_TO_IMAGE_IMAGE_PARAMS, TEXT_TO_IMAGE_PARAMS
from ..test_pipelines_common import (
    PipelineKarrasSchedulerTesterMixin,
    PipelineLatentTesterMixin,
    PipelineTesterMixin,
    assert_mean_pixel_difference,
)


enable_full_determinism()







class 		lowerCAmelCase__  ( __lowercase					,   __lowercase					,   __lowercase					,   unittest.TestCase   ):



 """simple docstring"""



 lowerCAmelCase__       = StableUnCLIPPipeline
 lowerCAmelCase__       = TEXT_TO_IMAGE_PARAMS
 lowerCAmelCase__       = TEXT_TO_IMAGE_BATCH_PARAMS
 lowerCAmelCase__       = TEXT_TO_IMAGE_IMAGE_PARAMS
 lowerCAmelCase__       = TEXT_TO_IMAGE_IMAGE_PARAMS

 # TODO(will) Expected attn_bias.stride(1) == 0 to be true, but got false
 lowerCAmelCase__       = False
 def  UpperCAmelCase__						(						self   :			Optional[int]     )    ->					Optional[Any]:


      """simple docstring"""




      __SCREAMING_SNAKE_CASE      = 32
      __SCREAMING_SNAKE_CASE      = embedder_hidden_size

      # prior components

      torch.manual_seed(0     )
      __SCREAMING_SNAKE_CASE      = CLIPTokenizer.from_pretrained("""hf-internal-testing/tiny-random-clip"""     )

      torch.manual_seed(0     )
      __SCREAMING_SNAKE_CASE      = CLIPTextModelWithProjection(
          CLIPTextConfig(
              bos_token_id=0   ,				eos_token_id=2   ,				hidden_size=_A   ,				projection_dim=_A   ,				intermediate_size=37   ,				layer_norm_eps=1E-05   ,				num_attention_heads=4   ,				num_hidden_layers=5   ,				pad_token_id=1   ,				vocab_size=1_000   ,				)     )

      torch.manual_seed(0     )
      __SCREAMING_SNAKE_CASE      = PriorTransformer(
          num_attention_heads=2   ,				attention_head_dim=12   ,				embedding_dim=_A   ,				num_layers=1   ,				)

      torch.manual_seed(0     )
      __SCREAMING_SNAKE_CASE      = DDPMScheduler(
          variance_type="""fixed_small_log"""   ,				prediction_type="""sample"""   ,				num_train_timesteps=1_000   ,				clip_sample=_A   ,				clip_sample_range=5.0   ,				beta_schedule="""squaredcos_cap_v2"""   ,				)

      # regular denoising components

      torch.manual_seed(0     )
      __SCREAMING_SNAKE_CASE      = StableUnCLIPImageNormalizer(embedding_dim=_A     )
      __SCREAMING_SNAKE_CASE      = DDPMScheduler(beta_schedule="""squaredcos_cap_v2"""     )

      torch.manual_seed(0     )
      __SCREAMING_SNAKE_CASE      = CLIPTokenizer.from_pretrained("""hf-internal-testing/tiny-random-clip"""     )

      torch.manual_seed(0     )
      __SCREAMING_SNAKE_CASE      = CLIPTextModel(
          CLIPTextConfig(
              bos_token_id=0   ,				eos_token_id=2   ,				hidden_size=_A   ,				projection_dim=32   ,				intermediate_size=37   ,				layer_norm_eps=1E-05   ,				num_attention_heads=4   ,				num_hidden_layers=5   ,				pad_token_id=1   ,				vocab_size=1_000   ,				)     )

      torch.manual_seed(0     )
      __SCREAMING_SNAKE_CASE      = UNetaDConditionModel(
          sample_size=32   ,				in_channels=4   ,				out_channels=4   ,				down_block_types=("""CrossAttnDownBlock2D""", """DownBlock2D""")   ,				up_block_types=("""UpBlock2D""", """CrossAttnUpBlock2D""")   ,				block_out_channels=(32, 64)   ,				attention_head_dim=(2, 4)   ,				class_embed_type="""projection"""   ,				projection_class_embeddings_input_dim=embedder_projection_dim * 2   ,				cross_attention_dim=_A   ,				layers_per_block=1   ,				upcast_attention=_A   ,				use_linear_projection=_A   ,				)

      torch.manual_seed(0     )
      __SCREAMING_SNAKE_CASE      = DDIMScheduler(
          beta_schedule="""scaled_linear"""   ,				beta_start=0.00085   ,				beta_end=0.012   ,				prediction_type="""v_prediction"""   ,				set_alpha_to_one=_A   ,				steps_offset=1   ,				)

      torch.manual_seed(0     )
      __SCREAMING_SNAKE_CASE      = AutoencoderKL()

      __SCREAMING_SNAKE_CASE      = {
          # prior components
          """prior_tokenizer""": prior_tokenizer,
          """prior_text_encoder""": prior_text_encoder,
          """prior""": prior,
          """prior_scheduler""": prior_scheduler,
          # image noising components
          """image_normalizer""": image_normalizer,
          """image_noising_scheduler""": image_noising_scheduler,
          # regular denoising components
          """tokenizer""": tokenizer,
          """text_encoder""": text_encoder,
          """unet""": unet,
          """scheduler""": scheduler,
          """vae""": vae,
      }

      return components
 def  UpperCAmelCase__						(						self   :			Dict   ,				__SCREAMING_SNAKE_CASE   :			Tuple   ,				__SCREAMING_SNAKE_CASE   :			Dict=0     )    ->					int:


      """simple docstring"""




      if str(_A     ).startswith("""mps"""     ):
           __SCREAMING_SNAKE_CASE      = torch.manual_seed(_A     )
      else:
           __SCREAMING_SNAKE_CASE      = torch.Generator(device=_A     ).manual_seed(_A     )
      __SCREAMING_SNAKE_CASE      = {
          """prompt""": """A painting of a squirrel eating a burger""",
          """generator""": generator,
          """num_inference_steps""": 2,
          """prior_num_inference_steps""": 2,
          """output_type""": """numpy""",
      }
      return inputs
 def  UpperCAmelCase__						(						self   :			Dict     )    ->					Optional[int]:


      """simple docstring"""




      __SCREAMING_SNAKE_CASE      = torch_device == """cpu"""

      self._test_attention_slicing_forward_pass(test_max_difference=_A     )


 def  UpperCAmelCase__						(						self   :			List[Any]     )    ->					Tuple:


      """simple docstring"""




      __SCREAMING_SNAKE_CASE      = torch_device in ["""cpu""", """mps"""]

      self._test_inference_batch_single_identical(test_max_difference=_A     )







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



 """simple docstring"""



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


      """simple docstring"""




      super().tearDown()
      gc.collect()
      torch.cuda.empty_cache()
 def  UpperCAmelCase__						(						self   :			List[str]     )    ->					Union[str, Any]:


      """simple docstring"""




      __SCREAMING_SNAKE_CASE      = load_numpy(
          """https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/stable_unclip/stable_unclip_2_1_l_anime_turtle_fp16.npy"""     )

      __SCREAMING_SNAKE_CASE      = StableUnCLIPPipeline.from_pretrained("""fusing/stable-unclip-2-1-l"""   ,				torch_dtype=torch.floataa     )
      pipe.to(_A     )
      pipe.set_progress_bar_config(disable=_A     )
      # stable unclip will oom when integration tests are run on a V100,
      # so turn on memory savings
      pipe.enable_attention_slicing()
      pipe.enable_sequential_cpu_offload()

      __SCREAMING_SNAKE_CASE      = torch.Generator(device="""cpu"""     ).manual_seed(0     )
      __SCREAMING_SNAKE_CASE      = pipe("""anime turle"""   ,				generator=_A   ,				output_type="""np"""     )

      __SCREAMING_SNAKE_CASE      = output.images[0]

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

      assert_mean_pixel_difference(_A   ,				_A     )


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


      """simple docstring"""




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

      __SCREAMING_SNAKE_CASE      = StableUnCLIPPipeline.from_pretrained("""fusing/stable-unclip-2-1-l"""   ,				torch_dtype=torch.floataa     )
      __SCREAMING_SNAKE_CASE      = pipe.to(_A     )
      pipe.set_progress_bar_config(disable=_A     )
      pipe.enable_attention_slicing()
      pipe.enable_sequential_cpu_offload()

      __SCREAMING_SNAKE_CASE      = pipe(
          """anime turtle"""   ,				prior_num_inference_steps=2   ,				num_inference_steps=2   ,				output_type="""np"""   ,				)

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







def 					_snake_case	(  __snake_case = 100						):
					_UpperCamelCase			     =     (n * (n + 1) // 2) ** 2
					_UpperCamelCase			     =     n * (n + 1) * (2 * n + 1) // 6
					return sum_cubes - sum_squares


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