

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__  (					__magic_name__				,							__magic_name__				,							unittest.TestCase			):



       '''simple docstring'''







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

       def     __UpperCamelCase     (							self     ):







             '''simple docstring'''



             return self._get_dummy_components()

       def     __UpperCamelCase     (							self ,   lowercase__ ,   lowercase__=0     ):







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







             '''simple docstring'''



             self._test_save_load_optional_components()

       @unittest.skipIf(torch_device != '''cuda''' ,   reason='''float16 requires CUDA'''     )
       def     __UpperCamelCase     (							self     ):







             '''simple docstring'''



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

       def     __UpperCamelCase     (							self     ):







             '''simple docstring'''



             self._test_attention_slicing_forward_pass(expected_max_diff=1E-2     )

       def     __UpperCamelCase     (							self     ):







             '''simple docstring'''



             self._test_save_load_local()

       def     __UpperCamelCase     (							self     ):







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







             '''simple docstring'''



             self._test_xformers_attention_forwardGenerator_pass(expected_max_diff=1E-3     )





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



       '''simple docstring'''







       def     __UpperCamelCase     (							self     ):







             '''simple docstring'''



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

       def     __UpperCamelCase     (							self     ):







             '''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     __UpperCamelCase     (							self ,   lowercase__ ,   lowercase__ ,   lowercase__ ,   lowercase__     ):







             '''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     __UpperCamelCase     (							self ,   lowercase__ ,   lowercase__ ,   lowercase__ ,   lowercase__     ):







             '''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     __UpperCamelCase     (							self ,   lowercase__ ,   lowercase__ ,   lowercase__ ,   lowercase__     ):







             '''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      A__					(			)					->Optional[int]:
      torch.cuda.empty_cache()
      torch.cuda.reset_max_memory_allocated()
      torch.cuda.reset_peak_memory_stats()





from typing import Dict, List, Optional, Tuple, Union

import numpy as np

from ...image_processing_utils import BaseImageProcessor, BatchFeature, get_size_dict
from ...image_transforms import (
    center_crop,
    flip_channel_order,
    get_resize_output_image_size,
    rescale,
    resize,
    to_channel_dimension_format,
)
from ...image_utils import (
    ChannelDimension,
    ImageInput,
    PILImageResampling,
    make_list_of_images,
    to_numpy_array,
    valid_images,
)
from ...utils import TensorType, is_torch_available, is_torch_tensor, is_vision_available, logging


if is_vision_available():
   import PIL

if is_torch_available():
   import torch


_lowerCamelCase      :					Dict					    =      logging.get_logger(__name__)





