

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 _UpperCamelCase				(				A__				,						A__				,						A__				,						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				:			List[Any]						)				->							List[str]:

				"""simple docstring"""

				SCREAMING_SNAKE_CASE :   Optional[int]					  =							32
				SCREAMING_SNAKE_CASE :   List[Any]					  =							embedder_hidden_size

				# prior components

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

				torch.manual_seed(0						)
				SCREAMING_SNAKE_CASE :   Any					  =							CLIPTextModelWithProjection(
				    CLIPTextConfig(
				        bos_token_id=0   ,   eos_token_id=2   ,   hidden_size=__snake_case   ,   projection_dim=__snake_case   ,   intermediate_size=37   ,   layer_norm_eps=1e-05   ,   num_attention_heads=4   ,   num_hidden_layers=5   ,   pad_token_id=1   ,   vocab_size=1000   ,   )						)

				torch.manual_seed(0						)
				SCREAMING_SNAKE_CASE :   Dict					  =							PriorTransformer(
				    num_attention_heads=2   ,   attention_head_dim=12   ,   embedding_dim=__snake_case   ,   num_layers=1   ,   )

				torch.manual_seed(0						)
				SCREAMING_SNAKE_CASE :   List[Any]					  =							DDPMScheduler(
				    variance_type="fixed_small_log"   ,   prediction_type="sample"   ,   num_train_timesteps=1000   ,   clip_sample=__snake_case   ,   clip_sample_range=5.0   ,   beta_schedule="squaredcos_cap_v2"   ,   )

				# regular denoising components

				torch.manual_seed(0						)
				SCREAMING_SNAKE_CASE :   Union[str, Any]					  =							StableUnCLIPImageNormalizer(embedding_dim=__snake_case						)
				SCREAMING_SNAKE_CASE :   List[Any]					  =							DDPMScheduler(beta_schedule="squaredcos_cap_v2"						)

				torch.manual_seed(0						)
				SCREAMING_SNAKE_CASE :   Union[str, Any]					  =							CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip"						)

				torch.manual_seed(0						)
				SCREAMING_SNAKE_CASE :   int					  =							CLIPTextModel(
				    CLIPTextConfig(
				        bos_token_id=0   ,   eos_token_id=2   ,   hidden_size=__snake_case   ,   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=1000   ,   )						)

				torch.manual_seed(0						)
				SCREAMING_SNAKE_CASE :   str					  =							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=__snake_case   ,   layers_per_block=1   ,   upcast_attention=__snake_case   ,   use_linear_projection=__snake_case   ,   )

				torch.manual_seed(0						)
				SCREAMING_SNAKE_CASE :   List[Any]					  =							DDIMScheduler(
				    beta_schedule="scaled_linear"   ,   beta_start=0.0_0085   ,   beta_end=0.012   ,   prediction_type="v_prediction"   ,   set_alpha_to_one=__snake_case   ,   steps_offset=1   ,   )

				torch.manual_seed(0						)
				SCREAMING_SNAKE_CASE :   str					  =							AutoencoderKL()

				SCREAMING_SNAKE_CASE :   Optional[int]					  =							{
				    # 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				:			Union[str, Any]   ,   a				:			int   ,   a				:			List[Any]=0						)				->							List[Any]:

				"""simple docstring"""

				if str(__snake_case						).startswith("mps"						):
					SCREAMING_SNAKE_CASE :   Dict					  =							torch.manual_seed(__snake_case						)
				else:
					SCREAMING_SNAKE_CASE :   List[Any]					  =							torch.Generator(device=__snake_case						).manual_seed(__snake_case						)
				SCREAMING_SNAKE_CASE :   Union[str, Any]					  =							{
				    "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				:			Optional[Any]						)				->							Optional[int]:

				"""simple docstring"""

				SCREAMING_SNAKE_CASE :   Tuple					  =							torch_device == "cpu"

				self._test_attention_slicing_forward_pass(test_max_difference=__snake_case						)






			def 				__UpperCamelCase			(      self				:			List[Any]						)				->							List[Any]:

				"""simple docstring"""

				SCREAMING_SNAKE_CASE :   Union[str, Any]					  =							torch_device in ["cpu", "mps"]

				self._test_inference_batch_single_identical(test_max_difference=__snake_case						)






@slow
@require_torch_gpu
class _UpperCamelCase				(				unittest.TestCase ):
			'''simple docstring'''







			def 				__UpperCamelCase			(      self				:			Dict						)				->							Tuple:

				"""simple docstring"""

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

				"""simple docstring"""

				SCREAMING_SNAKE_CASE :   List[Any]					  =							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 :   List[str]					  =							StableUnCLIPPipeline.from_pretrained("fusing/stable-unclip-2-1-l"   ,   torch_dtype=torch.floataa						)
				pipe.to(__snake_case						)
				pipe.set_progress_bar_config(disable=__snake_case						)
				# 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 :   Tuple					  =							torch.Generator(device="cpu"						).manual_seed(0						)
				SCREAMING_SNAKE_CASE :   Any					  =							pipe("anime turle"   ,   generator=__snake_case   ,   output_type="np"						)

				SCREAMING_SNAKE_CASE :   Optional[int]					  =							output.images[0]

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

				assert_mean_pixel_difference(__snake_case   ,   __snake_case						)






			def 				__UpperCamelCase			(      self				:			int						)				->							str:

				"""simple docstring"""

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

				SCREAMING_SNAKE_CASE :   List[str]					  =							StableUnCLIPPipeline.from_pretrained("fusing/stable-unclip-2-1-l"   ,   torch_dtype=torch.floataa						)
				SCREAMING_SNAKE_CASE :   List[Any]					  =							pipe.to(__snake_case						)
				pipe.set_progress_bar_config(disable=__snake_case						)
				pipe.enable_attention_slicing()
				pipe.enable_sequential_cpu_offload()

				SCREAMING_SNAKE_CASE :   Any					  =							pipe(
				    "anime turtle"   ,   prior_num_inference_steps=2   ,   num_inference_steps=2   ,   output_type="np"   ,   )

				SCREAMING_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'''




from manim import *
class 			_lowerCAmelCase		(			A__	):






     """simple docstring"""






     def      lowerCAmelCase			(   self							:			Optional[Any]						)->	Union[str, Any]:
        snake_case 						=   Rectangle(height=0.5   ,       width=0.5						)
        snake_case 						=   Rectangle(height=0.25   ,       width=0.25						)
        snake_case 						=   Rectangle(height=0.46   ,       width=0.46						).set_stroke(width=0						)

        snake_case 						=   [mem.copy() for i in range(6						)]
        snake_case 						=   [mem.copy() for i in range(6						)]
        snake_case 						=   VGroup(*__snake_case						).arrange(__snake_case   ,       buff=0						)
        snake_case 						=   VGroup(*__snake_case						).arrange(__snake_case   ,       buff=0						)
        snake_case 						=   VGroup(__snake_case   ,       __snake_case						).arrange(__snake_case   ,       buff=0						)
        snake_case 						=   Text("""CPU"""   ,       font_size=24						)
        snake_case 						=   Group(__snake_case   ,       __snake_case						).arrange(__snake_case   ,       buff=0.5   ,       aligned_edge=__snake_case						)
        cpu.move_to([-2.5, -0.5, 0]						)
        self.add(__snake_case						)

        snake_case 						=   [mem.copy() for i in range(4						)]
        snake_case 						=   VGroup(*__snake_case						).arrange(__snake_case   ,       buff=0						)
        snake_case 						=   Text("""GPU"""   ,       font_size=24						)
        snake_case 						=   Group(__snake_case   ,       __snake_case						).arrange(__snake_case   ,       buff=0.5   ,       aligned_edge=__snake_case						)
        gpu.move_to([-1, -1, 0]						)
        self.add(__snake_case						)

        snake_case 						=   [mem.copy() for i in range(6						)]
        snake_case 						=   VGroup(*__snake_case						).arrange(__snake_case   ,       buff=0						)

        snake_case 						=   Text("""Model"""   ,       font_size=24						)
        snake_case 						=   Group(__snake_case   ,       __snake_case						).arrange(__snake_case   ,       buff=0.5   ,       aligned_edge=__snake_case						)
        model.move_to([3, -1.0, 0]						)
        self.add(__snake_case						)

        snake_case 						=   []
        snake_case 						=   []
        snake_case 						=   []

        for i, rect in enumerate(__snake_case						):
           rect.set_stroke(__snake_case						)

           snake_case 						=   Rectangle(height=0.46 / 4   ,       width=0.46 / 3						).set_stroke(width=0.0						).set_fill(__snake_case   ,       opacity=0.7						)

           if i == 0:
              cpu_target.next_to(cpu_left_col_base[0].get_corner(DOWN + LEFT						)   ,       buff=0.02   ,       direction=__snake_case						)
              cpu_target.set_x(cpu_target.get_x() + 0.1						)
           elif i == 3:
              cpu_target.next_to(model_cpu_arr[0]   ,       direction=__snake_case   ,       buff=0.0						)
           else:
              cpu_target.next_to(model_cpu_arr[i - 1]   ,       direction=__snake_case   ,       buff=0.0						)
           self.add(__snake_case						)
           model_cpu_arr.append(__snake_case						)

        self.add(*__snake_case   ,       *__snake_case   ,       *__snake_case						)

        snake_case 						=   [mem.copy() for i in range(6						)]
        snake_case 						=   VGroup(*__snake_case						).arrange(__snake_case   ,       buff=0						)

        snake_case 						=   Text("""Loaded Checkpoint"""   ,       font_size=24						)
        snake_case 						=   Group(__snake_case   ,       __snake_case						).arrange(__snake_case   ,       buff=0.5   ,       aligned_edge=__snake_case						)
        checkpoint.move_to([3, 0.5, 0]						)

        self.add(__snake_case						)

        snake_case 						=   []
        snake_case 						=   []

        for i, rect in enumerate(__snake_case						):
           snake_case 						=   fill.copy().set_fill(__snake_case   ,       opacity=0.7						)
           target.move_to(__snake_case						)
           ckpt_arr.append(__snake_case						)

           snake_case 						=   target.copy()
           if i < 5:
              cpu_target.move_to(cpu_left_col_base[i + 1]						)
           else:
              cpu_target.move_to(cpu_right_col_base[i - 5]						)
           ckpt_cpu_arr.append(__snake_case						)
        self.add(*__snake_case   ,       *__snake_case						)

        snake_case 						=   Square(side_length=2.2						)
        key.move_to([-5, 2, 0]						)

        snake_case 						=   MarkupText(
            f'''<b>Key:</b>\n\n<span fgcolor=\'{YELLOW}\'>●</span> Empty Model'''   ,       font_size=18   ,       )

        key_text.move_to([-5, 2.4, 0]						)

        self.add(__snake_case   ,       __snake_case						)

        snake_case 						=   MarkupText(
            f'''<span fgcolor=\'{BLUE}\'>●</span> Checkpoint'''   ,       font_size=18   ,       )

        blue_text.next_to(__snake_case   ,       DOWN * 2.4   ,       aligned_edge=key_text.get_left()						)
        self.add(__snake_case						)

        snake_case 						=   MarkupText(
            f'''Based on the passed in configuration, weights are stored in\na variety of np.memmaps on disk or to a particular device.'''   ,       font_size=24   ,       )
        step_a.move_to([2, 2, 0]						)

        snake_case 						=   [meta_mem.copy() for i in range(6						)]
        snake_case 						=   [meta_mem.copy() for i in range(6						)]
        snake_case 						=   VGroup(*__snake_case						).arrange(__snake_case   ,       buff=0						)
        snake_case 						=   VGroup(*__snake_case						).arrange(__snake_case   ,       buff=0						)
        snake_case 						=   VGroup(__snake_case   ,       __snake_case						).arrange(__snake_case   ,       buff=0						)
        snake_case 						=   Text("""Disk"""   ,       font_size=24						)
        snake_case 						=   Group(__snake_case   ,       __snake_case						).arrange(__snake_case   ,       buff=0.5   ,       aligned_edge=__snake_case						)
        disk.move_to([-4.0, -1.25, 0]						)
        self.play(Write(__snake_case   ,       run_time=3						)   ,       Write(__snake_case   ,       run_time=1						)   ,       Create(__snake_case   ,       run_time=1						)						)

        snake_case 						=   []
        for i, rect in enumerate(__snake_case						):
           snake_case 						=   rect.copy()
           target.generate_target()
           target.target.move_to(disk_left_col_base[i]						).scale(0.5						)
           animations.append(MoveToTarget(__snake_case   ,       run_time=1.5						)						)
        self.play(*__snake_case						)

        self.play(FadeOut(__snake_case						)						)

        snake_case 						=   MarkupText(f'''Then, the checkpoint is removed from memory\nthrough garbage collection.'''   ,       font_size=24						)
        step_a.move_to([2, 2, 0]						)

        self.play(Write(__snake_case   ,       run_time=3						)						)

        self.play(
            FadeOut(__snake_case   ,       __snake_case   ,       *__snake_case   ,       *__snake_case						)   ,       )

        self.wait()

