

"""simple docstring"""






from manim import *







class    snake_case__	(			__a  ):







			def       __lowerCAmelCase					(      self      :			Tuple	):





								'''simple docstring'''



								UpperCAmelCase      : Union[str, Any]   =					Rectangle(height=0.5      ,	width=0.5	)
								UpperCAmelCase      : List[Any]   =					Rectangle(height=0.4_6      ,	width=0.4_6	).set_stroke(width=0	)

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

								UpperCAmelCase      : Optional[int]   =					[mem.copy() for i in range(1	)]
								UpperCAmelCase      : Dict   =					VGroup(*_lowercase	).arrange(_lowercase      ,	buff=0	)
								UpperCAmelCase      : Any   =					Text("GPU"      ,	font_size=24	)
								UpperCAmelCase      : List[str]   =					Group(_lowercase      ,	_lowercase	).arrange(_lowercase      ,	buff=0.5      ,	aligned_edge=_lowercase	)
								gpu.align_to(_lowercase      ,	_lowercase	)
								gpu.set_x(gpu.get_x() - 1	)

								self.add(_lowercase	)

								UpperCAmelCase      : Union[str, Any]   =					[mem.copy() for i in range(6	)]
								UpperCAmelCase      : str   =					VGroup(*_lowercase	).arrange(_lowercase      ,	buff=0	)

								UpperCAmelCase      : Union[str, Any]   =					Text("Model"      ,	font_size=24	)
								UpperCAmelCase      : str   =					Group(_lowercase      ,	_lowercase	).arrange(_lowercase      ,	buff=0.5      ,	aligned_edge=_lowercase	)
								model.move_to([3, -1.0, 0]	)

								self.play(
								    Create(_lowercase      ,	run_time=1	)      ,	Create(_lowercase      ,	run_time=1	)      ,	Create(_lowercase      ,	run_time=1	)      ,	)

								UpperCAmelCase      : Optional[Any]   =					MarkupText(
								    f"""First, an empty model skeleton is loaded\ninto <span fgcolor='{YELLOW}'>memory</span> without using much RAM."""      ,	font_size=24      ,	)

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

								UpperCAmelCase      : Union[str, 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]	)

								step_a.move_to([2, 2, 0]	)
								self.play(Write(_lowercase      ,	run_time=2.5	)      ,	Write(_lowercase	)      ,	Write(_lowercase	)	)

								self.add(_lowercase	)

								UpperCAmelCase      : Tuple   =					[]
								UpperCAmelCase      : Tuple   =					[]
								UpperCAmelCase      : str   =					[]
								for i, rect in enumerate(_lowercase	):
													UpperCAmelCase      : Tuple   =					Rectangle(height=0.4_6      ,	width=0.4_6	).set_stroke(width=0.0	).set_fill(_lowercase      ,	opacity=0.7	)
													cpu_target.move_to(_lowercase	)
													cpu_target.generate_target()
													UpperCAmelCase      : str   =					0.4_6 / 4
													UpperCAmelCase      : List[Any]   =					0.4_6 / 3

													if i == 0:
																		cpu_target.target.next_to(cpu_left_col_base[0].get_corner(DOWN + LEFT	)      ,	buff=0.0_2      ,	direction=_lowercase	)
																		cpu_target.target.set_x(cpu_target.target.get_x() + 0.1	)
													elif i == 3:
																		cpu_target.target.next_to(cpu_targs[0].target      ,	direction=_lowercase      ,	buff=0.0	)
													else:
																		cpu_target.target.next_to(cpu_targs[i - 1].target      ,	direction=_lowercase      ,	buff=0.0	)
													cpu_targs.append(_lowercase	)

													first_animations.append(rect.animate(run_time=0.5	).set_stroke(_lowercase	)	)
													second_animations.append(MoveToTarget(_lowercase      ,	run_time=1.5	)	)

								self.play(*_lowercase	)
								self.play(*_lowercase	)

								self.wait()



