

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  A						(							_UpperCAmelCase    ,							_UpperCAmelCase    ,							unittest.TestCase ):







    """simple docstring"""





    lowerCamelCase			=							IFPipeline
    lowerCamelCase			=							TEXT_TO_IMAGE_PARAMS - {'width', 'height', 'latents'}
    lowerCamelCase			=							TEXT_TO_IMAGE_BATCH_PARAMS
    lowerCamelCase			=							PipelineTesterMixin.required_optional_params - {'latents'}


    def 						snake_case__     (      self :       str		)->    List[Any]:





        '''simple docstring'''







        return self._get_dummy_components()


    def 						snake_case__     (      self :       Optional[Any],lowercase_ :       Tuple,lowercase_ :       List[str]=0		)->    List[str]:





        '''simple docstring'''







        if str(lowercase_		).startswith('mps'		):
            A__          =   torch.manual_seed(lowercase_		)
        else:
            A__          =   torch.Generator(device=lowercase_		).manual_seed(lowercase_		)

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

        return inputs


    def 						snake_case__     (      self :       List[str]		)->    Optional[int]:





        '''simple docstring'''







        self._test_save_load_optional_components()


    @unittest.skipIf(torch_device != 'cuda',reason='float16 requires CUDA'		)
    def 						snake_case__     (      self :       Optional[Any]		)->    Union[str, Any]:





        '''simple docstring'''







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


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





        '''simple docstring'''







        self._test_attention_slicing_forward_pass(expected_max_diff=1E-2		)


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





        '''simple docstring'''







        self._test_save_load_local()


    def 						snake_case__     (      self :       str		)->    Tuple:





        '''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 						snake_case__     (      self :       Any		)->    Any:





        '''simple docstring'''







        self._test_xformers_attention_forwardGenerator_pass(expected_max_diff=1E-3		)




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







    """simple docstring"""





    def 						snake_case__     (      self :       Any		)->    Optional[int]:





        '''simple docstring'''







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


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





        '''simple docstring'''







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

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

        # pre compute text embeddings and remove T5 to save memory

        pipe_a.text_encoder.to('cuda'		)

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

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

        A__          =   None
        A__          =   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(lowercase_,lowercase_,lowercase_,lowercase_		)

        pipe_a.remove_all_hooks()
        pipe_a.remove_all_hooks()

        # img2img

        A__          =   IFImgaImgPipeline(**pipe_a.components		)
        A__          =   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(lowercase_,lowercase_,lowercase_,lowercase_		)

        pipe_a.remove_all_hooks()
        pipe_a.remove_all_hooks()

        # inpainting

        A__          =   IFInpaintingPipeline(**pipe_a.components		)
        A__          =   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(lowercase_,lowercase_,lowercase_,lowercase_		)


    def 						snake_case__     (      self :       Tuple,lowercase_ :       List[Any],lowercase_ :       Optional[int],lowercase_ :       Any,lowercase_ :       Optional[Any]		)->    Dict:





        '''simple docstring'''







        _start_torch_memory_measurement()

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

        A__          =   output.images[0]

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

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

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

        # pipeline 2

        _start_torch_memory_measurement()

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

        A__          =   floats_tensor((1, 3, 6_4, 6_4),rng=random.Random(0		)		).to(lowercase_		)

        A__          =   pipe_a(
            prompt_embeds=lowercase_,negative_prompt_embeds=lowercase_,image=lowercase_,generator=lowercase_,num_inference_steps=2,output_type='np',)

        A__          =   output.images[0]

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

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

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


    def 						snake_case__     (      self :       Union[str, Any],lowercase_ :       List[Any],lowercase_ :       Optional[int],lowercase_ :       List[str],lowercase_ :       Any		)->    Optional[int]:





        '''simple docstring'''







        _start_torch_memory_measurement()

        A__          =   floats_tensor((1, 3, 6_4, 6_4),rng=random.Random(0		)		).to(lowercase_		)

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

        A__          =   pipe_a(
            prompt_embeds=lowercase_,negative_prompt_embeds=lowercase_,image=lowercase_,num_inference_steps=2,generator=lowercase_,output_type='np',)

        A__          =   output.images[0]

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

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

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

        # pipeline 2

        _start_torch_memory_measurement()

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

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

        A__          =   pipe_a(
            prompt_embeds=lowercase_,negative_prompt_embeds=lowercase_,image=lowercase_,original_image=lowercase_,generator=lowercase_,num_inference_steps=2,output_type='np',)

        A__          =   output.images[0]

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

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

        A__          =   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(lowercase_,lowercase_		)


    def 						snake_case__     (      self :       Optional[Any],lowercase_ :       str,lowercase_ :       int,lowercase_ :       List[str],lowercase_ :       int		)->    Tuple:





        '''simple docstring'''







        _start_torch_memory_measurement()

        A__          =   floats_tensor((1, 3, 6_4, 6_4),rng=random.Random(0		)		).to(lowercase_		)
        A__          =   floats_tensor((1, 3, 6_4, 6_4),rng=random.Random(1		)		).to(lowercase_		)

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

        A__          =   output.images[0]

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

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

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

        # pipeline 2

        _start_torch_memory_measurement()

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

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

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

        A__          =   output.images[0]

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

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

        A__          =   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(lowercase_,lowercase_		)




def 						_snake_case(    )							->      Dict:





    '''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  A						(							_UpperCAmelCase ):







    """simple docstring"""





    lowerCamelCase			=							(DPMSolverSinglestepScheduler,)
    lowerCamelCase			=							(('num_inference_steps', 25),)


    def 						snake_case__     (      self :       Tuple,**lowercase_ :       Dict		)->    Optional[int]:





        '''simple docstring'''







        A__          =   {
            'num_train_timesteps': 1_0_0_0,
            'beta_start': 0.0_001,
            'beta_end': 0.02,
            'beta_schedule': 'linear',
            'solver_order': 2,
            'prediction_type': 'epsilon',
            'thresholding': False,
            'sample_max_value': 1.0,
            'algorithm_type': 'dpmsolver++',
            'solver_type': 'midpoint',
            'lambda_min_clipped': -float('inf'		),
            'variance_type': None,
        }

        config.update(**lowercase_		)
        return config


    def 						snake_case__     (      self :       str,lowercase_ :       Optional[Any]=0,**lowercase_ :       Any		)->    List[Any]:





        '''simple docstring'''







        A__          =   dict(self.forward_default_kwargs		)
        A__          =   kwargs.pop('num_inference_steps',lowercase_		)
        A__          =   self.dummy_sample
        A__          =   0.1 * sample
        A__          =   [residual + 0.2, residual + 0.15, residual + 0.10]

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

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

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

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


    def 						snake_case__     (      self :       List[str]		)->    List[Any]:





        '''simple docstring'''







        pass


    def 						snake_case__     (      self :       Tuple,lowercase_ :       Union[str, Any]=0,**lowercase_ :       Union[str, Any]		)->    Union[str, Any]:





        '''simple docstring'''







        A__          =   dict(self.forward_default_kwargs		)
        A__          =   kwargs.pop('num_inference_steps',lowercase_		)
        A__          =   self.dummy_sample
        A__          =   0.1 * sample
        A__          =   [residual + 0.2, residual + 0.15, residual + 0.10]

        for scheduler_class in self.scheduler_classes:
            A__          =   self.get_scheduler_config()
            A__          =   scheduler_class(**lowercase_		)
            scheduler.set_timesteps(lowercase_		)

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

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

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

            A__          =   scheduler.step(lowercase_,lowercase_,lowercase_,**lowercase_		).prev_sample
            A__          =   new_scheduler.step(lowercase_,lowercase_,lowercase_,**lowercase_		).prev_sample

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


    def 						snake_case__     (      self :       Optional[Any],lowercase_ :       Optional[int]=None,**lowercase_ :       int		)->    int:





        '''simple docstring'''







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

        A__          =   self.scheduler_classes[0]
        A__          =   self.get_scheduler_config(**lowercase_		)
        A__          =   scheduler_class(**lowercase_		)

        A__          =   1_0
        A__          =   self.dummy_model()
        A__          =   self.dummy_sample_deter
        scheduler.set_timesteps(lowercase_		)

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

        return sample


    def 						snake_case__     (      self :       Any		)->    str:





        '''simple docstring'''







        A__          =   DPMSolverSinglestepScheduler(**self.get_scheduler_config()		)
        A__          =   5_0
        A__          =   self.dummy_model()
        A__          =   self.dummy_sample_deter
        scheduler.set_timesteps(lowercase_		)

        # make sure that the first t is uneven
        for i, t in enumerate(scheduler.timesteps[3:]		):
            A__          =   model(lowercase_,lowercase_		)
            A__          =   scheduler.step(lowercase_,lowercase_,lowercase_		).prev_sample

        A__          =   torch.mean(torch.abs(lowercase_		)		)

        assert abs(result_mean.item() - 0.2_574		) < 1E-3


    def 						snake_case__     (      self :       Optional[Any]		)->    List[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=lowercase_		)


    def 						snake_case__     (      self :       int		)->    Optional[Any]:





        '''simple docstring'''







        A__          =   DPMSolverSinglestepScheduler(**self.get_scheduler_config()		)
        A__          =   self.full_loop(scheduler=lowercase_		)
        A__          =   torch.mean(torch.abs(lowercase_		)		)

        assert abs(result_mean.item() - 0.2_791		) < 1E-3

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

        A__          =   self.full_loop(scheduler=lowercase_		)
        A__          =   torch.mean(torch.abs(lowercase_		)		)

        assert abs(result_mean.item() - 0.2_791		) < 1E-3


    def 						snake_case__     (      self :       Tuple		)->    Any:





        '''simple docstring'''







        self.check_over_configs(thresholding=lowercase_		)
        for order in [1, 2, 3]:
            for solver_type in ["midpoint", "heun"]:
                for threshold in [0.5, 1.0, 2.0]:
                    for prediction_type in ["epsilon", "sample"]:
                        self.check_over_configs(
                            thresholding=lowercase_,prediction_type=lowercase_,sample_max_value=lowercase_,algorithm_type='dpmsolver++',solver_order=lowercase_,solver_type=lowercase_,)


    def 						snake_case__     (      self :       List[Any]		)->    int:





        '''simple docstring'''







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


    def 						snake_case__     (      self :       Dict		)->    List[Any]:





        '''simple docstring'''







        for algorithm_type in ["dpmsolver", "dpmsolver++"]:
            for solver_type in ["midpoint", "heun"]:
                for order in [1, 2, 3]:
                    for prediction_type in ["epsilon", "sample"]:
                        self.check_over_configs(
                            solver_order=lowercase_,solver_type=lowercase_,prediction_type=lowercase_,algorithm_type=lowercase_,)
                        A__          =   self.full_loop(
                            solver_order=lowercase_,solver_type=lowercase_,prediction_type=lowercase_,algorithm_type=lowercase_,)
                        assert not torch.isnan(lowercase_		).any(), "Samples have nan numbers"


    def 						snake_case__     (      self :       Optional[int]		)->    Tuple:





        '''simple docstring'''







        self.check_over_configs(lower_order_final=lowercase_		)
        self.check_over_configs(lower_order_final=lowercase_		)


    def 						snake_case__     (      self :       Tuple		)->    Optional[int]:





        '''simple docstring'''







        self.check_over_configs(lambda_min_clipped=-float('inf'		)		)
        self.check_over_configs(lambda_min_clipped=-5.1		)


    def 						snake_case__     (      self :       Optional[Any]		)->    Tuple:





        '''simple docstring'''







        self.check_over_configs(variance_type=lowercase_		)
        self.check_over_configs(variance_type='learned_range'		)


    def 						snake_case__     (      self :       str		)->    Any:





        '''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=lowercase_,time_step=0		)


    def 						snake_case__     (      self :       Tuple		)->    Tuple:





        '''simple docstring'''







        A__          =   self.full_loop()
        A__          =   torch.mean(torch.abs(lowercase_		)		)

        assert abs(result_mean.item() - 0.2_791		) < 1E-3


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





        '''simple docstring'''







        A__          =   self.full_loop(use_karras_sigmas=lowercase_		)
        A__          =   torch.mean(torch.abs(lowercase_		)		)

        assert abs(result_mean.item() - 0.2_248		) < 1E-3


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





        '''simple docstring'''







        A__          =   self.full_loop(prediction_type='v_prediction'		)
        A__          =   torch.mean(torch.abs(lowercase_		)		)

        assert abs(result_mean.item() - 0.1_453		) < 1E-3


    def 						snake_case__     (      self :       Tuple		)->    int:





        '''simple docstring'''







        A__          =   self.full_loop(prediction_type='v_prediction',use_karras_sigmas=lowercase_		)
        A__          =   torch.mean(torch.abs(lowercase_		)		)

        assert abs(result_mean.item() - 0.0_649		) < 1E-3


    def 						snake_case__     (      self :       List[Any]		)->    int:





        '''simple docstring'''







        A__          =   self.scheduler_classes[0]
        A__          =   self.get_scheduler_config(thresholding=lowercase_,dynamic_thresholding_ratio=0		)
        A__          =   scheduler_class(**lowercase_		)

        A__          =   1_0
        A__          =   self.dummy_model()
        A__          =   self.dummy_sample_deter.half()
        scheduler.set_timesteps(lowercase_		)

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

        assert sample.dtype == torch.floataa

