
import argparse
import json
from collections import OrderedDict
from pathlib import Path

import requests
import torch
from huggingface_hub import hf_hub_download
from PIL import Image

from transformers import DetrConfig, DetrForObjectDetection, DetrForSegmentation, DetrImageProcessor
from transformers.utils import logging


logging.set_verbosity_info()
a__				:      Optional[Any]    		=     logging.get_logger(__name__)

# here we list all keys to be renamed (original name on the left, our name on the right)
a__				:      List[str]    		=     []
for i in range(6):
      # encoder layers: output projection, 2 feedforward neural networks and 2 layernorms
      rename_keys.append(
          (F"transformer.encoder.layers.{i}.self_attn.out_proj.weight", F"encoder.layers.{i}.self_attn.out_proj.weight")
      )
      rename_keys.append(
          (F"transformer.encoder.layers.{i}.self_attn.out_proj.bias", F"encoder.layers.{i}.self_attn.out_proj.bias")
      )
      rename_keys.append((F"transformer.encoder.layers.{i}.linear1.weight", F"encoder.layers.{i}.fc1.weight"))
      rename_keys.append((F"transformer.encoder.layers.{i}.linear1.bias", F"encoder.layers.{i}.fc1.bias"))
      rename_keys.append((F"transformer.encoder.layers.{i}.linear2.weight", F"encoder.layers.{i}.fc2.weight"))
      rename_keys.append((F"transformer.encoder.layers.{i}.linear2.bias", F"encoder.layers.{i}.fc2.bias"))
      rename_keys.append(
          (F"transformer.encoder.layers.{i}.norm1.weight", F"encoder.layers.{i}.self_attn_layer_norm.weight")
      )
      rename_keys.append((F"transformer.encoder.layers.{i}.norm1.bias", F"encoder.layers.{i}.self_attn_layer_norm.bias"))
      rename_keys.append((F"transformer.encoder.layers.{i}.norm2.weight", F"encoder.layers.{i}.final_layer_norm.weight"))
      rename_keys.append((F"transformer.encoder.layers.{i}.norm2.bias", F"encoder.layers.{i}.final_layer_norm.bias"))
      # decoder layers: 2 times output projection, 2 feedforward neural networks and 3 layernorms
      rename_keys.append(
          (F"transformer.decoder.layers.{i}.self_attn.out_proj.weight", F"decoder.layers.{i}.self_attn.out_proj.weight")
      )
      rename_keys.append(
          (F"transformer.decoder.layers.{i}.self_attn.out_proj.bias", F"decoder.layers.{i}.self_attn.out_proj.bias")
      )
      rename_keys.append(
          (
              F"transformer.decoder.layers.{i}.multihead_attn.out_proj.weight",
              F"decoder.layers.{i}.encoder_attn.out_proj.weight",
          )
      )
      rename_keys.append(
          (
              F"transformer.decoder.layers.{i}.multihead_attn.out_proj.bias",
              F"decoder.layers.{i}.encoder_attn.out_proj.bias",
          )
      )
      rename_keys.append((F"transformer.decoder.layers.{i}.linear1.weight", F"decoder.layers.{i}.fc1.weight"))
      rename_keys.append((F"transformer.decoder.layers.{i}.linear1.bias", F"decoder.layers.{i}.fc1.bias"))
      rename_keys.append((F"transformer.decoder.layers.{i}.linear2.weight", F"decoder.layers.{i}.fc2.weight"))
      rename_keys.append((F"transformer.decoder.layers.{i}.linear2.bias", F"decoder.layers.{i}.fc2.bias"))
      rename_keys.append(
          (F"transformer.decoder.layers.{i}.norm1.weight", F"decoder.layers.{i}.self_attn_layer_norm.weight")
      )
      rename_keys.append((F"transformer.decoder.layers.{i}.norm1.bias", F"decoder.layers.{i}.self_attn_layer_norm.bias"))
      rename_keys.append(
          (F"transformer.decoder.layers.{i}.norm2.weight", F"decoder.layers.{i}.encoder_attn_layer_norm.weight")
      )
      rename_keys.append(
          (F"transformer.decoder.layers.{i}.norm2.bias", F"decoder.layers.{i}.encoder_attn_layer_norm.bias")
      )
      rename_keys.append((F"transformer.decoder.layers.{i}.norm3.weight", F"decoder.layers.{i}.final_layer_norm.weight"))
      rename_keys.append((F"transformer.decoder.layers.{i}.norm3.bias", F"decoder.layers.{i}.final_layer_norm.bias"))

