


import gc
import unittest

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

from diffusers import AutoencoderKL, DDIMScheduler, LDMTextToImagePipeline, UNetaDConditionModel
from diffusers.utils.testing_utils import (
    enable_full_determinism,
    load_numpy,
    nightly,
    require_torch_gpu,
    slow,
    torch_device,
)

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


enable_full_determinism()

class 			UpperCamelCase__			(			UpperCAmelCase_				,    unittest.TestCase      ):
      a__ :  Union[str, Any]		    =    LDMTextToImagePipeline
      a__ :  Tuple		    =    TEXT_TO_IMAGE_PARAMS - {
          '''negative_prompt''',
          '''negative_prompt_embeds''',
          '''cross_attention_kwargs''',
          '''prompt_embeds''',
      }
      a__ :  Optional[Any]		    =    PipelineTesterMixin.required_optional_params - {
          '''num_images_per_prompt''',
          '''callback''',
          '''callback_steps''',
      }
      a__ :  Dict		    =    TEXT_TO_IMAGE_BATCH_PARAMS
      a__ :  Tuple		    =    False






      def 						__lowercase(  self       : Optional[Any] )	->			int:
           torch.manual_seed(0 )
           UpperCamelCase__				:	Any										=  UNetaDConditionModel(
               block_out_channels=(32, 64),		layers_per_block=2,		sample_size=32,		in_channels=4,		out_channels=4,		down_block_types=('''DownBlock2D''', '''CrossAttnDownBlock2D'''),		up_block_types=('''CrossAttnUpBlock2D''', '''UpBlock2D'''),		cross_attention_dim=32,		)
           UpperCamelCase__				:	Union[str, Any]										=  DDIMScheduler(
               beta_start=0.0_0085,		beta_end=0.012,		beta_schedule='''scaled_linear''',		clip_sample=lowerCamelCase__,		set_alpha_to_one=lowerCamelCase__,		)
           torch.manual_seed(0 )
           UpperCamelCase__				:	Any										=  AutoencoderKL(
               block_out_channels=(32, 64),		in_channels=3,		out_channels=3,		down_block_types=('''DownEncoderBlock2D''', '''DownEncoderBlock2D'''),		up_block_types=('''UpDecoderBlock2D''', '''UpDecoderBlock2D'''),		latent_channels=4,		)
           torch.manual_seed(0 )
           UpperCamelCase__				:	Union[str, Any]										=  CLIPTextConfig(
               bos_token_id=0,		eos_token_id=2,		hidden_size=32,		intermediate_size=37,		layer_norm_eps=1e-05,		num_attention_heads=4,		num_hidden_layers=5,		pad_token_id=1,		vocab_size=10_00,		)
           UpperCamelCase__				:	List[str]										=  CLIPTextModel(lowerCamelCase__ )
           UpperCamelCase__				:	List[str]										=  CLIPTokenizer.from_pretrained('''hf-internal-testing/tiny-random-clip''' )

           UpperCamelCase__				:	Any										=  {
               """unet""": unet,
               """scheduler""": scheduler,
               """vqvae""": vae,
               """bert""": text_encoder,
               """tokenizer""": tokenizer,
           }
           return components






      def 						__lowercase(  self       : Tuple,		__lowerCamelCase       : str,		__lowerCamelCase       : List[str]=0 )	->			str:

           if str(lowerCamelCase__ ).startswith('''mps''' ):
                UpperCamelCase__				:	str										=  torch.manual_seed(lowerCamelCase__ )
           else:
                UpperCamelCase__				:	Optional[Any]										=  torch.Generator(device=lowerCamelCase__ ).manual_seed(lowerCamelCase__ )
           UpperCamelCase__				:	Tuple										=  {
               """prompt""": """A painting of a squirrel eating a burger""",
               """generator""": generator,
               """num_inference_steps""": 2,
               """guidance_scale""": 6.0,
               """output_type""": """numpy""",
           }
           return inputs






      def 						__lowercase(  self       : Tuple )	->			Tuple:
           UpperCamelCase__				:	Dict										=  """cpu"""  # ensure determinism for the device-dependent torch.Generator

           UpperCamelCase__				:	List[str]										=  self.get_dummy_components()
           UpperCamelCase__				:	Optional[int]										=  LDMTextToImagePipeline(**lowerCamelCase__ )
           pipe.to(lowerCamelCase__ )
           pipe.set_progress_bar_config(disable=lowerCamelCase__ )

           UpperCamelCase__				:	List[Any]										=  self.get_dummy_inputs(lowerCamelCase__ )
           UpperCamelCase__				:	List[Any]										=  pipe(**lowerCamelCase__ ).images
           UpperCamelCase__				:	Tuple										=  image[0, -3:, -3:, -1]

           assert image.shape == (1, 16, 16, 3)
           UpperCamelCase__				:	int										=  np.array([0.6101, 0.6156, 0.5622, 0.4895, 0.6661, 0.3804, 0.5748, 0.6136, 0.5014] )

           assert np.abs(image_slice.flatten() - expected_slice ).max() < 1e-3

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






      def 						__lowercase(  self       : Any )	->			str:
           super().tearDown()
           gc.collect()
           torch.cuda.empty_cache()






      def 						__lowercase(  self       : Dict,		__lowerCamelCase       : List[Any],		__lowerCamelCase       : Tuple=torch.floataa,		__lowerCamelCase       : List[Any]=0 )	->			Union[str, Any]:
           UpperCamelCase__				:	Union[str, Any]										=  torch.manual_seed(lowerCamelCase__ )
           UpperCamelCase__				:	Tuple										=  np.random.RandomState(lowerCamelCase__ ).standard_normal((1, 4, 32, 32) )
           UpperCamelCase__				:	Any										=  torch.from_numpy(lowerCamelCase__ ).to(device=lowerCamelCase__,		dtype=lowerCamelCase__ )
           UpperCamelCase__				:	Optional[int]										=  {
               """prompt""": """A painting of a squirrel eating a burger""",
               """latents""": latents,
               """generator""": generator,
               """num_inference_steps""": 3,
               """guidance_scale""": 6.0,
               """output_type""": """numpy""",
           }
           return inputs






      def 						__lowercase(  self       : Optional[int] )	->			Dict:
           UpperCamelCase__				:	Any										=  LDMTextToImagePipeline.from_pretrained('''CompVis/ldm-text2im-large-256''' ).to(lowerCamelCase__ )
           pipe.set_progress_bar_config(disable=lowerCamelCase__ )

           UpperCamelCase__				:	Any										=  self.get_inputs(lowerCamelCase__ )
           UpperCamelCase__				:	Tuple										=  pipe(**lowerCamelCase__ ).images
           UpperCamelCase__				:	Optional[int]										=  image[0, -3:, -3:, -1].flatten()

           assert image.shape == (1, 2_56, 2_56, 3)
           UpperCamelCase__				:	Optional[Any]										=  np.array([0.5_1825, 0.5_2850, 0.5_2543, 0.5_4258, 0.5_2304, 0.5_2569, 0.5_4363, 0.5_5276, 0.5_6878] )
           UpperCamelCase__				:	Union[str, Any]										=  np.abs(expected_slice - image_slice ).max()
           assert max_diff < 1e-3

