"""simple docstring"""




import gc
import random
import unittest

import numpy as np
import torch
from transformers import XLMRobertaTokenizer

from diffusers import (
    AltDiffusionImgaImgPipeline,
    AutoencoderKL,
    PNDMScheduler,
    UNetaDConditionModel,
)
from diffusers.image_processor import VaeImageProcessor
from diffusers.pipelines.alt_diffusion.modeling_roberta_series import (
    RobertaSeriesConfig,
    RobertaSeriesModelWithTransformation,
)
from diffusers.utils import floats_tensor, load_image, load_numpy, slow, torch_device
from diffusers.utils.testing_utils import enable_full_determinism, require_torch_gpu


enable_full_determinism()
class     UpperCamelCase__(			unittest.TestCase ):





						def 				snake_case__	(       self     )					->				List[Any]:
													# clean up the VRAM after each test
													super().tearDown()
													gc.collect()
													torch.cuda.empty_cache()





						@property
						def 				snake_case__	(       self     )					->				Tuple:
													A__	     =							1
													A__	     =							3
													A__	     =							(32, 32)

													A__	     =							floats_tensor((batch_size, num_channels) + sizes		,rng=random.Random(0     )     ).to(__UpperCAmelCase     )
													return image





						@property
						def 				snake_case__	(       self     )					->				List[Any]:
													torch.manual_seed(0     )
													A__	     =							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		,)
													return model





						@property
						def 				snake_case__	(       self     )					->				Optional[int]:
													torch.manual_seed(0     )
													A__	     =							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		,)
													return model





						@property
						def 				snake_case__	(       self     )					->				List[str]:
													torch.manual_seed(0     )
													A__	     =							RobertaSeriesConfig(
													    hidden_size=32		,project_dim=32		,intermediate_size=37		,layer_norm_eps=1e-05		,num_attention_heads=4		,num_hidden_layers=5		,pad_token_id=1		,vocab_size=50_06		,)
													return RobertaSeriesModelWithTransformation(__UpperCAmelCase     )





						@property
						def 				snake_case__	(       self     )					->				Optional[int]:
													def extract(*__UpperCAmelCase		,**__UpperCAmelCase     ):
																				class     UpperCamelCase__:





																										def __init__(       self     )					->				Optional[Any]:
																																	A__	     =							torch.ones([0]     )





																										def 				snake_case__	(       self		,__UpperCAmelCase     )					->				Optional[int]:
																																	self.pixel_values.to(__UpperCAmelCase     )
																																	return self






																				return Out()

													return extract





						def 				snake_case__	(       self     )					->				Tuple:
													A__	     =							'cpu'  # ensure determinism for the device-dependent torch.Generator
													A__	     =							self.dummy_cond_unet
													A__	     =							PNDMScheduler(skip_prk_steps=__UpperCAmelCase     )
													A__	     =							self.dummy_vae
													A__	     =							self.dummy_text_encoder
													A__	     =							XLMRobertaTokenizer.from_pretrained('hf-internal-testing/tiny-xlm-roberta'     )
													A__	     =							77

													A__	     =							self.dummy_image.to(__UpperCAmelCase     )
													A__	     =							init_image / 2 + 0.5

													# make sure here that pndm scheduler skips prk
													A__	     =							AltDiffusionImgaImgPipeline(
													    unet=__UpperCAmelCase		,scheduler=__UpperCAmelCase		,vae=__UpperCAmelCase		,text_encoder=__UpperCAmelCase		,tokenizer=__UpperCAmelCase		,safety_checker=__UpperCAmelCase		,feature_extractor=self.dummy_extractor		,)
													A__	     =							VaeImageProcessor(vae_scale_factor=alt_pipe.vae_scale_factor		,do_normalize=__UpperCAmelCase     )
													A__	     =							alt_pipe.to(__UpperCAmelCase     )
													alt_pipe.set_progress_bar_config(disable=__UpperCAmelCase     )

													A__	     =							'A painting of a squirrel eating a burger'
													A__	     =							torch.Generator(device=__UpperCAmelCase     ).manual_seed(0     )
													A__	     =							alt_pipe(
													    [prompt]		,generator=__UpperCAmelCase		,guidance_scale=6.0		,num_inference_steps=2		,output_type='np'		,image=__UpperCAmelCase		,)

													A__	     =							output.images

													A__	     =							torch.Generator(device=__UpperCAmelCase     ).manual_seed(0     )
													A__	     =							alt_pipe(
													    [prompt]		,generator=__UpperCAmelCase		,guidance_scale=6.0		,num_inference_steps=2		,output_type='np'		,image=__UpperCAmelCase		,return_dict=__UpperCAmelCase		,)[0]

													A__	     =							image[0, -3:, -3:, -1]
													A__	     =							image_from_tuple[0, -3:, -3:, -1]

													assert image.shape == (1, 32, 32, 3)
													A__	     =							np.array([0.4_4_2_7, 0.3_7_3_1, 0.4_2_4_9, 0.4_9_4_1, 0.4_5_4_6, 0.4_1_4_8, 0.4_1_9_3, 0.4_6_6_6, 0.4_4_9_9]     )

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





						@unittest.skipIf(torch_device != 'cuda'		,'This test requires a GPU'     )
						def 				snake_case__	(       self     )					->				Tuple:
													A__	     =							self.dummy_cond_unet
													A__	     =							PNDMScheduler(skip_prk_steps=__UpperCAmelCase     )
													A__	     =							self.dummy_vae
													A__	     =							self.dummy_text_encoder
													A__	     =							XLMRobertaTokenizer.from_pretrained('hf-internal-testing/tiny-xlm-roberta'     )
													A__	     =							77

													A__	     =							self.dummy_image.to(__UpperCAmelCase     )

													# put models in fp16
													A__	     =							unet.half()
													A__	     =							vae.half()
													A__	     =							bert.half()

													# make sure here that pndm scheduler skips prk
													A__	     =							AltDiffusionImgaImgPipeline(
													    unet=__UpperCAmelCase		,scheduler=__UpperCAmelCase		,vae=__UpperCAmelCase		,text_encoder=__UpperCAmelCase		,tokenizer=__UpperCAmelCase		,safety_checker=__UpperCAmelCase		,feature_extractor=self.dummy_extractor		,)
													A__	     =							VaeImageProcessor(vae_scale_factor=alt_pipe.vae_scale_factor		,do_normalize=__UpperCAmelCase     )
													A__	     =							alt_pipe.to(__UpperCAmelCase     )
													alt_pipe.set_progress_bar_config(disable=__UpperCAmelCase     )

													A__	     =							'A painting of a squirrel eating a burger'
													A__	     =							torch.manual_seed(0     )
													A__	     =							alt_pipe(
													    [prompt]		,generator=__UpperCAmelCase		,num_inference_steps=2		,output_type='np'		,image=__UpperCAmelCase		,).images

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





						@unittest.skipIf(torch_device != 'cuda'		,'This test requires a GPU'     )
						def 				snake_case__	(       self     )					->				Optional[Any]:
													A__	     =							load_image(
													    'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'
													    '/img2img/sketch-mountains-input.jpg'     )
													# resize to resolution that is divisible by 8 but not 16 or 32
													A__	     =							init_image.resize((7_60, 5_04)     )

													A__	     =							'BAAI/AltDiffusion'
													A__	     =							AltDiffusionImgaImgPipeline.from_pretrained(
													    __UpperCAmelCase		,safety_checker=__UpperCAmelCase		,)
													pipe.to(__UpperCAmelCase     )
													pipe.set_progress_bar_config(disable=__UpperCAmelCase     )
													pipe.enable_attention_slicing()

													A__	     =							'A fantasy landscape, trending on artstation'

													A__	     =							torch.manual_seed(0     )
													A__	     =							pipe(
													    prompt=__UpperCAmelCase		,image=__UpperCAmelCase		,strength=0.7_5		,guidance_scale=7.5		,generator=__UpperCAmelCase		,output_type='np'		,)
													A__	     =							output.images[0]

													A__	     =							image[2_55:2_58, 3_83:3_86, -1]

													assert image.shape == (5_04, 7_60, 3)
													A__	     =							np.array([0.9_3_5_8, 0.9_3_9_7, 0.9_5_9_9, 0.9_9_0_1, 1.0_0_0_0, 1.0_0_0_0, 0.9_8_8_2, 1.0_0_0_0, 1.0_0_0_0]     )

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






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





						def 				snake_case__	(       self     )					->				Tuple:
													# clean up the VRAM after each test
													super().tearDown()
													gc.collect()
													torch.cuda.empty_cache()





						def 				snake_case__	(       self     )					->				Dict:
													A__	     =							load_image(
													    'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'
													    '/img2img/sketch-mountains-input.jpg'     )
													A__	     =							init_image.resize((7_68, 5_12)     )
													A__	     =							load_numpy(
													    'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/img2img/fantasy_landscape_alt.npy'     )

													A__	     =							'BAAI/AltDiffusion'
													A__	     =							AltDiffusionImgaImgPipeline.from_pretrained(
													    __UpperCAmelCase		,safety_checker=__UpperCAmelCase		,)
													pipe.to(__UpperCAmelCase     )
													pipe.set_progress_bar_config(disable=__UpperCAmelCase     )
													pipe.enable_attention_slicing()

													A__	     =							'A fantasy landscape, trending on artstation'

													A__	     =							torch.manual_seed(0     )
													A__	     =							pipe(
													    prompt=__UpperCAmelCase		,image=__UpperCAmelCase		,strength=0.7_5		,guidance_scale=7.5		,generator=__UpperCAmelCase		,output_type='np'		,)
													A__	     =							output.images[0]

													assert image.shape == (5_12, 7_68, 3)
													# img2img is flaky across GPUs even in fp32, so using MAE here
													assert np.abs(expected_image - image     ).max() < 1e-2