# convolutional projection + query embeddings + layernorm of decoder + class and bounding box heads
rename_keys.extend(
    [
        ('''input_proj.weight''', '''input_projection.weight'''),
        ('''input_proj.bias''', '''input_projection.bias'''),
        ('''query_embed.weight''', '''query_position_embeddings.weight'''),
        ('''transformer.decoder.norm.weight''', '''decoder.layernorm.weight'''),
        ('''transformer.decoder.norm.bias''', '''decoder.layernorm.bias'''),
        ('''class_embed.weight''', '''class_labels_classifier.weight'''),
        ('''class_embed.bias''', '''class_labels_classifier.bias'''),
        ('''bbox_embed.layers.0.weight''', '''bbox_predictor.layers.0.weight'''),
        ('''bbox_embed.layers.0.bias''', '''bbox_predictor.layers.0.bias'''),
        ('''bbox_embed.layers.1.weight''', '''bbox_predictor.layers.1.weight'''),
        ('''bbox_embed.layers.1.bias''', '''bbox_predictor.layers.1.bias'''),
        ('''bbox_embed.layers.2.weight''', '''bbox_predictor.layers.2.weight'''),
        ('''bbox_embed.layers.2.bias''', '''bbox_predictor.layers.2.bias'''),
    ]
)







def 				UpperCAmelCase_(    a__ ,						a__ ,						a__   ):



    """simple docstring"""



    SCREAMING_SNAKE_CASE				:				Any				      = state_dict.pop(a__   )
    SCREAMING_SNAKE_CASE				:				List[str]				      = val







def 				UpperCAmelCase_(    a__   ):



    """simple docstring"""



    SCREAMING_SNAKE_CASE				:				int				      = OrderedDict()
    for key, value in state_dict.items():
        if "backbone.0.body" in key:
            SCREAMING_SNAKE_CASE				:				Tuple				      = key.replace('''backbone.0.body''' ,						'''backbone.conv_encoder.model'''   )
            SCREAMING_SNAKE_CASE				:				int				      = value
        else:
            SCREAMING_SNAKE_CASE				:				int				      = value

    return new_state_dict