@nightly
@require_torch_gpu
class 			UpperCamelCase__			(			unittest.TestCase      ):






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






      def 						__lowercase(  self       : Tuple,		__lowerCamelCase       : str,		__lowerCamelCase       : Union[str, Any]=torch.floataa,		__lowerCamelCase       : Tuple=0 )	->			List[Any]:
           UpperCamelCase__				:	Union[str, Any]										=  torch.manual_seed(lowerCamelCase__ )
           UpperCamelCase__				:	Optional[Any]										=  np.random.RandomState(lowerCamelCase__ ).standard_normal((1, 4, 32, 32) )
           UpperCamelCase__				:	List[Any]										=  torch.from_numpy(lowerCamelCase__ ).to(device=lowerCamelCase__,		dtype=lowerCamelCase__ )
           UpperCamelCase__				:	Union[str, Any]										=  {
               """prompt""": """A painting of a squirrel eating a burger""",
               """latents""": latents,
               """generator""": generator,
               """num_inference_steps""": 50,
               """guidance_scale""": 6.0,
               """output_type""": """numpy""",
           }
           return inputs






      def 						__lowercase(  self       : Tuple )	->			Tuple:
           UpperCamelCase__				:	Tuple										=  LDMTextToImagePipeline.from_pretrained('''CompVis/ldm-text2im-large-256''' ).to(lowerCamelCase__ )
           pipe.set_progress_bar_config(disable=lowerCamelCase__ )

           UpperCamelCase__				:	Dict										=  self.get_inputs(lowerCamelCase__ )
           UpperCamelCase__				:	List[Any]										=  pipe(**lowerCamelCase__ ).images[0]

           UpperCamelCase__				:	Optional[Any]										=  load_numpy(
               '''https://huggingface.co/datasets/diffusers/test-arrays/resolve/main/ldm_text2img/ldm_large_256_ddim.npy''' )
           UpperCamelCase__				:	int										=  np.abs(expected_image - image ).max()
           assert max_diff < 1e-3












