


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 					__magic_name__		(			__UpperCAmelCase			,						unittest.TestCase     ):
       __A			:      Union[str, Any]        =   DiTPipeline
       __A			:      Dict        =   CLASS_CONDITIONED_IMAGE_GENERATION_PARAMS
       __A			:      Any        =   PipelineTesterMixin.required_optional_params - {
           "latents",
           "num_images_per_prompt",
           "callback",
           "callback_steps",
       }
       __A			:      Dict        =   CLASS_CONDITIONED_IMAGE_GENERATION_BATCH_PARAMS
       __A			:      List[str]        =   False






       def  __snake_case						(   self  :    Dict				):

              '''simple docstring'''






              torch.manual_seed(0				)
              lowercase		:int       							=				TransformeraDModel(
                  sample_size=1_6  ,     num_layers=2  ,     patch_size=4  ,     attention_head_dim=8  ,     num_attention_heads=2  ,     in_channels=4  ,     out_channels=8  ,     attention_bias=snake_case__  ,     activation_fn='''gelu-approximate'''  ,     num_embeds_ada_norm=1_0_0_0  ,     norm_type='''ada_norm_zero'''  ,     norm_elementwise_affine=snake_case__  ,     )
              lowercase		:Optional[int]       							=				AutoencoderKL()
              lowercase		:List[Any]       							=				DDIMScheduler()
              lowercase		:Tuple       							=				{'''transformer''': transformer.eval(), '''vae''': vae.eval(), '''scheduler''': scheduler}
              return components






       def  __snake_case						(   self  :    int  ,     snake_case__  :    Optional[int]  ,     snake_case__  :    Optional[int]=0				):

              '''simple docstring'''






              if str(snake_case__				).startswith('''mps'''				):
                     lowercase		:List[Any]       							=				torch.manual_seed(snake_case__				)
              else:
                     lowercase		:Dict       							=				torch.Generator(device=snake_case__				).manual_seed(snake_case__				)
              lowercase		:Optional[Any]       							=				{
                  '''class_labels''': [1],
                  '''generator''': generator,
                  '''num_inference_steps''': 2,
                  '''output_type''': '''numpy''',
              }
              return inputs






       def  __snake_case						(   self  :    Tuple				):

              '''simple docstring'''






              lowercase		:List[str]       							=				'''cpu'''

              lowercase		:Any       							=				self.get_dummy_components()
              lowercase		:Union[str, Any]       							=				self.pipeline_class(**snake_case__				)
              pipe.to(snake_case__				)
              pipe.set_progress_bar_config(disable=snake_case__				)

              lowercase		:int       							=				self.get_dummy_inputs(snake_case__				)
              lowercase		:str       							=				pipe(**snake_case__				).images
              lowercase		:Optional[Any]       							=				image[0, -3:, -3:, -1]

              self.assertEqual(image.shape  ,     (1, 1_6, 1_6, 3)				)
              lowercase		:int       							=				np.array([0.29_46, 0.66_01, 0.43_29, 0.32_96, 0.41_44, 0.53_19, 0.72_73, 0.50_13, 0.44_57]				)
              lowercase		:List[str]       							=				np.abs(image_slice.flatten() - expected_slice				).max()
              self.assertLessEqual(snake_case__  ,     1e-3				)






       def  __snake_case						(   self  :    int				):

              '''simple docstring'''






              self._test_inference_batch_single_identical(relax_max_difference=snake_case__  ,     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  __snake_case						(   self  :    Optional[Any]				):

              '''simple docstring'''






              self._test_xformers_attention_forwardGenerator_pass(expected_max_diff=1e-3				)







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






       def  __snake_case						(   self  :    List[str]				):

              '''simple docstring'''






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






       def  __snake_case						(   self  :    Tuple				):

              '''simple docstring'''






              lowercase		:str       							=				torch.manual_seed(0				)

              lowercase		:List[str]       							=				DiTPipeline.from_pretrained('''facebook/DiT-XL-2-256'''				)
              pipe.to('''cuda'''				)

              lowercase		:int       							=				['''vase''', '''umbrella''', '''white shark''', '''white wolf''']
              lowercase		:List[str]       							=				pipe.get_label_ids(snake_case__				)

              lowercase		:str       							=				pipe(snake_case__  ,     generator=snake_case__  ,     num_inference_steps=4_0  ,     output_type='''np'''				).images

              for word, image in zip(snake_case__  ,     snake_case__				):
                     lowercase		: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  __snake_case						(   self  :    Optional[Any]				):

              '''simple docstring'''






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

              lowercase		:List[Any]       							=				['''vase''', '''umbrella''']
              lowercase		:Dict       							=				pipe.get_label_ids(snake_case__				)

              lowercase		:Dict       							=				torch.manual_seed(0				)
              lowercase		:List[Any]       							=				pipe(snake_case__  ,     generator=snake_case__  ,     num_inference_steps=2_5  ,     output_type='''np'''				).images

              for word, image in zip(snake_case__  ,     snake_case__				):
                     lowercase		:Any       							=				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 typing import List, Optional, Tuple, Union

import torch

from ...models import UNetaDModel
from ...schedulers import ScoreSdeVeScheduler
from ...utils import randn_tensor
from ..pipeline_utils import DiffusionPipeline, ImagePipelineOutput
class 					__magic_name__		(			__UpperCAmelCase     ):
       __A			:      UNetaDModel
       __A			:      ScoreSdeVeScheduler






       def __init__(   self  :    Optional[Any]  ,     snake_case__  :    UNetaDModel  ,     snake_case__  :    ScoreSdeVeScheduler				):

              '''simple docstring'''






              super().__init__()
              self.register_modules(unet=snake_case__  ,     scheduler=snake_case__				)







       @torch.no_grad()
       def __call__(   self  :    Tuple  ,     snake_case__  :    int = 1  ,     snake_case__  :    int = 2_0_0_0  ,     snake_case__  :    Optional[Union[torch.Generator, List[torch.Generator]]] = None  ,     snake_case__  :    Optional[str] = "pil"  ,     snake_case__  :    bool = True  ,     **snake_case__  :    Optional[Any]  ,     ):

              '''simple docstring'''






              lowercase		:str       							=				self.unet.config.sample_size
              lowercase		:str       							=				(batch_size, 3, img_size, img_size)

              lowercase		:Any       							=				self.unet

              lowercase		:Dict       							=				randn_tensor(snake_case__  ,     generator=snake_case__				) * self.scheduler.init_noise_sigma
              lowercase		:int       							=				sample.to(self.device				)

              self.scheduler.set_timesteps(snake_case__				)
              self.scheduler.set_sigmas(snake_case__				)

              for i, t in enumerate(self.progress_bar(self.scheduler.timesteps				)				):
                     lowercase		:Optional[int]       							=				self.scheduler.sigmas[i] * torch.ones(shape[0]  ,     device=self.device				)

                     # correction step
                     for _ in range(self.scheduler.config.correct_steps				):
                            lowercase		:str       							=				self.unet(snake_case__  ,     snake_case__				).sample
                            lowercase		:Optional[Any]       							=				self.scheduler.step_correct(snake_case__  ,     snake_case__  ,     generator=snake_case__				).prev_sample

                     # prediction step
                     lowercase		:Optional[Any]       							=				model(snake_case__  ,     snake_case__				).sample
                     lowercase		:Dict       							=				self.scheduler.step_pred(snake_case__  ,     snake_case__  ,     snake_case__  ,     generator=snake_case__				)

                     lowercase							,  lowercase		:Tuple       							=				output.prev_sample, output.prev_sample_mean

              lowercase		:str       							=				sample_mean.clamp(0  ,     1				)
              lowercase		:int       							=				sample.cpu().permute(0  ,     2  ,     3  ,     1				).numpy()
              if output_type == "pil":
                     lowercase		:str       							=				self.numpy_to_pil(snake_case__				)

              if not return_dict:
                     return (sample,)

              return ImagePipelineOutput(images=snake_case__				)

