




import gc
import random
import unittest

import numpy as np
import torch
from transformers import CLIPImageProcessor, CLIPVisionConfig, CLIPVisionModel

from diffusers import HeunDiscreteScheduler, PriorTransformer, ShapEImgaImgPipeline
from diffusers.pipelines.shap_e import ShapERenderer
from diffusers.utils import floats_tensor, load_image, load_numpy, slow
from diffusers.utils.testing_utils import require_torch_gpu, torch_device

from ..test_pipelines_common import PipelineTesterMixin, assert_mean_pixel_difference




class       a__	( UpperCamelCase__       , unittest.TestCase				):
     a		:      Union[str, Any]		 =     ShapEImgaImgPipeline
     a		:      Any		 =     ["""image"""]
     a		:      Optional[int]		 =     ["""image"""]
     a		:      Tuple		 =     [
         """num_images_per_prompt""",
         """num_inference_steps""",
         """generator""",
         """latents""",
         """guidance_scale""",
         """frame_size""",
         """output_type""",
         """return_dict""",
     ]
     a		:      List[Any]		 =     False







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



          '''simple docstring'''



          return 32







     @property
     def       lowerCAmelCase_    (     self			)  ->   int:



          '''simple docstring'''



          return 32







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



          '''simple docstring'''



          return self.time_input_dim * 4







     @property
     def       lowerCAmelCase_    (     self			)  ->   Any:



          '''simple docstring'''



          return 8







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



          '''simple docstring'''



          torch.manual_seed(0			)
          a     =				CLIPVisionConfig(
              hidden_size=self.text_embedder_hidden_size  ,		image_size=64  ,		projection_dim=self.text_embedder_hidden_size  ,		intermediate_size=37  ,		num_attention_heads=4  ,		num_channels=3  ,		num_hidden_layers=5  ,		patch_size=1  ,		)

          a     =				CLIPVisionModel(A			)
          return model







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



          '''simple docstring'''



          a     =				CLIPImageProcessor(
              crop_size=224  ,		do_center_crop=A  ,		do_normalize=A  ,		do_resize=A  ,		image_mean=[0.4_8_1_4_5_4_6_6, 0.4_5_7_8_2_7_5, 0.4_0_8_2_1_0_7_3]  ,		image_std=[0.2_6_8_6_2_9_5_4, 0.2_6_1_3_0_2_5_8, 0.2_7_5_7_7_7_1_1]  ,		resample=3  ,		size=224  ,		)

          return image_processor







     @property
     def       lowerCAmelCase_    (     self			)  ->   str:



          '''simple docstring'''



          torch.manual_seed(0			)

          a     =				{
              "num_attention_heads": 2,
              "attention_head_dim": 16,
              "embedding_dim": self.time_input_dim,
              "num_embeddings": 32,
              "embedding_proj_dim": self.text_embedder_hidden_size,
              "time_embed_dim": self.time_embed_dim,
              "num_layers": 1,
              "clip_embed_dim": self.time_input_dim * 2,
              "additional_embeddings": 0,
              "time_embed_act_fn": "gelu",
              "norm_in_type": "layer",
              "embedding_proj_norm_type": "layer",
              "encoder_hid_proj_type": None,
              "added_emb_type": None,
          }

          a     =				PriorTransformer(**A			)
          return model







     @property
     def       lowerCAmelCase_    (     self			)  ->   Tuple:



          '''simple docstring'''



          torch.manual_seed(0			)

          a     =				{
              "param_shapes": (
                  (self.renderer_dim, 93),
                  (self.renderer_dim, 8),
                  (self.renderer_dim, 8),
                  (self.renderer_dim, 8),
              ),
              "d_latent": self.time_input_dim,
              "d_hidden": self.renderer_dim,
              "n_output": 12,
              "background": (
                  0.1,
                  0.1,
                  0.1,
              ),
          }
          a     =				ShapERenderer(**A			)
          return model







     def       lowerCAmelCase_    (     self			)  ->   Any:



          '''simple docstring'''



          a     =				self.dummy_prior
          a     =				self.dummy_image_encoder
          a     =				self.dummy_image_processor
          a     =				self.dummy_renderer

          a     =				HeunDiscreteScheduler(
              beta_schedule="exp"  ,		num_train_timesteps=1024  ,		prediction_type="sample"  ,		use_karras_sigmas=A  ,		clip_sample=A  ,		clip_sample_range=1.0  ,		)
          a     =				{
              "prior": prior,
              "image_encoder": image_encoder,
              "image_processor": image_processor,
              "renderer": renderer,
              "scheduler": scheduler,
          }

          return components







     def       lowerCAmelCase_    (     self  ,		A  ,		A=0			)  ->   List[Any]:



          '''simple docstring'''



          a     =				floats_tensor((1, 3, 64, 64)  ,		rng=random.Random(A			)			).to(A			)

          if str(A			).startswith("mps"			):
               a     =				torch.manual_seed(A			)
          else:
               a     =				torch.Generator(device=A			).manual_seed(A			)
          a     =				{
              "image": input_image,
              "generator": generator,
              "num_inference_steps": 1,
              "frame_size": 32,
              "output_type": "np",
          }
          return inputs







     def       lowerCAmelCase_    (     self			)  ->   Tuple:



          '''simple docstring'''



          a     =				"cpu"

          a     =				self.get_dummy_components()

          a     =				self.pipeline_class(**A			)
          a     =				pipe.to(A			)

          pipe.set_progress_bar_config(disable=A			)

          a     =				pipe(**self.get_dummy_inputs(A			)			)
          a     =				output.images[0]
          a     =				image[0, -3:, -3:, -1]

          assert image.shape == (20, 32, 32, 3)

          a     =				np.array(
              [
                  0.0_0_0_3_9_2_1_6,
                  0.0_0_0_3_9_2_1_6,
                  0.0_0_0_3_9_2_1_6,
                  0.0_0_0_3_9_2_1_6,
                  0.0_0_0_3_9_2_1_6,
                  0.0_0_0_3_9_2_1_6,
                  0.0_0_0_3_9_2_1_6,
                  0.0_0_0_3_9_2_1_6,
                  0.0_0_0_3_9_2_1_6,
              ]			)

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







     def       lowerCAmelCase_    (     self			)  ->   Any:



          '''simple docstring'''



          self._test_inference_batch_consistent(batch_sizes=[1, 2]			)







     def       lowerCAmelCase_    (     self			)  ->   Optional[Any]:



          '''simple docstring'''



          a     =				torch_device == "cpu"
          a     =				True
          self._test_inference_batch_single_identical(
              batch_size=2  ,		test_max_difference=A  ,		relax_max_difference=A  ,		)







     def       lowerCAmelCase_    (     self			)  ->   Any:



          '''simple docstring'''



          a     =				self.get_dummy_components()
          a     =				self.pipeline_class(**A			)
          a     =				pipe.to(A			)
          pipe.set_progress_bar_config(disable=A			)

          a     =				1
          a     =				2

          a     =				self.get_dummy_inputs(A			)

          for key in inputs.keys():
               if key in self.batch_params:
                    a     =				batch_size * [inputs[key]]

          a     =				pipe(**A  ,		num_images_per_prompt=A			)[0]

          assert images.shape[0] == batch_size * num_images_per_prompt







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







     def       lowerCAmelCase_    (     self			)  ->   Optional[Any]:



          '''simple docstring'''



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







     def       lowerCAmelCase_    (     self			)  ->   str:



          '''simple docstring'''



          a     =				load_image(
              "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main" "/shap_e/corgi.png"			)
          a     =				load_numpy(
              "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main"
              "/shap_e/test_shap_e_img2img_out.npy"			)
          a     =				ShapEImgaImgPipeline.from_pretrained("openai/shap-e-img2img"			)
          a     =				pipe.to(A			)
          pipe.set_progress_bar_config(disable=A			)

          a     =				torch.Generator(device=A			).manual_seed(0			)

          a     =				pipe(
              A  ,		generator=A  ,		guidance_scale=3.0  ,		num_inference_steps=64  ,		frame_size=64  ,		output_type="np"  ,		).images[0]

          assert images.shape == (20, 64, 64, 3)

          assert_mean_pixel_difference(A  ,		A			)