def 				UpperCAmelCase_(    a__ ,						a__=False   ):



    """simple docstring"""



    SCREAMING_SNAKE_CASE				:				Any				      = ''''''
    if is_panoptic:
        SCREAMING_SNAKE_CASE				:				Optional[Any]				      = '''detr.'''

    # first: transformer encoder
    for i in range(6   ):
        # read in weights + bias of input projection layer (in PyTorch's MultiHeadAttention, this is a single matrix + bias)
        SCREAMING_SNAKE_CASE				:				str				      = state_dict.pop(F"""{prefix}transformer.encoder.layers.{i}.self_attn.in_proj_weight"""   )
        SCREAMING_SNAKE_CASE				:				Dict				      = state_dict.pop(F"""{prefix}transformer.encoder.layers.{i}.self_attn.in_proj_bias"""   )
        # next, add query, keys and values (in that order) to the state dict
        SCREAMING_SNAKE_CASE				:				Any				      = in_proj_weight[:256, :]
        SCREAMING_SNAKE_CASE				:				Dict				      = in_proj_bias[:256]
        SCREAMING_SNAKE_CASE				:				Tuple				      = in_proj_weight[256:512, :]
        SCREAMING_SNAKE_CASE				:				Any				      = in_proj_bias[256:512]
        SCREAMING_SNAKE_CASE				:				Dict				      = in_proj_weight[-256:, :]
        SCREAMING_SNAKE_CASE				:				Any				      = in_proj_bias[-256:]
    # next: transformer decoder (which is a bit more complex because it also includes cross-attention)
    for i in range(6   ):
        # read in weights + bias of input projection layer of self-attention
        SCREAMING_SNAKE_CASE				:				int				      = state_dict.pop(F"""{prefix}transformer.decoder.layers.{i}.self_attn.in_proj_weight"""   )
        SCREAMING_SNAKE_CASE				:				int				      = state_dict.pop(F"""{prefix}transformer.decoder.layers.{i}.self_attn.in_proj_bias"""   )
        # next, add query, keys and values (in that order) to the state dict
        SCREAMING_SNAKE_CASE				:				Any				      = in_proj_weight[:256, :]
        SCREAMING_SNAKE_CASE				:				Optional[Any]				      = in_proj_bias[:256]
        SCREAMING_SNAKE_CASE				:				Optional[Any]				      = in_proj_weight[256:512, :]
        SCREAMING_SNAKE_CASE				:				Optional[int]				      = in_proj_bias[256:512]
        SCREAMING_SNAKE_CASE				:				Optional[int]				      = in_proj_weight[-256:, :]
        SCREAMING_SNAKE_CASE				:				Union[str, Any]				      = in_proj_bias[-256:]
        # read in weights + bias of input projection layer of cross-attention
        SCREAMING_SNAKE_CASE				:				Dict				      = state_dict.pop(
            F"""{prefix}transformer.decoder.layers.{i}.multihead_attn.in_proj_weight"""   )
        SCREAMING_SNAKE_CASE				:				Optional[Any]				      = state_dict.pop(F"""{prefix}transformer.decoder.layers.{i}.multihead_attn.in_proj_bias"""   )
        # next, add query, keys and values (in that order) of cross-attention to the state dict
        SCREAMING_SNAKE_CASE				:				Optional[Any]				      = in_proj_weight_cross_attn[:256, :]
        SCREAMING_SNAKE_CASE				:				List[Any]				      = in_proj_bias_cross_attn[:256]
        SCREAMING_SNAKE_CASE				:				List[str]				      = in_proj_weight_cross_attn[256:512, :]
        SCREAMING_SNAKE_CASE				:				List[str]				      = in_proj_bias_cross_attn[256:512]
        SCREAMING_SNAKE_CASE				:				Any				      = in_proj_weight_cross_attn[-256:, :]
        SCREAMING_SNAKE_CASE				:				Tuple				      = in_proj_bias_cross_attn[-256:]







def 				UpperCAmelCase_(    ):



    """simple docstring"""



    SCREAMING_SNAKE_CASE				:				Optional[Any]				      = '''http://images.cocodataset.org/val2017/000000039769.jpg'''
    SCREAMING_SNAKE_CASE				:				Union[str, Any]				      = Image.open(requests.get(a__ ,						stream=a__   ).raw   )

    return im







