


import gc
import unittest

import numpy as np
import torch

from diffusers import StableDiffusionKDiffusionPipeline
from diffusers.utils import slow, torch_device
from diffusers.utils.testing_utils import enable_full_determinism, require_torch_gpu


enable_full_determinism()






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




       """simple docstring"""



       def        _snake_case       (      self		)->Dict:




         '''simple docstring'''





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



       def        _snake_case       (      self		)->List[Any]:




         '''simple docstring'''





         A_      :				Optional[Any]		     =					StableDiffusionKDiffusionPipeline.from_pretrained('''CompVis/stable-diffusion-v1-4'''		)
         A_      :				Any		     =					sd_pipe.to(_SCREAMING_SNAKE_CASE		)
         sd_pipe.set_progress_bar_config(disable=_SCREAMING_SNAKE_CASE		)

         sd_pipe.set_scheduler('''sample_euler'''		)

         A_      :				Union[str, Any]		     =					'''A painting of a squirrel eating a burger'''
         A_      :				Dict		     =					torch.manual_seed(0		)
         A_      :				str		     =					sd_pipe([prompt]				,       generator=_SCREAMING_SNAKE_CASE				,       guidance_scale=9.0				,       num_inference_steps=20				,       output_type='''np'''		)

         A_      :				int		     =					output.images

         A_      :				Optional[Any]		     =					image[0, -3:, -3:, -1]

         assert image.shape == (1, 512, 512, 3)
         A_      :				List[str]		     =					np.array([0.0_4_4_7, 0.0_4_9_2, 0.0_4_6_8, 0.0_4_0_8, 0.0_3_8_3, 0.0_4_0_8, 0.0_3_5_4, 0.0_3_8_0, 0.0_3_3_9]		)

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



       def        _snake_case       (      self		)->Union[str, Any]:




         '''simple docstring'''





         A_      :				Optional[int]		     =					StableDiffusionKDiffusionPipeline.from_pretrained('''stabilityai/stable-diffusion-2-1-base'''		)
         A_      :				Any		     =					sd_pipe.to(_SCREAMING_SNAKE_CASE		)
         sd_pipe.set_progress_bar_config(disable=_SCREAMING_SNAKE_CASE		)

         sd_pipe.set_scheduler('''sample_euler'''		)

         A_      :				Tuple		     =					'''A painting of a squirrel eating a burger'''
         A_      :				Tuple		     =					torch.manual_seed(0		)
         A_      :				Union[str, Any]		     =					sd_pipe([prompt]				,       generator=_SCREAMING_SNAKE_CASE				,       guidance_scale=9.0				,       num_inference_steps=20				,       output_type='''np'''		)

         A_      :				List[str]		     =					output.images

         A_      :				Union[str, Any]		     =					image[0, -3:, -3:, -1]

         assert image.shape == (1, 512, 512, 3)
         A_      :				Tuple		     =					np.array([0.1_2_3_7, 0.1_3_2_0, 0.1_4_3_8, 0.1_3_5_9, 0.1_3_9_0, 0.1_1_3_2, 0.1_2_7_7, 0.1_1_7_5, 0.1_1_1_2]		)

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







       def        _snake_case       (      self		)->Tuple:




         '''simple docstring'''





         A_      :				Dict		     =					StableDiffusionKDiffusionPipeline.from_pretrained('''stabilityai/stable-diffusion-2-1-base'''		)
         A_      :				Optional[int]		     =					sd_pipe.to(_SCREAMING_SNAKE_CASE		)
         sd_pipe.set_progress_bar_config(disable=_SCREAMING_SNAKE_CASE		)

         sd_pipe.set_scheduler('''sample_dpmpp_2m'''		)

         A_      :				Optional[Any]		     =					'''A painting of a squirrel eating a burger'''
         A_      :				int		     =					torch.manual_seed(0		)
         A_      :				Dict		     =					sd_pipe(
             [prompt]				,       generator=_SCREAMING_SNAKE_CASE				,       guidance_scale=7.5				,       num_inference_steps=15				,       output_type='''np'''				,       use_karras_sigmas=_SCREAMING_SNAKE_CASE				,       )

         A_      :				Tuple		     =					output.images

         A_      :				Tuple		     =					image[0, -3:, -3:, -1]

         assert image.shape == (1, 512, 512, 3)
         A_      :				Optional[int]		     =					np.array(
             [0.1_1_3_8_1_6_8_9, 0.1_2_1_1_2_9_2_1, 0.1_3_8_9_4_5_7, 0.1_2_5_4_9_6_0_6, 0.1_2_4_4_9_6_4, 0.1_0_8_3_1_5_1_7, 0.1_1_5_6_2_8_6_6, 0.1_0_8_6_7_8_1_6, 0.1_0_4_9_9_0_4_8]		)

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








import unittest

import numpy as np
import torch

from diffusers import DDIMPipeline, DDIMScheduler, UNetaDModel
from diffusers.utils.testing_utils import enable_full_determinism, require_torch_gpu, slow, torch_device

from ..pipeline_params import UNCONDITIONAL_IMAGE_GENERATION_BATCH_PARAMS, UNCONDITIONAL_IMAGE_GENERATION_PARAMS
from ..test_pipelines_common import PipelineTesterMixin


enable_full_determinism()