import unittest

import torch

from diffusers import VQModel
from diffusers.utils import floats_tensor, torch_device
from diffusers.utils.testing_utils import enable_full_determinism

from .test_modeling_common import ModelTesterMixin, UNetTesterMixin


enable_full_determinism()




class 			__magic_name__	(		__a   ,						__a   ,						unittest.TestCase     ):





   """simple docstring"""




   lowerCAmelCase   :     Dict	      = VQModel
   lowerCAmelCase   :     List[str]	      = '''sample'''



   @property
   def 	lowerCAmelCase    (						self       :	List[Any]		,	_lowercase       :	Tuple=(32, 32)				):



    """simple docstring"""



    _UpperCamelCase:  List[Any]    						=			4
    _UpperCamelCase:  Tuple    						=			3

    _UpperCamelCase:  List[str]    						=			floats_tensor((batch_size, num_channels) + sizes				).to(_lowercase				)

    return {"sample": image}



   @property
   def 	lowerCAmelCase    (						self       :	List[Any]				):



    """simple docstring"""



    return (3, 32, 32)



   @property
   def 	lowerCAmelCase    (						self       :	str				):



    """simple docstring"""



    return (3, 32, 32)



   def 	lowerCAmelCase    (						self       :	Optional[int]				):



    """simple docstring"""



    _UpperCamelCase:  Any    						=			{
        '''block_out_channels''': [32, 64],
        '''in_channels''': 3,
        '''out_channels''': 3,
        '''down_block_types''': ['''DownEncoderBlock2D''', '''DownEncoderBlock2D'''],
        '''up_block_types''': ['''UpDecoderBlock2D''', '''UpDecoderBlock2D'''],
        '''latent_channels''': 3,
    }
    _UpperCamelCase:  Optional[Any]    						=			self.dummy_input
    return init_dict, inputs_dict



   def 	lowerCAmelCase    (						self       :	List[Any]				):



    """simple docstring"""



    pass



   def 	lowerCAmelCase    (						self       :	List[Any]				):



    """simple docstring"""



    pass



   def 	lowerCAmelCase    (						self       :	Optional[Any]				):



    """simple docstring"""



    _UpperCamelCase						,				_UpperCamelCase:  int    						=			VQModel.from_pretrained('''fusing/vqgan-dummy'''		,	output_loading_info=_lowercase				)
    self.assertIsNotNone(_lowercase				)
    self.assertEqual(len(loading_info['''missing_keys''']				)		,	0				)

    model.to(_lowercase				)
    _UpperCamelCase:  Tuple    						=			model(**self.dummy_input				)

    assert image is not None, "Make sure output is not None"



   def 	lowerCAmelCase    (						self       :	Optional[Any]				):



    """simple docstring"""



    _UpperCamelCase:  Dict    						=			VQModel.from_pretrained('''fusing/vqgan-dummy'''				)
    model.to(_lowercase				).eval()

    torch.manual_seed(0				)
    if torch.cuda.is_available():
     torch.cuda.manual_seed_all(0				)

    _UpperCamelCase:  int    						=			torch.randn(1		,	model.config.in_channels		,	model.config.sample_size		,	model.config.sample_size				)
    _UpperCamelCase:  Tuple    						=			image.to(_lowercase				)
    with torch.no_grad():
     _UpperCamelCase:  Optional[int]    						=			model(_lowercase				).sample

    _UpperCamelCase:  List[Any]    						=			output[0, -1, -3:, -3:].flatten().cpu()
    # fmt: off
    _UpperCamelCase:  Any    						=			torch.tensor([-0.0153, -0.4044, -0.1880, -0.5161, -0.2418, -0.4072, -0.1612, -0.0633, -0.0143]				)
    # fmt: on
    self.assertTrue(torch.allclose(_lowercase		,	_lowercase		,	atol=1E-3				)				)