@torch.no_grad()
def 				UpperCAmelCase_(    a__ ,						a__   ):



    """simple docstring"""



    SCREAMING_SNAKE_CASE				:				Dict				      = DetrConfig()
    # set backbone and dilation attributes
    if "resnet101" in model_name:
        SCREAMING_SNAKE_CASE				:				Any				      = '''resnet101'''
    if "dc5" in model_name:
        SCREAMING_SNAKE_CASE				:				Optional[Any]				      = True
    SCREAMING_SNAKE_CASE				:				str				      = '''panoptic''' in model_name
    if is_panoptic:
        SCREAMING_SNAKE_CASE				:				List[str]				      = 250
    else:
        SCREAMING_SNAKE_CASE				:				Dict				      = 91
        SCREAMING_SNAKE_CASE				:				Any				      = '''huggingface/label-files'''
        SCREAMING_SNAKE_CASE				:				Optional[Any]				      = '''coco-detection-id2label.json'''
        SCREAMING_SNAKE_CASE				:				int				      = json.load(open(hf_hub_download(a__ ,						a__ ,						repo_type='''dataset'''   ) ,						'''r'''   )   )
        SCREAMING_SNAKE_CASE				:				Union[str, Any]				      = {int(a__   ): v for k, v in idalabel.items()}
        SCREAMING_SNAKE_CASE				:				Optional[int]				      = idalabel
        SCREAMING_SNAKE_CASE				:				Any				      = {v: k for k, v in idalabel.items()}

    # load image processor
    SCREAMING_SNAKE_CASE				:				int				      = '''coco_panoptic''' if is_panoptic else '''coco_detection'''
    SCREAMING_SNAKE_CASE				:				Tuple				      = DetrImageProcessor(format=a__   )

    # prepare image
    SCREAMING_SNAKE_CASE				:				Optional[int]				      = prepare_img()
    SCREAMING_SNAKE_CASE				:				List[Any]				      = image_processor(images=a__ ,						return_tensors='''pt'''   )
    SCREAMING_SNAKE_CASE				:				Dict				      = encoding['''pixel_values''']

    logger.info(F"""Converting model {model_name}..."""   )

    # load original model from torch hub
    SCREAMING_SNAKE_CASE				:				Any				      = torch.hub.load('''facebookresearch/detr''' ,						a__ ,						pretrained=a__   ).eval()
    SCREAMING_SNAKE_CASE				:				str				      = detr.state_dict()
    # rename keys
    for src, dest in rename_keys:
        if is_panoptic:
            SCREAMING_SNAKE_CASE				:				Optional[int]				      = '''detr.''' + src
        rename_key(a__ ,						a__ ,						a__   )
    SCREAMING_SNAKE_CASE				:				Any				      = rename_backbone_keys(a__   )
    # query, key and value matrices need special treatment
    read_in_q_k_v(a__ ,						is_panoptic=a__   )
    # important: we need to prepend a prefix to each of the base model keys as the head models use different attributes for them
    SCREAMING_SNAKE_CASE				:				Tuple				      = '''detr.model.''' if is_panoptic else '''model.'''
    for key in state_dict.copy().keys():
        if is_panoptic:
            if (
                key.startswith('''detr'''   )
                and not key.startswith('''class_labels_classifier'''   )
                and not key.startswith('''bbox_predictor'''   )
            ):
                SCREAMING_SNAKE_CASE				:				List[str]				      = state_dict.pop(a__   )
                SCREAMING_SNAKE_CASE				:				int				      = val
            elif "class_labels_classifier" in key or "bbox_predictor" in key:
                SCREAMING_SNAKE_CASE				:				int				      = state_dict.pop(a__   )
                SCREAMING_SNAKE_CASE				:				List[str]				      = val
            elif key.startswith('''bbox_attention'''   ) or key.startswith('''mask_head'''   ):
                continue
            else:
                SCREAMING_SNAKE_CASE				:				Optional[int]				      = state_dict.pop(a__   )
                SCREAMING_SNAKE_CASE				:				Dict				      = val
        else:
            if not key.startswith('''class_labels_classifier'''   ) and not key.startswith('''bbox_predictor'''   ):
                SCREAMING_SNAKE_CASE				:				Optional[Any]				      = state_dict.pop(a__   )
                SCREAMING_SNAKE_CASE				:				Optional[Any]				      = val
    # finally, create HuggingFace model and load state dict
    SCREAMING_SNAKE_CASE				:				Optional[int]				      = DetrForSegmentation(a__   ) if is_panoptic else DetrForObjectDetection(a__   )
    model.load_state_dict(a__   )
    model.eval()
    # verify our conversion
    SCREAMING_SNAKE_CASE				:				str				      = detr(a__   )
    SCREAMING_SNAKE_CASE				:				Tuple				      = model(a__   )
    assert torch.allclose(outputs.logits ,						original_outputs['''pred_logits'''] ,						atol=1e-4   )
    assert torch.allclose(outputs.pred_boxes ,						original_outputs['''pred_boxes'''] ,						atol=1e-4   )
    if is_panoptic:
        assert torch.allclose(outputs.pred_masks ,						original_outputs['''pred_masks'''] ,						atol=1e-4   )

    # Save model and image processor
    logger.info(F"""Saving PyTorch model and image processor to {pytorch_dump_folder_path}..."""   )
    Path(a__   ).mkdir(exist_ok=a__   )
    model.save_pretrained(a__   )
    image_processor.save_pretrained(a__   )