from manim import *







class 		_A						( UpperCAmelCase_ ):

  def    a					(    self						:	Dict			):
     """simple docstring"""

     __UpperCamelCase	:			Tuple					=					Rectangle(height=0.5   ,				width=0.5			)
     __UpperCamelCase	:			Tuple					=					Rectangle(height=0.46   ,				width=0.46			).set_stroke(width=0			)

     __UpperCamelCase	:			Optional[Any]					=					Rectangle(height=0.25   ,				width=0.25			)

     __UpperCamelCase	:			Any					=					[mem.copy() for i in range(6			)]
     __UpperCamelCase	:			str					=					[mem.copy() for i in range(6			)]
     __UpperCamelCase	:			int					=					VGroup(*lowerCamelCase__			).arrange(lowerCamelCase__   ,				buff=0			)
     __UpperCamelCase	:			Optional[int]					=					VGroup(*lowerCamelCase__			).arrange(lowerCamelCase__   ,				buff=0			)
     __UpperCamelCase	:			Dict					=					VGroup(lowerCamelCase__   ,				lowerCamelCase__			).arrange(lowerCamelCase__   ,				buff=0			)
     __UpperCamelCase	:			List[Any]					=					Text("""CPU"""   ,				font_size=24			)
     __UpperCamelCase	:			Any					=					Group(lowerCamelCase__   ,				lowerCamelCase__			).arrange(lowerCamelCase__   ,				buff=0.5   ,				aligned_edge=lowerCamelCase__			)
     cpu.move_to([-2.5, -0.5, 0]			)
     self.add(lowerCamelCase__			)

     __UpperCamelCase	:			Tuple					=					[mem.copy() for i in range(4			)]
     __UpperCamelCase	:			Optional[Any]					=					VGroup(*lowerCamelCase__			).arrange(lowerCamelCase__   ,				buff=0			)
     __UpperCamelCase	:			Optional[Any]					=					Text("""GPU"""   ,				font_size=24			)
     __UpperCamelCase	:			Union[str, Any]					=					Group(lowerCamelCase__   ,				lowerCamelCase__			).arrange(lowerCamelCase__   ,				buff=0.5   ,				aligned_edge=lowerCamelCase__			)
     gpu.move_to([-1, -1, 0]			)
     self.add(lowerCamelCase__			)

     __UpperCamelCase	:			Dict					=					[mem.copy() for i in range(6			)]
     __UpperCamelCase	:			Tuple					=					VGroup(*lowerCamelCase__			).arrange(lowerCamelCase__   ,				buff=0			)

     __UpperCamelCase	:			Tuple					=					Text("""Model"""   ,				font_size=24			)
     __UpperCamelCase	:			List[str]					=					Group(lowerCamelCase__   ,				lowerCamelCase__			).arrange(lowerCamelCase__   ,				buff=0.5   ,				aligned_edge=lowerCamelCase__			)
     model.move_to([3, -1.0, 0]			)
     self.add(lowerCamelCase__			)

     __UpperCamelCase	:			Tuple					=					[]
     __UpperCamelCase	:			str					=					[]

     for i, rect in enumerate(lowerCamelCase__			):
        __UpperCamelCase	:			Optional[Any]					=					fill.copy().set_fill(lowerCamelCase__   ,				opacity=0.8			)
        target.move_to(lowerCamelCase__			)
        model_arr.append(lowerCamelCase__			)

        __UpperCamelCase	:			Tuple					=					Rectangle(height=0.46   ,				width=0.46			).set_stroke(width=0.0			).set_fill(lowerCamelCase__   ,				opacity=0.8			)
        cpu_target.move_to(cpu_left_col_base[i]			)
        model_cpu_arr.append(lowerCamelCase__			)

     self.add(*lowerCamelCase__   ,				*lowerCamelCase__			)

     __UpperCamelCase	:			Union[str, Any]					=					[meta_mem.copy() for i in range(6			)]
     __UpperCamelCase	:			Any					=					[meta_mem.copy() for i in range(6			)]
     __UpperCamelCase	:			int					=					VGroup(*lowerCamelCase__			).arrange(lowerCamelCase__   ,				buff=0			)
     __UpperCamelCase	:			Union[str, Any]					=					VGroup(*lowerCamelCase__			).arrange(lowerCamelCase__   ,				buff=0			)
     __UpperCamelCase	:			str					=					VGroup(lowerCamelCase__   ,				lowerCamelCase__			).arrange(lowerCamelCase__   ,				buff=0			)
     __UpperCamelCase	:			Optional[Any]					=					Text("""Disk"""   ,				font_size=24			)
     __UpperCamelCase	:			Optional[Any]					=					Group(lowerCamelCase__   ,				lowerCamelCase__			).arrange(lowerCamelCase__   ,				buff=0.5   ,				aligned_edge=lowerCamelCase__			)
     disk.move_to([-4, -1.25, 0]			)
     self.add(lowerCamelCase__   ,				lowerCamelCase__			)

     __UpperCamelCase	:			Union[str, Any]					=					Square(side_length=2.2			)
     key.move_to([-5, 2, 0]			)

     __UpperCamelCase	:			Optional[Any]					=					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(lowerCamelCase__   ,				lowerCamelCase__			)

     __UpperCamelCase	:			Optional[int]					=					MarkupText(
         f'<span fgcolor=\'{BLUE}\'>●</span> Checkpoint'   ,				font_size=18   ,				)

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

     __UpperCamelCase	:			Any					=					MarkupText(
         f'Now watch as an input is passed through the model\nand how the memory is utilized and handled.'   ,				font_size=24   ,				)
     step_a.move_to([2, 2, 0]			)

     self.play(Write(lowerCamelCase__			)			)

     __UpperCamelCase	:			List[str]					=					Square(0.3			)
     input.set_fill(lowerCamelCase__   ,				opacity=1.0			)
     input.set_stroke(width=0.0			)
     input.next_to(model_base[0]   ,				lowerCamelCase__   ,				buff=0.5			)

     self.play(Write(lowerCamelCase__			)			)

     input.generate_target()
     input.target.next_to(model_arr[0]   ,				direction=lowerCamelCase__   ,				buff=0.02			)
     self.play(MoveToTarget(lowerCamelCase__			)			)

     self.play(FadeOut(lowerCamelCase__			)			)

     __UpperCamelCase	:			Dict					=					Arrow(start=lowerCamelCase__   ,				end=lowerCamelCase__   ,				color=lowerCamelCase__   ,				buff=0.5			)
     a.next_to(model_arr[0].get_left()   ,				lowerCamelCase__   ,				buff=0.2			)

     model_cpu_arr[0].generate_target()
     model_cpu_arr[0].target.move_to(gpu_rect[0]			)

     __UpperCamelCase	:			int					=					MarkupText(
         f'As the input reaches a layer, the hook triggers\nand weights are moved from the CPU\nto the GPU and back.'   ,				font_size=24   ,				)
     step_a.move_to([2, 2, 0]			)

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

     __UpperCamelCase	:			Dict					=					{"""run_time""": 1, """fade_in""": True, """fade_out""": True, """buff""": 0.02}

     self.play(
         Write(lowerCamelCase__			)   ,				Circumscribe(model_arr[0]   ,				color=lowerCamelCase__   ,				**lowerCamelCase__			)   ,				Circumscribe(model_cpu_arr[0]   ,				color=lowerCamelCase__   ,				**lowerCamelCase__			)   ,				Circumscribe(gpu_rect[0]   ,				color=lowerCamelCase__   ,				**lowerCamelCase__			)   ,				)
     self.play(MoveToTarget(model_cpu_arr[0]			)			)

     __UpperCamelCase	:			int					=					a.copy()
     for i in range(6			):
        a_c.next_to(model_arr[i].get_right() + 0.02   ,				lowerCamelCase__   ,				buff=0.2			)

        input.generate_target()
        input.target.move_to(model_arr[i].get_right() + 0.02			)

        __UpperCamelCase	:			Tuple					=					AnimationGroup(
            FadeOut(lowerCamelCase__   ,				run_time=0.5			)   ,				MoveToTarget(lowerCamelCase__   ,				run_time=0.5			)   ,				FadeIn(lowerCamelCase__   ,				run_time=0.5			)   ,				lag_ratio=0.2			)

        self.play(lowerCamelCase__			)

        model_cpu_arr[i].generate_target()
        model_cpu_arr[i].target.move_to(cpu_left_col_base[i]			)

        if i < 5:
           model_cpu_arr[i + 1].generate_target()
           model_cpu_arr[i + 1].target.move_to(gpu_rect[0]			)
           if i >= 1:
              __UpperCamelCase	:			Any					=					0.7

           self.play(
               Circumscribe(model_arr[i]   ,				**lowerCamelCase__			)   ,				Circumscribe(cpu_left_col_base[i]   ,				**lowerCamelCase__			)   ,				Circumscribe(cpu_left_col_base[i + 1]   ,				color=lowerCamelCase__   ,				**lowerCamelCase__			)   ,				Circumscribe(gpu_rect[0]   ,				color=lowerCamelCase__   ,				**lowerCamelCase__			)   ,				Circumscribe(model_arr[i + 1]   ,				color=lowerCamelCase__   ,				**lowerCamelCase__			)   ,				)
           if i < 1:
              self.play(
                  MoveToTarget(model_cpu_arr[i]			)   ,				MoveToTarget(model_cpu_arr[i + 1]			)   ,				)
           else:
              self.play(
                  MoveToTarget(model_cpu_arr[i]   ,				run_time=0.7			)   ,				MoveToTarget(model_cpu_arr[i + 1]   ,				run_time=0.7			)   ,				)
        else:
           model_cpu_arr[i].generate_target()
           model_cpu_arr[i].target.move_to(cpu_left_col_base[-1]			)
           input.generate_target()
           input.target.next_to(model_arr[-1].get_right()   ,				RIGHT + 0.02   ,				buff=0.2			)

           self.play(
               Circumscribe(model_arr[-1]   ,				color=lowerCamelCase__   ,				**lowerCamelCase__			)   ,				Circumscribe(cpu_left_col_base[-1]   ,				color=lowerCamelCase__   ,				**lowerCamelCase__			)   ,				Circumscribe(gpu_rect[0]   ,				color=lowerCamelCase__   ,				**lowerCamelCase__			)   ,				)

           self.play(MoveToTarget(model_cpu_arr[i]			)			)

        __UpperCamelCase	:			str					=					a_c
        __UpperCamelCase	:			List[Any]					=					a_c.copy()

     input.generate_target()
     input.target.next_to(model_base[-1]   ,				RIGHT + 0.02   ,				buff=0.5			)
     self.play(
         FadeOut(lowerCamelCase__			)   ,				FadeOut(lowerCamelCase__   ,				run_time=0.5			)   ,				)

     __UpperCamelCase	:			List[str]					=					MarkupText(f'Inference on a model too large for GPU memory\nis successfully completed.'   ,				font_size=24			)
     step_a.move_to([2, 2, 0]			)

     self.play(Write(lowerCamelCase__   ,				run_time=3			)   ,				MoveToTarget(lowerCamelCase__			)			)

     self.wait()



