


import gc
import random
import unittest

import torch

from diffusers import (
    IFImgaImgPipeline,
    IFImgaImgSuperResolutionPipeline,
    IFInpaintingPipeline,
    IFInpaintingSuperResolutionPipeline,
    IFPipeline,
    IFSuperResolutionPipeline,
)
from diffusers.models.attention_processor import AttnAddedKVProcessor
from diffusers.utils.import_utils import is_xformers_available
from diffusers.utils.testing_utils import floats_tensor, load_numpy, require_torch_gpu, skip_mps, slow, torch_device

from ..pipeline_params import TEXT_TO_IMAGE_BATCH_PARAMS, TEXT_TO_IMAGE_PARAMS
from ..test_pipelines_common import PipelineTesterMixin, assert_mean_pixel_difference
from . import IFPipelineTesterMixin
@skip_mps
class 	lowerCamelCase_       (				lowerCAmelCase__     ,     lowerCAmelCase__     ,     unittest.TestCase				):



     '''simple docstring'''







     __UpperCAmelCase            =				IFPipeline
     __UpperCAmelCase            =				TEXT_TO_IMAGE_PARAMS - {"width", "height", "latents"}
     __UpperCAmelCase            =				TEXT_TO_IMAGE_BATCH_PARAMS
     __UpperCAmelCase            =				PipelineTesterMixin.required_optional_params - {"latents"}






     def A      (				self						)   ->			str:







        '''simple docstring'''

        return self._get_dummy_components()






     def A      (				self ,       snake_case_ ,       snake_case_=0						)   ->			Union[str, Any]:







        '''simple docstring'''

        if str(__A						).startswith('''mps'''						):
           __lowercase 					=   torch.manual_seed(__A						)
        else:
           __lowercase 					=   torch.Generator(device=__A						).manual_seed(__A						)

        __lowercase 					=   {
            '''prompt''': '''A painting of a squirrel eating a burger''',
            '''generator''': generator,
            '''num_inference_steps''': 2,
            '''output_type''': '''numpy''',
        }

        return inputs






     def A      (				self						)   ->			int:







        '''simple docstring'''

        self._test_save_load_optional_components()






     @unittest.skipIf(torch_device != '''cuda''' ,       reason='''float16 requires CUDA'''						)
     def A      (				self						)   ->			Tuple:







        '''simple docstring'''

        super().test_save_load_floataa(expected_max_diff=1e-1						)






     def A      (				self						)   ->			Optional[Any]:







        '''simple docstring'''

        self._test_attention_slicing_forward_pass(expected_max_diff=1e-2						)






     def A      (				self						)   ->			Dict:







        '''simple docstring'''

        self._test_save_load_local()






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







        '''simple docstring'''

        self._test_inference_batch_single_identical(
            expected_max_diff=1e-2 ,       )






     @unittest.skipIf(
         torch_device != '''cuda''' or not is_xformers_available() ,       reason='''XFormers attention is only available with CUDA and `xformers` installed''' ,       )
     def A      (				self						)   ->			List[str]:







        '''simple docstring'''

        self._test_xformers_attention_forwardGenerator_pass(expected_max_diff=1e-3						)




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



     '''simple docstring'''







     def A      (				self						)   ->			Any:







        '''simple docstring'''

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






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







        '''simple docstring'''

        __lowercase 					=   IFPipeline.from_pretrained('''DeepFloyd/IF-I-XL-v1.0''' ,       variant='''fp16''' ,       torch_dtype=torch.floataa						)

        __lowercase 					=   IFSuperResolutionPipeline.from_pretrained(
            '''DeepFloyd/IF-II-L-v1.0''' ,       variant='''fp16''' ,       torch_dtype=torch.floataa ,       text_encoder=__A ,       tokenizer=__A						)

        # pre compute text embeddings and remove T5 to save memory

        pipe_a.text_encoder.to('''cuda'''						)

        __lowercase					, __lowercase 					=   pipe_a.encode_prompt('''anime turtle''' ,       device='''cuda'''						)

        del pipe_a.tokenizer
        del pipe_a.text_encoder
        gc.collect()

        __lowercase 					=   None
        __lowercase 					=   None

        pipe_a.enable_model_cpu_offload()
        pipe_a.enable_model_cpu_offload()

        pipe_a.unet.set_attn_processor(AttnAddedKVProcessor()						)
        pipe_a.unet.set_attn_processor(AttnAddedKVProcessor()						)

        self._test_if(__A ,       __A ,       __A ,       __A						)

        pipe_a.remove_all_hooks()
        pipe_a.remove_all_hooks()

        # img2img

        __lowercase 					=   IFImgaImgPipeline(**pipe_a.components						)
        __lowercase 					=   IFImgaImgSuperResolutionPipeline(**pipe_a.components						)

        pipe_a.enable_model_cpu_offload()
        pipe_a.enable_model_cpu_offload()

        pipe_a.unet.set_attn_processor(AttnAddedKVProcessor()						)
        pipe_a.unet.set_attn_processor(AttnAddedKVProcessor()						)

        self._test_if_imgaimg(__A ,       __A ,       __A ,       __A						)

        pipe_a.remove_all_hooks()
        pipe_a.remove_all_hooks()

        # inpainting

        __lowercase 					=   IFInpaintingPipeline(**pipe_a.components						)
        __lowercase 					=   IFInpaintingSuperResolutionPipeline(**pipe_a.components						)

        pipe_a.enable_model_cpu_offload()
        pipe_a.enable_model_cpu_offload()

        pipe_a.unet.set_attn_processor(AttnAddedKVProcessor()						)
        pipe_a.unet.set_attn_processor(AttnAddedKVProcessor()						)

        self._test_if_inpainting(__A ,       __A ,       __A ,       __A						)






     def A      (				self ,       snake_case_ ,       snake_case_ ,       snake_case_ ,       snake_case_						)   ->			Optional[int]:







        '''simple docstring'''

        _start_torch_memory_measurement()

        __lowercase 					=   torch.Generator(device='''cpu'''						).manual_seed(0						)
        __lowercase 					=   pipe_a(
            prompt_embeds=__A ,       negative_prompt_embeds=__A ,       num_inference_steps=2 ,       generator=__A ,       output_type='''np''' ,       )

        __lowercase 					=   output.images[0]

        assert image.shape == (6_4, 6_4, 3)

        __lowercase 					=   torch.cuda.max_memory_allocated()
        assert mem_bytes < 1_3 * 1_0**9

        __lowercase 					=   load_numpy(
            '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/if/test_if.npy'''						)
        assert_mean_pixel_difference(__A ,       __A						)

        # pipeline 2

        _start_torch_memory_measurement()

        __lowercase 					=   torch.Generator(device='''cpu'''						).manual_seed(0						)

        __lowercase 					=   floats_tensor((1, 3, 6_4, 6_4) ,       rng=random.Random(0						)						).to(__A						)

        __lowercase 					=   pipe_a(
            prompt_embeds=__A ,       negative_prompt_embeds=__A ,       image=__A ,       generator=__A ,       num_inference_steps=2 ,       output_type='''np''' ,       )

        __lowercase 					=   output.images[0]

        assert image.shape == (2_5_6, 2_5_6, 3)

        __lowercase 					=   torch.cuda.max_memory_allocated()
        assert mem_bytes < 4 * 1_0**9

        __lowercase 					=   load_numpy(
            '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/if/test_if_superresolution_stage_II.npy'''						)
        assert_mean_pixel_difference(__A ,       __A						)






     def A      (				self ,       snake_case_ ,       snake_case_ ,       snake_case_ ,       snake_case_						)   ->			List[Any]:







        '''simple docstring'''

        _start_torch_memory_measurement()

        __lowercase 					=   floats_tensor((1, 3, 6_4, 6_4) ,       rng=random.Random(0						)						).to(__A						)

        __lowercase 					=   torch.Generator(device='''cpu'''						).manual_seed(0						)

        __lowercase 					=   pipe_a(
            prompt_embeds=__A ,       negative_prompt_embeds=__A ,       image=__A ,       num_inference_steps=2 ,       generator=__A ,       output_type='''np''' ,       )

        __lowercase 					=   output.images[0]

        assert image.shape == (6_4, 6_4, 3)

        __lowercase 					=   torch.cuda.max_memory_allocated()
        assert mem_bytes < 1_0 * 1_0**9

        __lowercase 					=   load_numpy(
            '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/if/test_if_img2img.npy'''						)
        assert_mean_pixel_difference(__A ,       __A						)

        # pipeline 2

        _start_torch_memory_measurement()

        __lowercase 					=   torch.Generator(device='''cpu'''						).manual_seed(0						)

        __lowercase 					=   floats_tensor((1, 3, 2_5_6, 2_5_6) ,       rng=random.Random(0						)						).to(__A						)
        __lowercase 					=   floats_tensor((1, 3, 6_4, 6_4) ,       rng=random.Random(0						)						).to(__A						)

        __lowercase 					=   pipe_a(
            prompt_embeds=__A ,       negative_prompt_embeds=__A ,       image=__A ,       original_image=__A ,       generator=__A ,       num_inference_steps=2 ,       output_type='''np''' ,       )

        __lowercase 					=   output.images[0]

        assert image.shape == (2_5_6, 2_5_6, 3)

        __lowercase 					=   torch.cuda.max_memory_allocated()
        assert mem_bytes < 4 * 1_0**9

        __lowercase 					=   load_numpy(
            '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/if/test_if_img2img_superresolution_stage_II.npy'''						)
        assert_mean_pixel_difference(__A ,       __A						)






     def A      (				self ,       snake_case_ ,       snake_case_ ,       snake_case_ ,       snake_case_						)   ->			Optional[int]:







        '''simple docstring'''

        _start_torch_memory_measurement()

        __lowercase 					=   floats_tensor((1, 3, 6_4, 6_4) ,       rng=random.Random(0						)						).to(__A						)
        __lowercase 					=   floats_tensor((1, 3, 6_4, 6_4) ,       rng=random.Random(1						)						).to(__A						)

        __lowercase 					=   torch.Generator(device='''cpu'''						).manual_seed(0						)
        __lowercase 					=   pipe_a(
            prompt_embeds=__A ,       negative_prompt_embeds=__A ,       image=__A ,       mask_image=__A ,       num_inference_steps=2 ,       generator=__A ,       output_type='''np''' ,       )

        __lowercase 					=   output.images[0]

        assert image.shape == (6_4, 6_4, 3)

        __lowercase 					=   torch.cuda.max_memory_allocated()
        assert mem_bytes < 1_0 * 1_0**9

        __lowercase 					=   load_numpy(
            '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/if/test_if_inpainting.npy'''						)
        assert_mean_pixel_difference(__A ,       __A						)

        # pipeline 2

        _start_torch_memory_measurement()

        __lowercase 					=   torch.Generator(device='''cpu'''						).manual_seed(0						)

        __lowercase 					=   floats_tensor((1, 3, 6_4, 6_4) ,       rng=random.Random(0						)						).to(__A						)
        __lowercase 					=   floats_tensor((1, 3, 2_5_6, 2_5_6) ,       rng=random.Random(0						)						).to(__A						)
        __lowercase 					=   floats_tensor((1, 3, 2_5_6, 2_5_6) ,       rng=random.Random(1						)						).to(__A						)

        __lowercase 					=   pipe_a(
            prompt_embeds=__A ,       negative_prompt_embeds=__A ,       image=__A ,       mask_image=__A ,       original_image=__A ,       generator=__A ,       num_inference_steps=2 ,       output_type='''np''' ,       )

        __lowercase 					=   output.images[0]

        assert image.shape == (2_5_6, 2_5_6, 3)

        __lowercase 					=   torch.cuda.max_memory_allocated()
        assert mem_bytes < 4 * 1_0**9

        __lowercase 					=   load_numpy(
            '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/if/test_if_inpainting_superresolution_stage_II.npy'''						)
        assert_mean_pixel_difference(__A ,       __A						)