if __name__ == "__main__":
      a__				:      int    		=     argparse.ArgumentParser()

      parser.add_argument(
          '''--model_name''', default='''detr_resnet50''', type=str, help='''Name of the DETR model you\'d like to convert.'''
      )
      parser.add_argument(
          '''--pytorch_dump_folder_path''', default=None, type=str, help='''Path to the folder to output PyTorch model.'''
      )
      a__				:      int    		=     parser.parse_args()
      convert_detr_checkpoint(args.model_name, args.pytorch_dump_folder_path)


import gc
import unittest

import numpy as np
import torch
from transformers import CLIPTextConfig, CLIPTextModel, CLIPTokenizer

from diffusers import (
    AutoencoderKL,
    DDIMScheduler,
    EulerAncestralDiscreteScheduler,
    PNDMScheduler,
    StableDiffusionModelEditingPipeline,
    UNetaDConditionModel,
)
from diffusers.utils import slow, torch_device
from diffusers.utils.testing_utils import enable_full_determinism, require_torch_gpu, skip_mps

from ..pipeline_params import TEXT_TO_IMAGE_BATCH_PARAMS, TEXT_TO_IMAGE_IMAGE_PARAMS, TEXT_TO_IMAGE_PARAMS
from ..test_pipelines_common import PipelineKarrasSchedulerTesterMixin, PipelineLatentTesterMixin, PipelineTesterMixin


enable_full_determinism()