class        lowerCAmelCase__  (					__magic_name__			):



       '''simple docstring'''







       lowercase_						=    ["""pixel_values"""]

       def __init__(							self ,   lowercase__ = True ,   lowercase__ = None ,   lowercase__ = PILImageResampling.BILINEAR ,   lowercase__ = True ,   lowercase__ = 1 / 2_5_5 ,   lowercase__ = True ,   lowercase__ = None ,   lowercase__ = True ,   **lowercase__ ,   ):







             '''simple docstring'''



             super().__init__(**lowercase__     )
             __A    =size if size is not None else {'''shortest_edge''': 2_2_4}
             __A    =get_size_dict(lowercase__ ,   default_to_square=lowercase__     )
             __A    =crop_size if crop_size is not None else {'''height''': 2_5_6, '''width''': 2_5_6}
             __A    =get_size_dict(lowercase__ ,   param_name='''crop_size'''     )

             __A    =do_resize
             __A    =size
             __A    =resample
             __A    =do_rescale
             __A    =rescale_factor
             __A    =do_center_crop
             __A    =crop_size
             __A    =do_flip_channel_order

       def     __UpperCamelCase     (							self ,   lowercase__ ,   lowercase__ ,   lowercase__ = PIL.Image.BILINEAR ,   lowercase__ = None ,   **lowercase__ ,   ):







             '''simple docstring'''



             __A    =get_size_dict(lowercase__ ,   default_to_square=lowercase__     )
             if "shortest_edge" not in size:
                   raise ValueError(f'''The `size` dictionary must contain the key `shortest_edge`. Got {size.keys()}'''     )
             __A    =get_resize_output_image_size(lowercase__ ,   size=size['''shortest_edge'''] ,   default_to_square=lowercase__     )
             return resize(lowercase__ ,   size=lowercase__ ,   resample=lowercase__ ,   data_format=lowercase__ ,   **lowercase__     )

       def     __UpperCamelCase     (							self ,   lowercase__ ,   lowercase__ ,   lowercase__ = None ,   **lowercase__ ,   ):







             '''simple docstring'''



             __A    =get_size_dict(lowercase__     )
             if "height" not in size or "width" not in size:
                   raise ValueError(f'''The `size` dictionary must contain the keys `height` and `width`. Got {size.keys()}'''     )
             return center_crop(lowercase__ ,   size=(size['''height'''], size['''width''']) ,   data_format=lowercase__ ,   **lowercase__     )

       def     __UpperCamelCase     (							self ,   lowercase__ ,   lowercase__ ,   lowercase__ = None ,   **lowercase__ ,   ):







             '''simple docstring'''



             return rescale(lowercase__ ,   scale=lowercase__ ,   data_format=lowercase__ ,   **lowercase__     )

       def     __UpperCamelCase     (							self ,   lowercase__ ,   lowercase__ = None     ):







             '''simple docstring'''



             return flip_channel_order(lowercase__ ,   data_format=lowercase__     )

       def     __UpperCamelCase     (							self ,   lowercase__ ,   lowercase__ = None ,   lowercase__ = None ,   lowercase__ = None ,   lowercase__ = None ,   lowercase__ = None ,   lowercase__ = None ,   lowercase__ = None ,   lowercase__ = None ,   lowercase__ = None ,   lowercase__ = ChannelDimension.FIRST ,   **lowercase__ ,   ):







             '''simple docstring'''



             __A    =do_resize if do_resize is not None else self.do_resize
             __A    =resample if resample is not None else self.resample
             __A    =do_rescale if do_rescale is not None else self.do_rescale
             __A    =rescale_factor if rescale_factor is not None else self.rescale_factor
             __A    =do_center_crop if do_center_crop is not None else self.do_center_crop
             __A    =(
                 do_flip_channel_order if do_flip_channel_order is not None else self.do_flip_channel_order
             )

             __A    =size if size is not None else self.size
             __A    =get_size_dict(lowercase__ ,   default_to_square=lowercase__     )
             __A    =crop_size if crop_size is not None else self.crop_size
             __A    =get_size_dict(lowercase__ ,   param_name='''crop_size'''     )

             __A    =make_list_of_images(lowercase__     )

             if not valid_images(lowercase__     ):
                   raise ValueError(
                       '''Invalid image type. Must be of type PIL.Image.Image, numpy.ndarray, '''
                       '''torch.Tensor, tf.Tensor or jax.ndarray.'''     )

             if do_resize and size is None:
                   raise ValueError('''Size must be specified if do_resize is True.'''     )

             if do_rescale and rescale_factor is None:
                   raise ValueError('''Rescale factor must be specified if do_rescale is True.'''     )

             if do_center_crop and crop_size is None:
                   raise ValueError('''Crop size must be specified if do_center_crop is True.'''     )

             # All transformations expect numpy arrays.
             __A    =[to_numpy_array(lowercase__     ) for image in images]

             if do_resize:
                   __A    =[self.resize(image=lowercase__ ,   size=lowercase__ ,   resample=lowercase__     ) for image in images]

             if do_center_crop:
                   __A    =[self.center_crop(image=lowercase__ ,   size=lowercase__     ) for image in images]

             if do_rescale:
                   __A    =[self.rescale(image=lowercase__ ,   scale=lowercase__     ) for image in images]

             # the pretrained checkpoints assume images are BGR, not RGB
             if do_flip_channel_order:
                   __A    =[self.flip_channel_order(image=lowercase__     ) for image in images]

             __A    =[to_channel_dimension_format(lowercase__ ,   lowercase__     ) for image in images]

             __A    ={'''pixel_values''': images}
             return BatchFeature(data=lowercase__ ,   tensor_type=lowercase__     )


       def     __UpperCamelCase     (							self ,   lowercase__ ,   lowercase__ = None     ):







             '''simple docstring'''



             __A    =outputs.logits

             # Resize logits and compute semantic segmentation maps
             if target_sizes is not None:
                   if len(lowercase__     ) != len(lowercase__     ):
                         raise ValueError(
                             '''Make sure that you pass in as many target sizes as the batch dimension of the logits'''     )

                   if is_torch_tensor(lowercase__     ):
                         __A    =target_sizes.numpy()

                   __A    =[]

                   for idx in range(len(lowercase__     )     ):
                         __A    =torch.nn.functional.interpolate(
                             logits[idx].unsqueeze(dim=0     ) ,   size=target_sizes[idx] ,   mode='''bilinear''' ,   align_corners=lowercase__     )
                         __A    =resized_logits[0].argmax(dim=0     )
                         semantic_segmentation.append(lowercase__     )
             else:
                   __A    =logits.argmax(dim=1     )
                   __A    =[semantic_segmentation[i] for i in range(semantic_segmentation.shape[0]     )]

             return semantic_segmentation


