'''simple docstring'''


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 					__SCREAMING_SNAKE_CASE						(	lowerCamelCase_						,     unittest.TestCase	):







  '''simple docstring'''






  lowerCamelCase_		:str						   =		ShapEImgaImgPipeline
  lowerCamelCase_		:Dict						   =		['''image''']
  lowerCamelCase_		:int						   =		['''image''']
  lowerCamelCase_		:str						   =		[
      '''num_images_per_prompt''',
      '''num_inference_steps''',
      '''generator''',
      '''latents''',
      '''guidance_scale''',
      '''frame_size''',
      '''output_type''',
      '''return_dict''',
  ]
  lowerCamelCase_		:Tuple						   =		False


  @property
  def _UpperCamelCase					(			self						):




       '''simple docstring'''

       return 3_2


  @property
  def _UpperCamelCase					(			self						):




       '''simple docstring'''

       return 3_2


  @property
  def _UpperCamelCase					(			self						):




       '''simple docstring'''

       return self.time_input_dim * 4


  @property
  def _UpperCamelCase					(			self						):




       '''simple docstring'''

       return 8


  @property
  def _UpperCamelCase					(			self						):




       '''simple docstring'''

       torch.manual_seed(0						)
       UpperCAmelCase_							:       str       =				CLIPVisionConfig(
           hidden_size=self.text_embedder_hidden_size				,		image_size=6_4				,		projection_dim=self.text_embedder_hidden_size				,		intermediate_size=3_7				,		num_attention_heads=4				,		num_channels=3				,		num_hidden_layers=5				,		patch_size=1				,		)

       UpperCAmelCase_							:       Optional[Any]       =				CLIPVisionModel(lowerCamelCase_						)
       return model


  @property
  def _UpperCamelCase					(			self						):




       '''simple docstring'''

       UpperCAmelCase_							:       Union[str, Any]       =				CLIPImageProcessor(
           crop_size=2_2_4				,		do_center_crop=lowerCamelCase_				,		do_normalize=lowerCamelCase_				,		do_resize=lowerCamelCase_				,		image_mean=[0.48_14_54_66, 0.4_57_82_75, 0.40_82_10_73]				,		image_std=[0.26_86_29_54, 0.26_13_02_58, 0.27_57_77_11]				,		resample=3				,		size=2_2_4				,		)

       return image_processor


  @property
  def _UpperCamelCase					(			self						):




       '''simple docstring'''

       torch.manual_seed(0						)

       UpperCAmelCase_							:       List[Any]       =				{
           'num_attention_heads': 2,
           'attention_head_dim': 1_6,
           'embedding_dim': self.time_input_dim,
           'num_embeddings': 3_2,
           '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,
       }

       UpperCAmelCase_							:       List[str]       =				PriorTransformer(**lowerCamelCase_						)
       return model


  @property
  def _UpperCamelCase					(			self						):




       '''simple docstring'''

       torch.manual_seed(0						)

       UpperCAmelCase_							:       int       =				{
           'param_shapes': (
               (self.renderer_dim, 9_3),
               (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': 1_2,
           'background': (
               0.1,
               0.1,
               0.1,
           ),
       }
       UpperCAmelCase_							:       int       =				ShapERenderer(**lowerCamelCase_						)
       return model


  def _UpperCamelCase					(			self						):




       '''simple docstring'''

       UpperCAmelCase_							:       Tuple       =				self.dummy_prior
       UpperCAmelCase_							:       Optional[int]       =				self.dummy_image_encoder
       UpperCAmelCase_							:       str       =				self.dummy_image_processor
       UpperCAmelCase_							:       Any       =				self.dummy_renderer

       UpperCAmelCase_							:       Optional[int]       =				HeunDiscreteScheduler(
           beta_schedule='exp'				,		num_train_timesteps=1_0_2_4				,		prediction_type='sample'				,		use_karras_sigmas=lowerCamelCase_				,		clip_sample=lowerCamelCase_				,		clip_sample_range=1.0				,		)
       UpperCAmelCase_							:       Union[str, Any]       =				{
           'prior': prior,
           'image_encoder': image_encoder,
           'image_processor': image_processor,
           'renderer': renderer,
           'scheduler': scheduler,
       }

       return components


  def _UpperCamelCase					(			self				,		snake_case_				,		snake_case_=0						):




       '''simple docstring'''

       UpperCAmelCase_							:       Tuple       =				floats_tensor((1, 3, 6_4, 6_4)				,		rng=random.Random(lowerCamelCase_						)						).to(lowerCamelCase_						)

       if str(lowerCamelCase_						).startswith('mps'						):
            UpperCAmelCase_							:       List[Any]       =				torch.manual_seed(lowerCamelCase_						)
       else:
            UpperCAmelCase_							:       Any       =				torch.Generator(device=lowerCamelCase_						).manual_seed(lowerCamelCase_						)
       UpperCAmelCase_							:       List[Any]       =				{
           'image': input_image,
           'generator': generator,
           'num_inference_steps': 1,
           'frame_size': 3_2,
           'output_type': 'np',
       }
       return inputs


  def _UpperCamelCase					(			self						):




       '''simple docstring'''

       UpperCAmelCase_							:       Union[str, Any]       =				'cpu'

       UpperCAmelCase_							:       str       =				self.get_dummy_components()

       UpperCAmelCase_							:       List[str]       =				self.pipeline_class(**lowerCamelCase_						)
       UpperCAmelCase_							:       Tuple       =				pipe.to(lowerCamelCase_						)

       pipe.set_progress_bar_config(disable=lowerCamelCase_						)

       UpperCAmelCase_							:       str       =				pipe(**self.get_dummy_inputs(lowerCamelCase_						)						)
       UpperCAmelCase_							:       Tuple       =				output.images[0]
       UpperCAmelCase_							:       Optional[int]       =				image[0, -3:, -3:, -1]

       assert image.shape == (2_0, 3_2, 3_2, 3)

       UpperCAmelCase_							:       Union[str, Any]       =				np.array(
           [
               0.00_03_92_16,
               0.00_03_92_16,
               0.00_03_92_16,
               0.00_03_92_16,
               0.00_03_92_16,
               0.00_03_92_16,
               0.00_03_92_16,
               0.00_03_92_16,
               0.00_03_92_16,
           ]						)

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


  def _UpperCamelCase					(			self						):




       '''simple docstring'''

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


  def _UpperCamelCase					(			self						):




       '''simple docstring'''

       UpperCAmelCase_							:       Optional[Any]       =				torch_device == 'cpu'
       UpperCAmelCase_							:       Dict       =				True
       self._test_inference_batch_single_identical(
           batch_size=2				,		test_max_difference=lowerCamelCase_				,		relax_max_difference=lowerCamelCase_				,		)






  def _UpperCamelCase					(			self						):




       '''simple docstring'''

       UpperCAmelCase_							:       Optional[int]       =				self.get_dummy_components()
       UpperCAmelCase_							:       Tuple       =				self.pipeline_class(**lowerCamelCase_						)
       UpperCAmelCase_							:       Optional[Any]       =				pipe.to(lowerCamelCase_						)
       pipe.set_progress_bar_config(disable=lowerCamelCase_						)

       UpperCAmelCase_							:       Dict       =				1
       UpperCAmelCase_							:       Tuple       =				2

       UpperCAmelCase_							:       Dict       =				self.get_dummy_inputs(lowerCamelCase_						)

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

       UpperCAmelCase_							:       int       =				pipe(**lowerCamelCase_				,		num_images_per_prompt=lowerCamelCase_						)[0]

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







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







  '''simple docstring'''






  def _UpperCamelCase					(			self						):




       '''simple docstring'''

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






  def _UpperCamelCase					(			self						):




       '''simple docstring'''

       UpperCAmelCase_							:       int       =				load_image(
           'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main' '/shap_e/corgi.png'						)
       UpperCAmelCase_							:       Optional[Any]       =				load_numpy(
           'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'
           '/shap_e/test_shap_e_img2img_out.npy'						)
       UpperCAmelCase_							:       Optional[Any]       =				ShapEImgaImgPipeline.from_pretrained('openai/shap-e-img2img'						)
       UpperCAmelCase_							:       Optional[int]       =				pipe.to(lowerCamelCase_						)
       pipe.set_progress_bar_config(disable=lowerCamelCase_						)

       UpperCAmelCase_							:       Any       =				torch.Generator(device=lowerCamelCase_						).manual_seed(0						)

       UpperCAmelCase_							:       Dict       =				pipe(
           lowerCamelCase_				,		generator=lowerCamelCase_				,		guidance_scale=3.0				,		num_inference_steps=6_4				,		frame_size=6_4				,		output_type='np'				,		).images[0]

       assert images.shape == (2_0, 6_4, 6_4, 3)

       assert_mean_pixel_difference(lowerCamelCase_				,		lowerCamelCase_						)







'''simple docstring'''


from manim import *






class 					__SCREAMING_SNAKE_CASE						(	lowerCamelCase_	):







  '''simple docstring'''






  def _UpperCamelCase					(			self						):




       '''simple docstring'''

       UpperCAmelCase_							:       Dict       =				Rectangle(height=0.5				,		width=0.5						)
       UpperCAmelCase_							:       Any       =				Rectangle(height=0.46				,		width=0.46						).set_stroke(width=0						)

       UpperCAmelCase_							:       List[str]       =				Rectangle(height=0.25				,		width=0.25						)

       UpperCAmelCase_							:       Any       =				[mem.copy() for i in range(6						)]
       UpperCAmelCase_							:       Union[str, Any]       =				[mem.copy() for i in range(6						)]
       UpperCAmelCase_							:       Optional[int]       =				VGroup(*snake_case_						).arrange(snake_case_				,		buff=0						)
       UpperCAmelCase_							:       Any       =				VGroup(*snake_case_						).arrange(snake_case_				,		buff=0						)
       UpperCAmelCase_							:       str       =				VGroup(snake_case_				,		snake_case_						).arrange(snake_case_				,		buff=0						)
       UpperCAmelCase_							:       Any       =				Text('CPU'				,		font_size=2_4						)
       UpperCAmelCase_							:       Tuple       =				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_						)

       UpperCAmelCase_							:       str       =				[mem.copy() for i in range(4						)]
       UpperCAmelCase_							:       Dict       =				VGroup(*snake_case_						).arrange(snake_case_				,		buff=0						)
       UpperCAmelCase_							:       List[str]       =				Text('GPU'				,		font_size=2_4						)
       UpperCAmelCase_							:       Optional[int]       =				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_						)

       UpperCAmelCase_							:       str       =				[mem.copy() for i in range(6						)]
       UpperCAmelCase_							:       Tuple       =				VGroup(*snake_case_						).arrange(snake_case_				,		buff=0						)

       UpperCAmelCase_							:       str       =				Text('Model'				,		font_size=2_4						)
       UpperCAmelCase_							:       Optional[int]       =				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_						)

       UpperCAmelCase_							:       str       =				[]
       UpperCAmelCase_							:       Optional[Any]       =				[]

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

            UpperCAmelCase_							:       int       =				Rectangle(height=0.46				,		width=0.46						).set_stroke(width=0.0						).set_fill(snake_case_				,		opacity=0.8						)
            cpu_target.move_to(cpu_left_col_base[i]						)
            model_cpu_arr.append(snake_case_						)

       self.add(*snake_case_				,		*snake_case_						)

       UpperCAmelCase_							:       List[Any]       =				[meta_mem.copy() for i in range(6						)]
       UpperCAmelCase_							:       List[str]       =				[meta_mem.copy() for i in range(6						)]
       UpperCAmelCase_							:       Tuple       =				VGroup(*snake_case_						).arrange(snake_case_				,		buff=0						)
       UpperCAmelCase_							:       Dict       =				VGroup(*snake_case_						).arrange(snake_case_				,		buff=0						)
       UpperCAmelCase_							:       Optional[Any]       =				VGroup(snake_case_				,		snake_case_						).arrange(snake_case_				,		buff=0						)
       UpperCAmelCase_							:       Tuple       =				Text('Disk'				,		font_size=2_4						)
       UpperCAmelCase_							:       Union[str, Any]       =				Group(snake_case_				,		snake_case_						).arrange(snake_case_				,		buff=0.5				,		aligned_edge=snake_case_						)
       disk.move_to([-4, -1.25, 0]						)
       self.add(snake_case_				,		snake_case_						)

       UpperCAmelCase_							:       List[Any]       =				Square(side_length=2.2						)
       key.move_to([-5, 2, 0]						)

       UpperCAmelCase_							:       Any       =				MarkupText(
           F'''<b>Key:</b>\n\n<span fgcolor=\'{YELLOW}\'>●</span> Empty Model'''				,		font_size=1_8				,		)

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

       self.add(snake_case_				,		snake_case_						)

       UpperCAmelCase_							:       Dict       =				MarkupText(
           F'''<span fgcolor=\'{BLUE}\'>●</span> Checkpoint'''				,		font_size=1_8				,		)

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

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

       self.play(Write(snake_case_						)						)

       UpperCAmelCase_							:       Tuple       =				Square(0.3						)
       input.set_fill(snake_case_				,		opacity=1.0						)
       input.set_stroke(width=0.0						)
       input.next_to(model_base[0]				,		snake_case_				,		buff=0.5						)

       self.play(Write(snake_case_						)						)

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

       self.play(FadeOut(snake_case_						)						)

       UpperCAmelCase_							:       Any       =				Arrow(start=snake_case_				,		end=snake_case_				,		color=snake_case_				,		buff=0.5						)
       a.next_to(model_arr[0].get_left()				,		snake_case_				,		buff=0.2						)

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

       UpperCAmelCase_							:       List[str]       =				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=2_4				,		)
       step_a.move_to([2, 2, 0]						)

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

       UpperCAmelCase_							:       List[Any]       =				{'run_time': 1, 'fade_in': True, 'fade_out': True, 'buff': 0.02}

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

       UpperCAmelCase_							:       Union[str, Any]       =				a.copy()
       for i in range(6						):
            a_c.next_to(model_arr[i].get_right() + 0.02				,		snake_case_				,		buff=0.2						)

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

            UpperCAmelCase_							:       Tuple       =				AnimationGroup(
                FadeOut(snake_case_				,		run_time=0.5						)				,		MoveToTarget(snake_case_				,		run_time=0.5						)				,		FadeIn(snake_case_				,		run_time=0.5						)				,		lag_ratio=0.2						)

            self.play(snake_case_						)

            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]				,		**snake_case_						)				,		Circumscribe(cpu_left_col_base[i]				,		**snake_case_						)				,		Circumscribe(cpu_left_col_base[i + 1]				,		color=snake_case_				,		**snake_case_						)				,		Circumscribe(gpu_rect[0]				,		color=snake_case_				,		**snake_case_						)				,		Circumscribe(model_arr[i + 1]				,		color=snake_case_				,		**snake_case_						)				,		)
                 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=snake_case_				,		**snake_case_						)				,		Circumscribe(cpu_left_col_base[-1]				,		color=snake_case_				,		**snake_case_						)				,		Circumscribe(gpu_rect[0]				,		color=snake_case_				,		**snake_case_						)				,		)

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

            UpperCAmelCase_							:       Any       =				a_c
            UpperCAmelCase_							:       int       =				a_c.copy()

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

       UpperCAmelCase_							:       Optional[Any]       =				MarkupText(F'''Inference on a model too large for GPU memory\nis successfully completed.'''				,		font_size=2_4						)
       step_a.move_to([2, 2, 0]						)

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

       self.wait()






