
"""simple docstring"""




import gc
import unittest

import numpy as np
import torch

from diffusers import AutoencoderKL, DDIMScheduler, DiTPipeline, DPMSolverMultistepScheduler, TransformeraDModel
from diffusers.utils import is_xformers_available, load_numpy, slow, torch_device
from diffusers.utils.testing_utils import enable_full_determinism, require_torch_gpu

from ..pipeline_params import (
    CLASS_CONDITIONED_IMAGE_GENERATION_BATCH_PARAMS,
    CLASS_CONDITIONED_IMAGE_GENERATION_PARAMS,
)
from ..test_pipelines_common import PipelineTesterMixin


enable_full_determinism()



class 			lowerCAmelCase_		(a__				,      unittest.TestCase						):


    """simple docstring"""



    __UpperCamelCase   :		Optional[Any]       =     DiTPipeline
    __UpperCamelCase   :		Optional[int]       =     CLASS_CONDITIONED_IMAGE_GENERATION_PARAMS
    __UpperCamelCase   :		Tuple       =     PipelineTesterMixin.required_optional_params - {
        '''latents''',
        '''num_images_per_prompt''',
        '''callback''',
        '''callback_steps''',
    }
    __UpperCamelCase   :		Tuple       =     CLASS_CONDITIONED_IMAGE_GENERATION_BATCH_PARAMS
    __UpperCamelCase   :		Any       =     False



    def 			__magic_name__					(self					)      ->	List[str]:



           """simple docstring"""




           torch.manual_seed(0					)
           SCREAMING_SNAKE_CASE__		: List[str]     =     TransformeraDModel(
               sample_size=16     , num_layers=2     , patch_size=4     , attention_head_dim=8     , num_attention_heads=2     , in_channels=4     , out_channels=8     , attention_bias=_lowerCAmelCase     , activation_fn="""gelu-approximate"""     , num_embeds_ada_norm=10_00     , norm_type="""ada_norm_zero"""     , norm_elementwise_affine=_lowerCAmelCase     , )
           SCREAMING_SNAKE_CASE__		: Dict     =     AutoencoderKL()
           SCREAMING_SNAKE_CASE__		: Dict     =     DDIMScheduler()
           SCREAMING_SNAKE_CASE__		: str     =     {"""transformer""": transformer.eval(), """vae""": vae.eval(), """scheduler""": scheduler}
           return components



    def 			__magic_name__					(self     , SCREAMING_SNAKE_CASE__     , SCREAMING_SNAKE_CASE__=0					)      ->	List[str]:



           """simple docstring"""




           if str(_lowerCAmelCase					).startswith("""mps"""					):
                  SCREAMING_SNAKE_CASE__		: List[str]     =     torch.manual_seed(_lowerCAmelCase					)
           else:
                  SCREAMING_SNAKE_CASE__		: List[Any]     =     torch.Generator(device=_lowerCAmelCase					).manual_seed(_lowerCAmelCase					)
           SCREAMING_SNAKE_CASE__		: Tuple     =     {
               """class_labels""": [1],
               """generator""": generator,
               """num_inference_steps""": 2,
               """output_type""": """numpy""",
           }
           return inputs



    def 			__magic_name__					(self					)      ->	List[str]:



           """simple docstring"""




           SCREAMING_SNAKE_CASE__		: Optional[int]     =     """cpu"""

           SCREAMING_SNAKE_CASE__		: Optional[Any]     =     self.get_dummy_components()
           SCREAMING_SNAKE_CASE__		: Tuple     =     self.pipeline_class(**_lowerCAmelCase					)
           pipe.to(_lowerCAmelCase					)
           pipe.set_progress_bar_config(disable=_lowerCAmelCase					)

           SCREAMING_SNAKE_CASE__		: int     =     self.get_dummy_inputs(_lowerCAmelCase					)
           SCREAMING_SNAKE_CASE__		: Any     =     pipe(**_lowerCAmelCase					).images
           SCREAMING_SNAKE_CASE__		: Dict     =     image[0, -3:, -3:, -1]

           self.assertEqual(image.shape     , (1, 16, 16, 3)					)
           SCREAMING_SNAKE_CASE__		: str     =     np.array([0.2946, 0.6601, 0.4329, 0.3296, 0.4144, 0.5319, 0.7273, 0.5013, 0.4457]					)
           SCREAMING_SNAKE_CASE__		: Tuple     =     np.abs(image_slice.flatten() - expected_slice					).max()
           self.assertLessEqual(_lowerCAmelCase     , 1E-3					)



    def 			__magic_name__					(self					)      ->	List[str]:



           """simple docstring"""




           self._test_inference_batch_single_identical(relax_max_difference=_lowerCAmelCase     , expected_max_diff=1E-3					)



    @unittest.skipIf(
        torch_device != """cuda""" or not is_xformers_available()     , reason="""XFormers attention is only available with CUDA and `xformers` installed"""     , )
    def 			__magic_name__					(self					)      ->	Tuple:



           """simple docstring"""




           self._test_xformers_attention_forwardGenerator_pass(expected_max_diff=1E-3					)






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


    """simple docstring"""



    def 			__magic_name__					(self					)      ->	int:



           """simple docstring"""




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



    def 			__magic_name__					(self					)      ->	str:



           """simple docstring"""




           SCREAMING_SNAKE_CASE__		: List[Any]     =     torch.manual_seed(0					)

           SCREAMING_SNAKE_CASE__		: Any     =     DiTPipeline.from_pretrained("""facebook/DiT-XL-2-256"""					)
           pipe.to("""cuda"""					)

           SCREAMING_SNAKE_CASE__		: int     =     ["""vase""", """umbrella""", """white shark""", """white wolf"""]
           SCREAMING_SNAKE_CASE__		: int     =     pipe.get_label_ids(_lowerCAmelCase					)

           SCREAMING_SNAKE_CASE__		: str     =     pipe(_lowerCAmelCase     , generator=_lowerCAmelCase     , num_inference_steps=40     , output_type="""np"""					).images

           for word, image in zip(_lowerCAmelCase     , _lowerCAmelCase					):
                  SCREAMING_SNAKE_CASE__		: str     =     load_numpy(
                      F'''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/dit/{word}.npy'''					)
                  assert np.abs((expected_image - image).max()					) < 1E-2



    def 			__magic_name__					(self					)      ->	str:



           """simple docstring"""




           SCREAMING_SNAKE_CASE__		: Optional[int]     =     DiTPipeline.from_pretrained("""facebook/DiT-XL-2-512"""					)
           SCREAMING_SNAKE_CASE__		: str     =     DPMSolverMultistepScheduler.from_config(pipe.scheduler.config					)
           pipe.to("""cuda"""					)

           SCREAMING_SNAKE_CASE__		: Any     =     ["""vase""", """umbrella"""]
           SCREAMING_SNAKE_CASE__		: List[Any]     =     pipe.get_label_ids(_lowerCAmelCase					)

           SCREAMING_SNAKE_CASE__		: Optional[Any]     =     torch.manual_seed(0					)
           SCREAMING_SNAKE_CASE__		: str     =     pipe(_lowerCAmelCase     , generator=_lowerCAmelCase     , num_inference_steps=25     , output_type="""np"""					).images

           for word, image in zip(_lowerCAmelCase     , _lowerCAmelCase					):
                  SCREAMING_SNAKE_CASE__		: List[str]     =     load_numpy(
                      """https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main"""
                      F'''/dit/{word}_512.npy'''					)

                  assert np.abs((expected_image - image).max()					) < 1E-1