from manim import *




class       a__	( UpperCamelCase__				):







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



          '''simple docstring'''



          a     =				Rectangle(height=0.5  ,		width=0.5			)
          a     =				Rectangle(height=0.4_6  ,		width=0.4_6			).set_stroke(width=0			)

          a     =				Rectangle(height=0.2_5  ,		width=0.2_5			)

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

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

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

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

          a     =				[]
          a     =				[]

          for i, rect in enumerate(A			):
               a     =				fill.copy().set_fill(A  ,		opacity=0.8			)
               target.move_to(A			)
               model_arr.append(A			)

               a     =				Rectangle(height=0.4_6  ,		width=0.4_6			).set_stroke(width=0.0			).set_fill(A  ,		opacity=0.8			)
               cpu_target.move_to(cpu_left_col_base[i]			)
               model_cpu_arr.append(A			)

          self.add(*A  ,		*A			)

          a     =				[meta_mem.copy() for i in range(6			)]
          a     =				[meta_mem.copy() for i in range(6			)]
          a     =				VGroup(*A			).arrange(A  ,		buff=0			)
          a     =				VGroup(*A			).arrange(A  ,		buff=0			)
          a     =				VGroup(A  ,		A			).arrange(A  ,		buff=0			)
          a     =				Text("Disk"  ,		font_size=24			)
          a     =				Group(A  ,		A			).arrange(A  ,		buff=0.5  ,		aligned_edge=A			)
          disk.move_to([-4, -1.2_5, 0]			)
          self.add(A  ,		A			)

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

          a     =				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(A  ,		A			)

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

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

          a     =				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(A			)			)

          a     =				Square(0.3			)
          input.set_fill(A  ,		opacity=1.0			)
          input.set_stroke(width=0.0			)
          input.next_to(model_base[0]  ,		A  ,		buff=0.5			)

          self.play(Write(A			)			)

          input.generate_target()
          input.target.next_to(model_arr[0]  ,		direction=A  ,		buff=0.0_2			)
          self.play(MoveToTarget(A			)			)

          self.play(FadeOut(A			)			)

          a     =				Arrow(start=A  ,		end=A  ,		color=A  ,		buff=0.5			)
          a.next_to(model_arr[0].get_left()  ,		A  ,		buff=0.2			)

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

          a     =				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(A  ,		run_time=3			)			)

          a     =				{"run_time": 1, "fade_in": True, "fade_out": True, "buff": 0.0_2}

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

          a     =				a.copy()
          for i in range(6			):
               a_c.next_to(model_arr[i].get_right() + 0.0_2  ,		A  ,		buff=0.2			)

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

               a     =				AnimationGroup(
                   FadeOut(A  ,		run_time=0.5			)  ,		MoveToTarget(A  ,		run_time=0.5			)  ,		FadeIn(A  ,		run_time=0.5			)  ,		lag_ratio=0.2			)

               self.play(A			)

               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:
                         a     =				0.7

                    self.play(
                        Circumscribe(model_arr[i]  ,		**A			)  ,		Circumscribe(cpu_left_col_base[i]  ,		**A			)  ,		Circumscribe(cpu_left_col_base[i + 1]  ,		color=A  ,		**A			)  ,		Circumscribe(gpu_rect[0]  ,		color=A  ,		**A			)  ,		Circumscribe(model_arr[i + 1]  ,		color=A  ,		**A			)  ,		)
                    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.0_2  ,		buff=0.2			)

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

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

               a     =				a_c
               a     =				a_c.copy()

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

          a     =				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(A  ,		run_time=3			)  ,		MoveToTarget(A			)			)

          self.wait()