"""simple docstring"""




from manim import *
class     UpperCamelCase__(			__A ):





						def 				snake_case__	(       self     )					->				List[str]:
													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(*__UpperCAmelCase     ).arrange(__UpperCAmelCase		,buff=0     )
													A__	     =							VGroup(*__UpperCAmelCase     ).arrange(__UpperCAmelCase		,buff=0     )
													A__	     =							VGroup(__UpperCAmelCase		,__UpperCAmelCase     ).arrange(__UpperCAmelCase		,buff=0     )
													A__	     =							Text('CPU'		,font_size=24     )
													A__	     =							Group(__UpperCAmelCase		,__UpperCAmelCase     ).arrange(__UpperCAmelCase		,buff=0.5		,aligned_edge=__UpperCAmelCase     )
													cpu.move_to([-2.5, -0.5, 0]     )
													self.add(__UpperCAmelCase     )

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

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

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

													A__	     =							[]
													A__	     =							[]

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

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

													self.add(*__UpperCAmelCase		,*__UpperCAmelCase     )

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

													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(__UpperCAmelCase		,__UpperCAmelCase     )

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

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

													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(__UpperCAmelCase     )     )

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

													self.play(Write(__UpperCAmelCase     )     )

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

													self.play(FadeOut(__UpperCAmelCase     )     )

													A__	     =							Arrow(start=__UpperCAmelCase		,end=__UpperCAmelCase		,color=__UpperCAmelCase		,buff=0.5     )
													a.next_to(model_arr[0].get_left()		,__UpperCAmelCase		,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(__UpperCAmelCase		,run_time=3     )     )

													A__	     =							{'run_time': 1, 'fade_in': True, 'fade_out': True, 'buff': 0.0_2}

													self.play(
													    Write(__UpperCAmelCase     )		,Circumscribe(model_arr[0]		,color=__UpperCAmelCase		,**__UpperCAmelCase     )		,Circumscribe(model_cpu_arr[0]		,color=__UpperCAmelCase		,**__UpperCAmelCase     )		,Circumscribe(gpu_rect[0]		,color=__UpperCAmelCase		,**__UpperCAmelCase     )		,)
													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		,__UpperCAmelCase		,buff=0.2     )

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

																				A__	     =							AnimationGroup(
																				    FadeOut(__UpperCAmelCase		,run_time=0.5     )		,MoveToTarget(__UpperCAmelCase		,run_time=0.5     )		,FadeIn(__UpperCAmelCase		,run_time=0.5     )		,lag_ratio=0.2     )

																				self.play(__UpperCAmelCase     )

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

																											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(__UpperCAmelCase     )		,FadeOut(__UpperCAmelCase		,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(__UpperCAmelCase		,run_time=3     )		,MoveToTarget(__UpperCAmelCase     )     )

													self.wait()