"""simple docstring"""




from manim import *



class 			lowerCAmelCase_		(a__						):


				"""simple docstring"""



				def 			__magic_name__					(self					)      ->	List[str]:



											"""simple docstring"""




											SCREAMING_SNAKE_CASE__		: Dict     =     Rectangle(height=0.5     , width=0.5					)
											SCREAMING_SNAKE_CASE__		: Union[str, Any]     =     Rectangle(height=0.46     , width=0.46					).set_stroke(width=0					)

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

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

											self.add(SCREAMING_SNAKE_CASE__					)

											SCREAMING_SNAKE_CASE__		: List[Any]     =     [mem.copy() for i in range(6					)]
											SCREAMING_SNAKE_CASE__		: List[Any]     =     VGroup(*SCREAMING_SNAKE_CASE__					).arrange(SCREAMING_SNAKE_CASE__     , buff=0					)

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

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

											SCREAMING_SNAKE_CASE__		: Tuple     =     MarkupText(
											    F'''First, an empty model skeleton is loaded\ninto <span fgcolor=\'{YELLOW}\'>memory</span> without using much RAM.'''     , font_size=24     , )

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

											SCREAMING_SNAKE_CASE__		: Tuple     =     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(SCREAMING_SNAKE_CASE__     , run_time=2.5					)     , Write(SCREAMING_SNAKE_CASE__					)     , Write(SCREAMING_SNAKE_CASE__					)					)

											self.add(SCREAMING_SNAKE_CASE__					)

											SCREAMING_SNAKE_CASE__		: Optional[int]     =     []
											SCREAMING_SNAKE_CASE__		: Union[str, Any]     =     []
											SCREAMING_SNAKE_CASE__		: str     =     []
											for i, rect in enumerate(SCREAMING_SNAKE_CASE__					):
																		SCREAMING_SNAKE_CASE__		: Union[str, Any]     =     Rectangle(height=0.46     , width=0.46					).set_stroke(width=0.0					).set_fill(SCREAMING_SNAKE_CASE__     , opacity=0.7					)
																		cpu_target.move_to(SCREAMING_SNAKE_CASE__					)
																		cpu_target.generate_target()
																		SCREAMING_SNAKE_CASE__		: Tuple     =     0.46 / 4
																		SCREAMING_SNAKE_CASE__		: int     =     0.46 / 3

																		if i == 0:
																									cpu_target.target.next_to(cpu_left_col_base[0].get_corner(DOWN + LEFT					)     , buff=0.02     , direction=SCREAMING_SNAKE_CASE__					)
																									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=SCREAMING_SNAKE_CASE__     , buff=0.0					)
																		else:
																									cpu_target.target.next_to(cpu_targs[i - 1].target     , direction=SCREAMING_SNAKE_CASE__     , buff=0.0					)
																		cpu_targs.append(SCREAMING_SNAKE_CASE__					)

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

											self.play(*SCREAMING_SNAKE_CASE__					)
											self.play(*SCREAMING_SNAKE_CASE__					)

											self.wait()




