





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 ):
     __A	  =       IFPipeline
     __A	  =       TEXT_TO_IMAGE_PARAMS - {"width", "height", "latents"}
     __A	  =       TEXT_TO_IMAGE_BATCH_PARAMS
     __A	  =       PipelineTesterMixin.required_optional_params - {"latents"}

     def 	_UpperCamelCase     (      self      ):





            """simple docstring"""

            return self._get_dummy_components()

     def 	_UpperCamelCase     (      self,       A,       A=0      ):





            """simple docstring"""

            if str(A      ).startswith('''mps'''      ):
                   _UpperCamelCase              =						torch.manual_seed(A      )
            else:
                   _UpperCamelCase              =						torch.Generator(device=A      ).manual_seed(A      )

            _UpperCamelCase              =						{
                '''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  __A(						unittest.TestCase ):

     def 	_UpperCamelCase     (      self      ):





            """simple docstring"""

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

     def 	_UpperCamelCase     (      self      ):





            """simple docstring"""

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

            _UpperCamelCase              =						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'''      )

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

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

            _UpperCamelCase              =						None
            _UpperCamelCase              =						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

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

            _UpperCamelCase              =						IFInpaintingPipeline(**pipe_a.components      )
            _UpperCamelCase              =						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 	_UpperCamelCase     (      self,       A,       A,       A,       A      ):





            """simple docstring"""

            _start_torch_memory_measurement()

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

            _UpperCamelCase              =						output.images[0]

            assert image.shape == (64, 64, 3)

            _UpperCamelCase              =						torch.cuda.max_memory_allocated()
            assert mem_bytes < 13 * 10**9

            _UpperCamelCase              =						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()

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

            _UpperCamelCase              =						floats_tensor((1, 3, 64, 64),       rng=random.Random(0      )      ).to(A      )

            _UpperCamelCase              =						pipe_a(
                prompt_embeds=A,       negative_prompt_embeds=A,       image=A,       generator=A,       num_inference_steps=2,       output_type='''np''',       )

            _UpperCamelCase              =						output.images[0]

            assert image.shape == (256, 256, 3)

            _UpperCamelCase              =						torch.cuda.max_memory_allocated()
            assert mem_bytes < 4 * 10**9

            _UpperCamelCase              =						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 	_UpperCamelCase     (      self,       A,       A,       A,       A      ):





            """simple docstring"""

            _start_torch_memory_measurement()

            _UpperCamelCase              =						floats_tensor((1, 3, 64, 64),       rng=random.Random(0      )      ).to(A      )

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

            _UpperCamelCase              =						pipe_a(
                prompt_embeds=A,       negative_prompt_embeds=A,       image=A,       num_inference_steps=2,       generator=A,       output_type='''np''',       )

            _UpperCamelCase              =						output.images[0]

            assert image.shape == (64, 64, 3)

            _UpperCamelCase              =						torch.cuda.max_memory_allocated()
            assert mem_bytes < 10 * 10**9

            _UpperCamelCase              =						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()

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

            _UpperCamelCase              =						floats_tensor((1, 3, 256, 256),       rng=random.Random(0      )      ).to(A      )
            _UpperCamelCase              =						floats_tensor((1, 3, 64, 64),       rng=random.Random(0      )      ).to(A      )

            _UpperCamelCase              =						pipe_a(
                prompt_embeds=A,       negative_prompt_embeds=A,       image=A,       original_image=A,       generator=A,       num_inference_steps=2,       output_type='''np''',       )

            _UpperCamelCase              =						output.images[0]

            assert image.shape == (256, 256, 3)

            _UpperCamelCase              =						torch.cuda.max_memory_allocated()
            assert mem_bytes < 4 * 10**9

            _UpperCamelCase              =						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 	_UpperCamelCase     (      self,       A,       A,       A,       A      ):





            """simple docstring"""

            _start_torch_memory_measurement()

            _UpperCamelCase              =						floats_tensor((1, 3, 64, 64),       rng=random.Random(0      )      ).to(A      )
            _UpperCamelCase              =						floats_tensor((1, 3, 64, 64),       rng=random.Random(1      )      ).to(A      )

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

            _UpperCamelCase              =						output.images[0]

            assert image.shape == (64, 64, 3)

            _UpperCamelCase              =						torch.cuda.max_memory_allocated()
            assert mem_bytes < 10 * 10**9

            _UpperCamelCase              =						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()

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

            _UpperCamelCase              =						floats_tensor((1, 3, 64, 64),       rng=random.Random(0      )      ).to(A      )
            _UpperCamelCase              =						floats_tensor((1, 3, 256, 256),       rng=random.Random(0      )      ).to(A      )
            _UpperCamelCase              =						floats_tensor((1, 3, 256, 256),       rng=random.Random(1      )      ).to(A      )

            _UpperCamelCase              =						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''',       )

            _UpperCamelCase              =						output.images[0]

            assert image.shape == (256, 256, 3)

            _UpperCamelCase              =						torch.cuda.max_memory_allocated()
            assert mem_bytes < 4 * 10**9

            _UpperCamelCase              =						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     SCREAMING_SNAKE_CASE			(      ):
       torch.cuda.empty_cache()
       torch.cuda.reset_max_memory_allocated()
       torch.cuda.reset_peak_memory_stats()













import copy
from collections import OrderedDict
from typing import Dict, Mapping

from packaging import version

from ...configuration_utils import PretrainedConfig
from ...onnx import OnnxConfig
from ...utils import logging
from ..auto import CONFIG_MAPPING


lowercase    : Optional[int]    	=  logging.get_logger(__name__)

lowercase    : List[Any]    	=  {
    """facebook/detr-resnet-50""": """https://huggingface.co/facebook/detr-resnet-50/resolve/main/config.json""",
    # See all DETR models at https://huggingface.co/models?filter=detr
}

class  __A(						__UpperCAmelCase ):
     __A	  =       "detr"
     __A	  =       ["past_key_values"]
     __A	  =       {
         "hidden_size": "d_model",
         "num_attention_heads": "encoder_attention_heads",
     }

     def __init__(      self,       A=True,       A=None,       A=3,       A=100,       A=6,       A=2048,       A=8,       A=6,       A=2048,       A=8,       A=0.0,       A=0.0,       A=True,       A="relu",       A=256,       A=0.1,       A=0.0,       A=0.0,       A=0.02,       A=1.0,       A=False,       A="sine",       A="resnet50",       A=True,       A=False,       A=1,       A=5,       A=2,       A=1,       A=1,       A=5,       A=2,       A=0.1,       **A,       ):





            """simple docstring"""

            if backbone_config is not None and use_timm_backbone:
                   raise ValueError('''You can\'t specify both `backbone_config` and `use_timm_backbone`.'''      )

            if not use_timm_backbone:
                   if backbone_config is None:
                          logger.info('''`backbone_config` is `None`. Initializing the config with the default `ResNet` backbone.'''      )
                          _UpperCamelCase              =						CONFIG_MAPPING['''resnet'''](out_features=['''stage4''']      )
                   elif isinstance(A,       A      ):
                          _UpperCamelCase              =						backbone_config.get('''model_type'''      )
                          _UpperCamelCase              =						CONFIG_MAPPING[backbone_model_type]
                          _UpperCamelCase              =						config_class.from_dict(A      )
                   # set timm attributes to None
                   _UpperCamelCase       ,		_UpperCamelCase       ,		_UpperCamelCase              =						None, None, None

            _UpperCamelCase              =						use_timm_backbone
            _UpperCamelCase              =						backbone_config
            _UpperCamelCase              =						num_channels
            _UpperCamelCase              =						num_queries
            _UpperCamelCase              =						d_model
            _UpperCamelCase              =						encoder_ffn_dim
            _UpperCamelCase              =						encoder_layers
            _UpperCamelCase              =						encoder_attention_heads
            _UpperCamelCase              =						decoder_ffn_dim
            _UpperCamelCase              =						decoder_layers
            _UpperCamelCase              =						decoder_attention_heads
            _UpperCamelCase              =						dropout
            _UpperCamelCase              =						attention_dropout
            _UpperCamelCase              =						activation_dropout
            _UpperCamelCase              =						activation_function
            _UpperCamelCase              =						init_std
            _UpperCamelCase              =						init_xavier_std
            _UpperCamelCase              =						encoder_layerdrop
            _UpperCamelCase              =						decoder_layerdrop
            _UpperCamelCase              =						encoder_layers
            _UpperCamelCase              =						auxiliary_loss
            _UpperCamelCase              =						position_embedding_type
            _UpperCamelCase              =						backbone
            _UpperCamelCase              =						use_pretrained_backbone
            _UpperCamelCase              =						dilation
            # Hungarian matcher
            _UpperCamelCase              =						class_cost
            _UpperCamelCase              =						bbox_cost
            _UpperCamelCase              =						giou_cost
            # Loss coefficients
            _UpperCamelCase              =						mask_loss_coefficient
            _UpperCamelCase              =						dice_loss_coefficient
            _UpperCamelCase              =						bbox_loss_coefficient
            _UpperCamelCase              =						giou_loss_coefficient
            _UpperCamelCase              =						eos_coefficient
            super().__init__(is_encoder_decoder=A,       **A      )

     @property
     def 	_UpperCamelCase     (      self      ):





            """simple docstring"""

            return self.encoder_attention_heads

     @property
     def 	_UpperCamelCase     (      self      ):





            """simple docstring"""

            return self.d_model

     @classmethod
     def 	_UpperCamelCase     (      cls,       A,       **A      ):





            """simple docstring"""

            return cls(backbone_config=A,       **A      )



     def 	_UpperCamelCase     (      self      ):





            """simple docstring"""

            _UpperCamelCase              =						copy.deepcopy(self.__dict__      )
            if output["backbone_config"] is not None:
                   _UpperCamelCase              =						self.backbone_config.to_dict()
            _UpperCamelCase              =						self.__class__.model_type
            return output






class  __A(						__UpperCAmelCase ):
     __A	  =       version.parse("1.11" )

     @property
     def 	_UpperCamelCase     (      self      ):





            """simple docstring"""

            return OrderedDict(
                [
                    ('''pixel_values''', {0: '''batch''', 1: '''num_channels''', 2: '''height''', 3: '''width'''}),
                    ('''pixel_mask''', {0: '''batch'''}),
                ]      )

     @property
     def 	_UpperCamelCase     (      self      ):





            """simple docstring"""

            return 1E-5



     @property
     def 	_UpperCamelCase     (      self      ):





            """simple docstring"""

            return 12