def 		lowercase_  (     ):


   '''simple docstring'''







   torch.cuda.empty_cache()
   torch.cuda.reset_max_memory_allocated()
   torch.cuda.reset_peak_memory_stats()






import tempfile

import torch

from diffusers import (
    DEISMultistepScheduler,
    DPMSolverMultistepScheduler,
    DPMSolverSinglestepScheduler,
    UniPCMultistepScheduler,
)

from .test_schedulers import SchedulerCommonTest
class 	lowerCamelCase_       (				lowerCAmelCase__				):



					'''simple docstring'''







					__UpperCAmelCase            =				(DEISMultistepScheduler,)
					__UpperCAmelCase            =				(("num_inference_steps", 25),)






					def A      (				self ,       **snake_case_						)   ->			int:







								'''simple docstring'''

								__lowercase 					=   {
								    '''num_train_timesteps''': 1_0_0_0,
								    '''beta_start''': 0.0_0_0_1,
								    '''beta_end''': 0.0_2,
								    '''beta_schedule''': '''linear''',
								    '''solver_order''': 2,
								}

								config.update(**snake_case_						)
								return config






					def A      (				self ,       snake_case_=0 ,       **snake_case_						)   ->			List[Any]:







								'''simple docstring'''

								__lowercase 					=   dict(self.forward_default_kwargs						)
								__lowercase 					=   kwargs.pop('''num_inference_steps''' ,       snake_case_						)
								__lowercase 					=   self.dummy_sample
								__lowercase 					=   0.1 * sample
								__lowercase 					=   [residual + 0.2, residual + 0.1_5, residual + 0.1_0]

								for scheduler_class in self.scheduler_classes:
											__lowercase 					=   self.get_scheduler_config(**snake_case_						)
											__lowercase 					=   scheduler_class(**snake_case_						)
											scheduler.set_timesteps(snake_case_						)
											# copy over dummy past residuals
											__lowercase 					=   dummy_past_residuals[: scheduler.config.solver_order]

											with tempfile.TemporaryDirectory() as tmpdirname:
														scheduler.save_config(snake_case_						)
														__lowercase 					=   scheduler_class.from_pretrained(snake_case_						)
														new_scheduler.set_timesteps(snake_case_						)
														# copy over dummy past residuals
														__lowercase 					=   dummy_past_residuals[: new_scheduler.config.solver_order]

											__lowercase					, __lowercase 					=   sample, sample
											for t in range(snake_case_ ,       time_step + scheduler.config.solver_order + 1						):
														__lowercase 					=   scheduler.step(snake_case_ ,       snake_case_ ,       snake_case_ ,       **snake_case_						).prev_sample
														__lowercase 					=   new_scheduler.step(snake_case_ ,       snake_case_ ,       snake_case_ ,       **snake_case_						).prev_sample

														assert torch.sum(torch.abs(output - new_output						)						) < 1e-5, "Scheduler outputs are not identical"






					def A      (				self						)   ->			Dict:







								'''simple docstring'''

								pass






					def A      (				self ,       snake_case_=0 ,       **snake_case_						)   ->			List[str]:







								'''simple docstring'''

								__lowercase 					=   dict(self.forward_default_kwargs						)
								__lowercase 					=   kwargs.pop('''num_inference_steps''' ,       snake_case_						)
								__lowercase 					=   self.dummy_sample
								__lowercase 					=   0.1 * sample
								__lowercase 					=   [residual + 0.2, residual + 0.1_5, residual + 0.1_0]

								for scheduler_class in self.scheduler_classes:
											__lowercase 					=   self.get_scheduler_config()
											__lowercase 					=   scheduler_class(**snake_case_						)
											scheduler.set_timesteps(snake_case_						)

											# copy over dummy past residuals (must be after setting timesteps)
											__lowercase 					=   dummy_past_residuals[: scheduler.config.solver_order]

											with tempfile.TemporaryDirectory() as tmpdirname:
														scheduler.save_config(snake_case_						)
														__lowercase 					=   scheduler_class.from_pretrained(snake_case_						)
														# copy over dummy past residuals
														new_scheduler.set_timesteps(snake_case_						)

														# copy over dummy past residual (must be after setting timesteps)
														__lowercase 					=   dummy_past_residuals[: new_scheduler.config.solver_order]

											__lowercase 					=   scheduler.step(snake_case_ ,       snake_case_ ,       snake_case_ ,       **snake_case_						).prev_sample
											__lowercase 					=   new_scheduler.step(snake_case_ ,       snake_case_ ,       snake_case_ ,       **snake_case_						).prev_sample

											assert torch.sum(torch.abs(output - new_output						)						) < 1e-5, "Scheduler outputs are not identical"






					def A      (				self ,       snake_case_=None ,       **snake_case_						)   ->			Optional[int]:







								'''simple docstring'''

								if scheduler is None:
											__lowercase 					=   self.scheduler_classes[0]
											__lowercase 					=   self.get_scheduler_config(**snake_case_						)
											__lowercase 					=   scheduler_class(**snake_case_						)

								__lowercase 					=   self.scheduler_classes[0]
								__lowercase 					=   self.get_scheduler_config(**snake_case_						)
								__lowercase 					=   scheduler_class(**snake_case_						)

								__lowercase 					=   1_0
								__lowercase 					=   self.dummy_model()
								__lowercase 					=   self.dummy_sample_deter
								scheduler.set_timesteps(snake_case_						)

								for i, t in enumerate(scheduler.timesteps						):
											__lowercase 					=   model(snake_case_ ,       snake_case_						)
											__lowercase 					=   scheduler.step(snake_case_ ,       snake_case_ ,       snake_case_						).prev_sample

								return sample






					def A      (				self						)   ->			str:







								'''simple docstring'''

								__lowercase 					=   dict(self.forward_default_kwargs						)

								__lowercase 					=   kwargs.pop('''num_inference_steps''' ,       snake_case_						)

								for scheduler_class in self.scheduler_classes:
											__lowercase 					=   self.get_scheduler_config()
											__lowercase 					=   scheduler_class(**snake_case_						)

											__lowercase 					=   self.dummy_sample
											__lowercase 					=   0.1 * sample

											if num_inference_steps is not None and hasattr(snake_case_ ,       '''set_timesteps'''						):
														scheduler.set_timesteps(snake_case_						)
											elif num_inference_steps is not None and not hasattr(snake_case_ ,       '''set_timesteps'''						):
														__lowercase 					=   num_inference_steps

											# copy over dummy past residuals (must be done after set_timesteps)
											__lowercase 					=   [residual + 0.2, residual + 0.1_5, residual + 0.1_0]
											__lowercase 					=   dummy_past_residuals[: scheduler.config.solver_order]

											__lowercase 					=   scheduler.timesteps[5]
											__lowercase 					=   scheduler.timesteps[6]

											__lowercase 					=   scheduler.step(snake_case_ ,       snake_case_ ,       snake_case_ ,       **snake_case_						).prev_sample
											__lowercase 					=   scheduler.step(snake_case_ ,       snake_case_ ,       snake_case_ ,       **snake_case_						).prev_sample

											self.assertEqual(output_a.shape ,       sample.shape						)
											self.assertEqual(output_a.shape ,       output_a.shape						)






					def A      (				self						)   ->			str:







								'''simple docstring'''

								__lowercase 					=   DEISMultistepScheduler(**self.get_scheduler_config()						)
								__lowercase 					=   self.full_loop(scheduler=snake_case_						)
								__lowercase 					=   torch.mean(torch.abs(snake_case_						)						)

								assert abs(result_mean.item() - 0.2_3_9_1_6						) < 1e-3

								__lowercase 					=   DPMSolverSinglestepScheduler.from_config(scheduler.config						)
								__lowercase 					=   DPMSolverMultistepScheduler.from_config(scheduler.config						)
								__lowercase 					=   UniPCMultistepScheduler.from_config(scheduler.config						)
								__lowercase 					=   DEISMultistepScheduler.from_config(scheduler.config						)

								__lowercase 					=   self.full_loop(scheduler=snake_case_						)
								__lowercase 					=   torch.mean(torch.abs(snake_case_						)						)

								assert abs(result_mean.item() - 0.2_3_9_1_6						) < 1e-3






					def A      (				self						)   ->			Optional[Any]:







								'''simple docstring'''

								for timesteps in [2_5, 5_0, 1_0_0, 9_9_9, 1_0_0_0]:
											self.check_over_configs(num_train_timesteps=snake_case_						)






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







								'''simple docstring'''

								self.check_over_configs(thresholding=snake_case_						)
								for order in [1, 2, 3]:
											for solver_type in ["logrho"]:
														for threshold in [0.5, 1.0, 2.0]:
																	for prediction_type in ["epsilon", "sample"]:
																				self.check_over_configs(
																				    thresholding=snake_case_ ,       prediction_type=snake_case_ ,       sample_max_value=snake_case_ ,       algorithm_type='''deis''' ,       solver_order=snake_case_ ,       solver_type=snake_case_ ,       )






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







								'''simple docstring'''

								for prediction_type in ["epsilon", "v_prediction"]:
											self.check_over_configs(prediction_type=snake_case_						)






					def A      (				self						)   ->			Optional[Any]:







								'''simple docstring'''

								for algorithm_type in ["deis"]:
											for solver_type in ["logrho"]:
														for order in [1, 2, 3]:
																	for prediction_type in ["epsilon", "sample"]:
																				self.check_over_configs(
																				    solver_order=snake_case_ ,       solver_type=snake_case_ ,       prediction_type=snake_case_ ,       algorithm_type=snake_case_ ,       )
																				__lowercase 					=   self.full_loop(
																				    solver_order=snake_case_ ,       solver_type=snake_case_ ,       prediction_type=snake_case_ ,       algorithm_type=snake_case_ ,       )
																				assert not torch.isnan(snake_case_						).any(), "Samples have nan numbers"






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







								'''simple docstring'''

								self.check_over_configs(lower_order_final=snake_case_						)
								self.check_over_configs(lower_order_final=snake_case_						)






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







								'''simple docstring'''

								for num_inference_steps in [1, 2, 3, 5, 1_0, 5_0, 1_0_0, 9_9_9, 1_0_0_0]:
											self.check_over_forward(num_inference_steps=snake_case_ ,       time_step=0						)






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







								'''simple docstring'''

								__lowercase 					=   self.full_loop()
								__lowercase 					=   torch.mean(torch.abs(snake_case_						)						)

								assert abs(result_mean.item() - 0.2_3_9_1_6						) < 1e-3






					def A      (				self						)   ->			Dict:







								'''simple docstring'''

								__lowercase 					=   self.full_loop(prediction_type='''v_prediction'''						)
								__lowercase 					=   torch.mean(torch.abs(snake_case_						)						)

								assert abs(result_mean.item() - 0.0_9_1						) < 1e-3






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







								'''simple docstring'''

								__lowercase 					=   self.scheduler_classes[0]
								__lowercase 					=   self.get_scheduler_config(thresholding=snake_case_ ,       dynamic_thresholding_ratio=0						)
								__lowercase 					=   scheduler_class(**snake_case_						)

								__lowercase 					=   1_0
								__lowercase 					=   self.dummy_model()
								__lowercase 					=   self.dummy_sample_deter.half()
								scheduler.set_timesteps(snake_case_						)

								for i, t in enumerate(scheduler.timesteps						):
											__lowercase 					=   model(snake_case_ ,       snake_case_						)
											__lowercase 					=   scheduler.step(snake_case_ ,       snake_case_ ,       snake_case_						).prev_sample

								assert sample.dtype == torch.floataa