@skip_mps
class     a_       (			a__				,    a__				,    a__				,    unittest.TestCase							):



       """simple docstring"""
       __SCREAMING_SNAKE_CASE       :					List[Any]						=  StableDiffusionModelEditingPipeline
       __SCREAMING_SNAKE_CASE       :					Optional[int]						=  TEXT_TO_IMAGE_PARAMS
       __SCREAMING_SNAKE_CASE       :					int						=  TEXT_TO_IMAGE_BATCH_PARAMS
       __SCREAMING_SNAKE_CASE       :					Any						=  TEXT_TO_IMAGE_IMAGE_PARAMS
       __SCREAMING_SNAKE_CASE       :					List[Any]						=  TEXT_TO_IMAGE_IMAGE_PARAMS






       def 	__lowerCAmelCase     (		self   )					->int:
           torch.manual_seed(0   )
           SCREAMING_SNAKE_CASE				:				Dict				      = UNetaDConditionModel(
               block_out_channels=(32, 64)				,       layers_per_block=2				,       sample_size=32				,       in_channels=4				,       out_channels=4				,       down_block_types=('''DownBlock2D''', '''CrossAttnDownBlock2D''')				,       up_block_types=('''CrossAttnUpBlock2D''', '''UpBlock2D''')				,       cross_attention_dim=32				,       )
           SCREAMING_SNAKE_CASE				:				Dict				      = DDIMScheduler()
           torch.manual_seed(0   )
           SCREAMING_SNAKE_CASE				:				Union[str, Any]				      = AutoencoderKL(
               block_out_channels=[32, 64]				,       in_channels=3				,       out_channels=3				,       down_block_types=['''DownEncoderBlock2D''', '''DownEncoderBlock2D''']				,       up_block_types=['''UpDecoderBlock2D''', '''UpDecoderBlock2D''']				,       latent_channels=4				,       )
           torch.manual_seed(0   )
           SCREAMING_SNAKE_CASE				:				List[str]				      = CLIPTextConfig(
               bos_token_id=0				,       eos_token_id=2				,       hidden_size=32				,       intermediate_size=37				,       layer_norm_eps=1e-05				,       num_attention_heads=4				,       num_hidden_layers=5				,       pad_token_id=1				,       vocab_size=1000				,       )
           SCREAMING_SNAKE_CASE				:				Dict				      = CLIPTextModel(_lowerCamelCase   )
           SCREAMING_SNAKE_CASE				:				Tuple				      = CLIPTokenizer.from_pretrained('''hf-internal-testing/tiny-random-clip'''   )

           SCREAMING_SNAKE_CASE				:				List[Any]				      = {
               '''unet''': unet,
               '''scheduler''': scheduler,
               '''vae''': vae,
               '''text_encoder''': text_encoder,
               '''tokenizer''': tokenizer,
               '''safety_checker''': None,
               '''feature_extractor''': None,
           }
           return components






       def 	__lowerCAmelCase     (		self				,       _lowerCamelCase				,       _lowerCamelCase=0   )					->Optional[int]:
           SCREAMING_SNAKE_CASE				:				Any				      = torch.manual_seed(_lowerCamelCase   )
           SCREAMING_SNAKE_CASE				:				Union[str, Any]				      = {
               '''prompt''': '''A field of roses''',
               '''generator''': generator,
               # Setting height and width to None to prevent OOMs on CPU.
               '''height''': None,
               '''width''': None,
               '''num_inference_steps''': 2,
               '''guidance_scale''': 6.0,
               '''output_type''': '''numpy''',
           }
           return inputs






       def 	__lowerCAmelCase     (		self   )					->Optional[int]:
           SCREAMING_SNAKE_CASE				:				Tuple				      = '''cpu'''  # ensure determinism for the device-dependent torch.Generator
           SCREAMING_SNAKE_CASE				:				Any				      = self.get_dummy_components()
           SCREAMING_SNAKE_CASE				:				Optional[Any]				      = StableDiffusionModelEditingPipeline(**_lowerCamelCase   )
           SCREAMING_SNAKE_CASE				:				List[Any]				      = sd_pipe.to(_lowerCamelCase   )
           sd_pipe.set_progress_bar_config(disable=_lowerCamelCase   )

           SCREAMING_SNAKE_CASE				:				Tuple				      = self.get_dummy_inputs(_lowerCamelCase   )
           SCREAMING_SNAKE_CASE				:				Tuple				      = sd_pipe(**_lowerCamelCase   ).images
           SCREAMING_SNAKE_CASE				:				Any				      = image[0, -3:, -3:, -1]
           assert image.shape == (1, 64, 64, 3)

           SCREAMING_SNAKE_CASE				:				List[str]				      = np.array([0.4_7_5_5, 0.5_1_3_2, 0.4_9_7_6, 0.3_9_0_4, 0.3_5_5_4, 0.4_7_6_5, 0.5_1_3_9, 0.5_1_5_8, 0.4_8_8_9]   )

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






       def 	__lowerCAmelCase     (		self   )					->List[Any]:
           SCREAMING_SNAKE_CASE				:				Dict				      = '''cpu'''  # ensure determinism for the device-dependent torch.Generator
           SCREAMING_SNAKE_CASE				:				Tuple				      = self.get_dummy_components()
           SCREAMING_SNAKE_CASE				:				Optional[int]				      = StableDiffusionModelEditingPipeline(**_lowerCamelCase   )
           SCREAMING_SNAKE_CASE				:				str				      = sd_pipe.to(_lowerCamelCase   )
           sd_pipe.set_progress_bar_config(disable=_lowerCamelCase   )

           SCREAMING_SNAKE_CASE				:				Any				      = self.get_dummy_inputs(_lowerCamelCase   )
           SCREAMING_SNAKE_CASE				:				List[str]				      = '''french fries'''
           SCREAMING_SNAKE_CASE				:				Optional[Any]				      = sd_pipe(**_lowerCamelCase				,       negative_prompt=_lowerCamelCase   )
           SCREAMING_SNAKE_CASE				:				Any				      = output.images
           SCREAMING_SNAKE_CASE				:				Union[str, Any]				      = image[0, -3:, -3:, -1]

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

           SCREAMING_SNAKE_CASE				:				Union[str, Any]				      = np.array([0.4_9_9_2, 0.5_1_0_1, 0.5_0_0_4, 0.3_9_4_9, 0.3_6_0_4, 0.4_7_3_5, 0.5_2_1_6, 0.5_2_0_4, 0.4_9_1_3]   )

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






       def 	__lowerCAmelCase     (		self   )					->Union[str, Any]:
           SCREAMING_SNAKE_CASE				:				List[str]				      = '''cpu'''  # ensure determinism for the device-dependent torch.Generator
           SCREAMING_SNAKE_CASE				:				Union[str, Any]				      = self.get_dummy_components()
           SCREAMING_SNAKE_CASE				:				List[Any]				      = EulerAncestralDiscreteScheduler(
               beta_start=0.0_0_0_8_5				,       beta_end=0.0_1_2				,       beta_schedule='''scaled_linear'''   )
           SCREAMING_SNAKE_CASE				:				List[Any]				      = StableDiffusionModelEditingPipeline(**_lowerCamelCase   )
           SCREAMING_SNAKE_CASE				:				Optional[Any]				      = sd_pipe.to(_lowerCamelCase   )
           sd_pipe.set_progress_bar_config(disable=_lowerCamelCase   )

           SCREAMING_SNAKE_CASE				:				Union[str, Any]				      = self.get_dummy_inputs(_lowerCamelCase   )
           SCREAMING_SNAKE_CASE				:				Union[str, Any]				      = sd_pipe(**_lowerCamelCase   ).images
           SCREAMING_SNAKE_CASE				:				Optional[int]				      = image[0, -3:, -3:, -1]

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

           SCREAMING_SNAKE_CASE				:				str				      = np.array([0.4_7_4_7, 0.5_3_7_2, 0.4_7_7_9, 0.4_9_8_2, 0.5_5_4_3, 0.4_8_1_6, 0.5_2_3_8, 0.4_9_0_4, 0.5_0_2_7]   )

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






       def 	__lowerCAmelCase     (		self   )					->Any:
           SCREAMING_SNAKE_CASE				:				Optional[int]				      = '''cpu'''  # ensure determinism for the device-dependent torch.Generator
           SCREAMING_SNAKE_CASE				:				List[Any]				      = self.get_dummy_components()
           SCREAMING_SNAKE_CASE				:				Optional[int]				      = PNDMScheduler()
           SCREAMING_SNAKE_CASE				:				Tuple				      = StableDiffusionModelEditingPipeline(**_lowerCamelCase   )
           SCREAMING_SNAKE_CASE				:				Optional[Any]				      = sd_pipe.to(_lowerCamelCase   )
           sd_pipe.set_progress_bar_config(disable=_lowerCamelCase   )

           SCREAMING_SNAKE_CASE				:				Optional[Any]				      = self.get_dummy_inputs(_lowerCamelCase   )
           # the pipeline does not expect pndm so test if it raises error.
           with self.assertRaises(_lowerCamelCase   ):
               SCREAMING_SNAKE_CASE				:				Dict				      = sd_pipe(**_lowerCamelCase   ).images






       def 	__lowerCAmelCase     (		self   )					->Dict:
           super().test_inference_batch_single_identical(expected_max_diff=5e-3   )






       def 	__lowerCAmelCase     (		self   )					->List[Any]:
           super().test_attention_slicing_forward_pass(expected_max_diff=5e-3   )






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



       """simple docstring"""






       def 	__lowerCAmelCase     (		self   )					->Dict:
           super().tearDown()
           gc.collect()
           torch.cuda.empty_cache()






       def 	__lowerCAmelCase     (		self				,       _lowerCamelCase=0   )					->Dict:
           SCREAMING_SNAKE_CASE				:				str				      = torch.manual_seed(_lowerCamelCase   )
           SCREAMING_SNAKE_CASE				:				List[Any]				      = {
               '''prompt''': '''A field of roses''',
               '''generator''': generator,
               '''num_inference_steps''': 3,
               '''guidance_scale''': 7.5,
               '''output_type''': '''numpy''',
           }
           return inputs






       def 	__lowerCAmelCase     (		self   )					->List[str]:
           SCREAMING_SNAKE_CASE				:				Union[str, Any]				      = '''CompVis/stable-diffusion-v1-4'''
           SCREAMING_SNAKE_CASE				:				List[str]				      = StableDiffusionModelEditingPipeline.from_pretrained(_lowerCamelCase				,       safety_checker=_lowerCamelCase   )
           pipe.to(_lowerCamelCase   )
           pipe.set_progress_bar_config(disable=_lowerCamelCase   )
           pipe.enable_attention_slicing()

           SCREAMING_SNAKE_CASE				:				Any				      = self.get_inputs()
           SCREAMING_SNAKE_CASE				:				Union[str, Any]				      = pipe(**_lowerCamelCase   ).images
           SCREAMING_SNAKE_CASE				:				Optional[Any]				      = image[0, -3:, -3:, -1].flatten()

           assert image.shape == (1, 512, 512, 3)

           SCREAMING_SNAKE_CASE				:				str				      = np.array(
               [0.6_7_4_9_4_9_6, 0.6_3_8_6_4_5_3, 0.5_1_4_4_3_2_6_7, 0.6_6_0_9_4_9_0_5, 0.6_1_9_2_1_2_1_5, 0.5_4_9_1_3_3_2, 0.5_7_4_4_4_1_7, 0.5_8_0_7_5_1_0_6, 0.5_1_7_4_6_5_8]   )

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

           # make sure image changes after editing
           pipe.edit_model('''A pack of roses'''				,       '''A pack of blue roses'''   )

           SCREAMING_SNAKE_CASE				:				Any				      = pipe(**_lowerCamelCase   ).images
           SCREAMING_SNAKE_CASE				:				List[str]				      = image[0, -3:, -3:, -1].flatten()

           assert image.shape == (1, 512, 512, 3)

           assert np.abs(expected_slice - image_slice   ).max() > 1e-1






       def 	__lowerCAmelCase     (		self   )					->str:
           torch.cuda.empty_cache()
           torch.cuda.reset_max_memory_allocated()
           torch.cuda.reset_peak_memory_stats()

           SCREAMING_SNAKE_CASE				:				Optional[int]				      = '''CompVis/stable-diffusion-v1-4'''
           SCREAMING_SNAKE_CASE				:				Any				      = DDIMScheduler.from_pretrained(_lowerCamelCase				,       subfolder='''scheduler'''   )
           SCREAMING_SNAKE_CASE				:				Optional[int]				      = StableDiffusionModelEditingPipeline.from_pretrained(
               _lowerCamelCase				,       scheduler=_lowerCamelCase				,       safety_checker=_lowerCamelCase   )
           SCREAMING_SNAKE_CASE				:				Dict				      = pipe.to(_lowerCamelCase   )
           pipe.set_progress_bar_config(disable=_lowerCamelCase   )
           pipe.enable_attention_slicing(1   )
           pipe.enable_sequential_cpu_offload()

           SCREAMING_SNAKE_CASE				:				List[Any]				      = self.get_inputs()
           SCREAMING_SNAKE_CASE				:				Union[str, Any]				      = pipe(**_lowerCamelCase   )

           SCREAMING_SNAKE_CASE				:				Optional[Any]				      = torch.cuda.max_memory_allocated()
           # make sure that less than 4.4 GB is allocated
           assert mem_bytes < 4.4 * 10**9