class   _lowerCamelCase (					UpperCamelCase					,							unittest.TestCase			):




       """simple docstring"""
       snake_case      =    DDIMPipeline
       snake_case      =    UNCONDITIONAL_IMAGE_GENERATION_PARAMS
       snake_case      =    PipelineTesterMixin.required_optional_params - {
           "num_images_per_prompt",
           "latents",
           "callback",
           "callback_steps",
       }
       snake_case      =    UNCONDITIONAL_IMAGE_GENERATION_BATCH_PARAMS
       snake_case      =    False



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




         '''simple docstring'''





         torch.manual_seed(0		)
         A_      :				List[str]		     =					UNetaDModel(
             block_out_channels=(32, 64)				,       layers_per_block=2				,       sample_size=32				,       in_channels=3				,       out_channels=3				,       down_block_types=('''DownBlock2D''', '''AttnDownBlock2D''')				,       up_block_types=('''AttnUpBlock2D''', '''UpBlock2D''')				,       )
         A_      :				Optional[Any]		     =					DDIMScheduler()
         A_      :				str		     =					{'''unet''': unet, '''scheduler''': scheduler}
         return components



       def        _snake_case       (      self				,       _SCREAMING_SNAKE_CASE				,       _SCREAMING_SNAKE_CASE=0		)->Optional[Any]:




         '''simple docstring'''





         if str(_SCREAMING_SNAKE_CASE		).startswith('''mps'''		):
           A_      :				Any		     =					torch.manual_seed(_SCREAMING_SNAKE_CASE		)
         else:
           A_      :				Optional[int]		     =					torch.Generator(device=_SCREAMING_SNAKE_CASE		).manual_seed(_SCREAMING_SNAKE_CASE		)
         A_      :				Any		     =					{
             '''batch_size''': 1,
             '''generator''': generator,
             '''num_inference_steps''': 2,
             '''output_type''': '''numpy''',
         }
         return inputs



       def        _snake_case       (      self		)->List[Any]:




         '''simple docstring'''





         A_      :				Optional[int]		     =					'''cpu'''

         A_      :				Dict		     =					self.get_dummy_components()
         A_      :				str		     =					self.pipeline_class(**_SCREAMING_SNAKE_CASE		)
         pipe.to(_SCREAMING_SNAKE_CASE		)
         pipe.set_progress_bar_config(disable=_SCREAMING_SNAKE_CASE		)

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

         self.assertEqual(image.shape				,       (1, 32, 32, 3)		)
         A_      :				List[Any]		     =					np.array(
             [1.000e00, 5.717e-01, 4.717e-01, 1.000e00, 0.000e00, 1.000e00, 3.000e-04, 0.000e00, 9.000e-04]		)
         A_      :				str		     =					np.abs(image_slice.flatten() - expected_slice		).max()
         self.assertLessEqual(_SCREAMING_SNAKE_CASE				,       1e-3		)



       def        _snake_case       (      self		)->Union[str, Any]:




         '''simple docstring'''





         super().test_dict_tuple_outputs_equivalent(expected_max_difference=3e-3		)



       def        _snake_case       (      self		)->Optional[int]:




         '''simple docstring'''





         super().test_save_load_local(expected_max_difference=3e-3		)



       def        _snake_case       (      self		)->Optional[int]:




         '''simple docstring'''





         super().test_save_load_optional_components(expected_max_difference=3e-3		)







       def        _snake_case       (      self		)->Any:




         '''simple docstring'''





         super().test_inference_batch_single_identical(expected_max_diff=3e-3		)






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




       """simple docstring"""



       def        _snake_case       (      self		)->Union[str, Any]:




         '''simple docstring'''





         A_      :				int		     =					'''google/ddpm-cifar10-32'''

         A_      :				Tuple		     =					UNetaDModel.from_pretrained(_SCREAMING_SNAKE_CASE		)
         A_      :				str		     =					DDIMScheduler()

         A_      :				str		     =					DDIMPipeline(unet=_SCREAMING_SNAKE_CASE				,       scheduler=_SCREAMING_SNAKE_CASE		)
         ddim.to(_SCREAMING_SNAKE_CASE		)
         ddim.set_progress_bar_config(disable=_SCREAMING_SNAKE_CASE		)

         A_      :				Optional[int]		     =					torch.manual_seed(0		)
         A_      :				Any		     =					ddim(generator=_SCREAMING_SNAKE_CASE				,       eta=0.0				,       output_type='''numpy'''		).images

         A_      :				Any		     =					image[0, -3:, -3:, -1]

         assert image.shape == (1, 32, 32, 3)
         A_      :				Any		     =					np.array([0.1_7_2_3, 0.1_6_1_7, 0.1_6_0_0, 0.1_6_2_6, 0.1_4_9_7, 0.1_5_1_3, 0.1_5_0_5, 0.1_4_4_2, 0.1_4_5_3]		)

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







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




         '''simple docstring'''





         A_      :				Tuple		     =					'''google/ddpm-ema-bedroom-256'''

         A_      :				int		     =					UNetaDModel.from_pretrained(_SCREAMING_SNAKE_CASE		)
         A_      :				Any		     =					DDIMScheduler.from_pretrained(_SCREAMING_SNAKE_CASE		)

         A_      :				Optional[Any]		     =					DDIMPipeline(unet=_SCREAMING_SNAKE_CASE				,       scheduler=_SCREAMING_SNAKE_CASE		)
         ddpm.to(_SCREAMING_SNAKE_CASE		)
         ddpm.set_progress_bar_config(disable=_SCREAMING_SNAKE_CASE		)

         A_      :				Dict		     =					torch.manual_seed(0		)
         A_      :				List[str]		     =					ddpm(generator=_SCREAMING_SNAKE_CASE				,       output_type='''numpy'''		).images

         A_      :				Any		     =					image[0, -3:, -3:, -1]

         assert image.shape == (1, 256, 256, 3)
         A_      :				Tuple		     =					np.array([0.0_0_6_0, 0.0_2_0_1, 0.0_3_4_4, 0.0_0_2_4, 0.0_0_1_8, 0.0_0_0_2, 0.0_0_2_2, 0.0_0_0_0, 0.0_0_6_9]		)

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




