



import argparse
import json

import torch
from huggingface_hub import hf_hub_download
from PIL import Image
from timm.models import create_model

from transformers import (
    BeitImageProcessor,
    DataaVecVisionConfig,
    DataaVecVisionForImageClassification,
    DataaVecVisionModel,
)







def    lowerCAmelCase_  (lowerCAmelCase__:		Optional[Any]		,	lowerCAmelCase__:		List[str]=False		,	lowerCAmelCase__:		List[str]=False		,	lowerCAmelCase__:		Union[str, Any]="data2vec."     ):





   """simple docstring"""


   UpperCAmelCase_:		str       			=    """backbone.""" if is_semantic else """"""

   UpperCAmelCase_:		Tuple       			=    []
   for i in range(config.num_hidden_layers     ):
      # encoder layers: output projection, 2 feedforward neural networks and 2 layernorms
      rename_keys.append(
          (F'{prefix}blocks.{i}.norm1.weight', F'{hf_prefix}encoder.layer.{i}.layernorm_before.weight')     )
      rename_keys.append((F'{prefix}blocks.{i}.norm1.bias', F'{hf_prefix}encoder.layer.{i}.layernorm_before.bias')     )
      rename_keys.append(
          (F'{prefix}blocks.{i}.attn.proj.weight', F'{hf_prefix}encoder.layer.{i}.attention.output.dense.weight')     )
      rename_keys.append(
          (F'{prefix}blocks.{i}.attn.proj.bias', F'{hf_prefix}encoder.layer.{i}.attention.output.dense.bias')     )
      rename_keys.append(
          (F'{prefix}blocks.{i}.norm2.weight', F'{hf_prefix}encoder.layer.{i}.layernorm_after.weight')     )
      rename_keys.append((F'{prefix}blocks.{i}.norm2.bias', F'{hf_prefix}encoder.layer.{i}.layernorm_after.bias')     )
      rename_keys.append(
          (F'{prefix}blocks.{i}.mlp.fc1.weight', F'{hf_prefix}encoder.layer.{i}.intermediate.dense.weight')     )
      rename_keys.append(
          (F'{prefix}blocks.{i}.mlp.fc1.bias', F'{hf_prefix}encoder.layer.{i}.intermediate.dense.bias')     )
      rename_keys.append((F'{prefix}blocks.{i}.mlp.fc2.weight', F'{hf_prefix}encoder.layer.{i}.output.dense.weight')     )
      rename_keys.append((F'{prefix}blocks.{i}.mlp.fc2.bias', F'{hf_prefix}encoder.layer.{i}.output.dense.bias')     )

   # projection layer + position embeddings
   rename_keys.extend(
       [
           (F'{prefix}cls_token', F'{hf_prefix}embeddings.cls_token'),
           (F'{prefix}patch_embed.proj.weight', F'{hf_prefix}embeddings.patch_embeddings.projection.weight'),
           (F'{prefix}patch_embed.proj.bias', F'{hf_prefix}embeddings.patch_embeddings.projection.bias'),
       ]     )

   if has_lm_head:
      # mask token + shared relative position bias + layernorm
      rename_keys.extend(
          [
              ("""mask_token""", F'{hf_prefix}embeddings.mask_token'),
              (
                  """rel_pos_bias.relative_position_bias_table""",
                  F'{hf_prefix}encoder.relative_position_bias.relative_position_bias_table',
              ),
              (
                  """rel_pos_bias.relative_position_index""",
                  F'{hf_prefix}encoder.relative_position_bias.relative_position_index',
              ),
              ("""norm.weight""", """layernorm.weight"""),
              ("""norm.bias""", """layernorm.bias"""),
          ]     )
   elif is_semantic:
      # semantic segmentation classification heads
      rename_keys.extend(
          [
              ("""decode_head.conv_seg.weight""", """decode_head.classifier.weight"""),
              ("""decode_head.conv_seg.bias""", """decode_head.classifier.bias"""),
              ("""auxiliary_head.conv_seg.weight""", """auxiliary_head.classifier.weight"""),
              ("""auxiliary_head.conv_seg.bias""", """auxiliary_head.classifier.bias"""),
          ]     )
   else:
      # layernorm + classification head
      rename_keys.extend(
          [
              ("""fc_norm.weight""", F'{hf_prefix}pooler.layernorm.weight'),
              ("""fc_norm.bias""", F'{hf_prefix}pooler.layernorm.bias'),
              ("""head.weight""", """classifier.weight"""),
              ("""head.bias""", """classifier.bias"""),
          ]     )

   return rename_keys







def    lowerCAmelCase_  (lowerCAmelCase__:		Dict		,	lowerCAmelCase__:		Union[str, Any]		,	lowerCAmelCase__:		str=False		,	lowerCAmelCase__:		Union[str, Any]=False		,	lowerCAmelCase__:		List[str]="data2vec_vision."     ):





   """simple docstring"""


   for i in range(config.num_hidden_layers     ):
      UpperCAmelCase_:		Any       			=    """backbone.""" if is_semantic else """"""
      # queries, keys and values
      UpperCAmelCase_:		Union[str, Any]       			=    state_dict.pop(F'{prefix}blocks.{i}.attn.qkv.weight'     )
      UpperCAmelCase_:		Optional[int]       			=    state_dict.pop(F'{prefix}blocks.{i}.attn.q_bias'     )
      UpperCAmelCase_:		str       			=    state_dict.pop(F'{prefix}blocks.{i}.attn.v_bias'     )

      UpperCAmelCase_:		str       			=    in_proj_weight[
          : config.hidden_size, :
      ]
      UpperCAmelCase_:		Dict       			=    q_bias
      UpperCAmelCase_:		List[Any]       			=    in_proj_weight[
          config.hidden_size : config.hidden_size * 2, :
      ]
      UpperCAmelCase_:		Union[str, Any]       			=    in_proj_weight[
          -config.hidden_size :, :
      ]
      UpperCAmelCase_:		Dict       			=    v_bias

      # gamma_1 and gamma_2
      # we call them lambda because otherwise they are renamed when using .from_pretrained
      UpperCAmelCase_:		Tuple       			=    state_dict.pop(F'{prefix}blocks.{i}.gamma_1'     )
      UpperCAmelCase_:		Any       			=    state_dict.pop(F'{prefix}blocks.{i}.gamma_2'     )

      UpperCAmelCase_:		Any       			=    gamma_a
      UpperCAmelCase_:		List[Any]       			=    gamma_a

      # relative_position bias table + index
      if not has_lm_head:
         # each layer has its own relative position bias
         UpperCAmelCase_:		Optional[Any]       			=    state_dict.pop(F'{prefix}blocks.{i}.attn.relative_position_bias_table'     )
         UpperCAmelCase_:		Tuple       			=    state_dict.pop(F'{prefix}blocks.{i}.attn.relative_position_index'     )

         UpperCAmelCase_:		str       			=    table
         UpperCAmelCase_:		int       			=    index







def    lowerCAmelCase_  ():





   """simple docstring"""


   UpperCAmelCase_:		Optional[int]       			=    argparse.ArgumentParser(
       """Convert Data2VecVision to HF for image classification and pretraining"""		,	add_help=lowerCAmelCase__     )
   parser.add_argument("""--hf_checkpoint_name"""		,	type=lowerCAmelCase__     )
   parser.add_argument("""--input_size"""		,	default=2_2_4		,	type=lowerCAmelCase__		,	help="""images input size"""     )
   parser.add_argument("""--beit_checkpoint"""		,	default=""""""		,	help="""beit checkpoint"""     )

   return parser.parse_args()







def    lowerCAmelCase_  (lowerCAmelCase__:		Optional[Any]		,	lowerCAmelCase__:		int		,	lowerCAmelCase__:		List[str]     ):





   """simple docstring"""


   def load_state_dict(lowerCAmelCase__:		int		,	lowerCAmelCase__:		Tuple		,	lowerCAmelCase__:		Tuple=""		,	lowerCAmelCase__:		List[Any]="relative_position_index"     ):
      UpperCAmelCase_:		str       			=    []
      UpperCAmelCase_:		Dict       			=    []
      UpperCAmelCase_:		Any       			=    []
      # copy state_dict so _load_from_state_dict can modify it
      UpperCAmelCase_:		str       			=    getattr(lowerCAmelCase__		,	"""_metadata"""		,	lowerCAmelCase__     )
      UpperCAmelCase_:		Union[str, Any]       			=    state_dict.copy()
      if metadata is not None:
         UpperCAmelCase_:		Union[str, Any]       			=    metadata

      def load(lowerCAmelCase__:		Tuple		,	lowerCAmelCase__:		List[str]=""     ):
         UpperCAmelCase_:		Dict       			=    {} if metadata is None else metadata.get(prefix[:-1]		,	{}     )
         module._load_from_state_dict(
             lowerCAmelCase__		,	lowerCAmelCase__		,	lowerCAmelCase__		,	lowerCAmelCase__		,	lowerCAmelCase__		,	lowerCAmelCase__		,	lowerCAmelCase__     )
         for name, child in module._modules.items():
            if child is not None:
               load(lowerCAmelCase__		,	prefix + name + """."""     )

      load(lowerCAmelCase__		,	prefix=lowerCAmelCase__     )

      UpperCAmelCase_:		Tuple       			=    []
      UpperCAmelCase_:		Optional[Any]       			=    []
      for key in missing_keys:
         UpperCAmelCase_:		Dict       			=    True
         for ignore_key in ignore_missing.split("""|"""     ):
            if ignore_key in key:
               UpperCAmelCase_:		Tuple       			=    False
               break
         if keep_flag:
            warn_missing_keys.append(lowerCAmelCase__     )
         else:
            ignore_missing_keys.append(lowerCAmelCase__     )

      UpperCAmelCase_:		Union[str, Any]       			=    warn_missing_keys

      if len(lowerCAmelCase__     ) > 0:
         print(
             """Weights of {} not initialized from pretrained model: {}""".format(
                 model.__class__.__name__		,	lowerCAmelCase__     )     )
      if len(lowerCAmelCase__     ) > 0:
         print("""Weights from pretrained model not used in {}: {}""".format(model.__class__.__name__		,	lowerCAmelCase__     )     )
      if len(lowerCAmelCase__     ) > 0:
         print(
             """Ignored weights of {} not initialized from pretrained model: {}""".format(
                 model.__class__.__name__		,	lowerCAmelCase__     )     )
      if len(lowerCAmelCase__     ) > 0:
         print("""\n""".join(lowerCAmelCase__     )     )

   UpperCAmelCase_:		List[Any]       			=    {
       """pretrained""": False,
       """use_shared_rel_pos_bias""": True,
       """use_abs_pos_emb""": False,
       """init_values""": 0.1,
   }

   if is_finetuned:
      model_kwargs.update(
          {
              """num_classes""": 1_0_0_0,
              """use_mean_pooling""": True,
              """init_scale""": 0.001,
              """use_rel_pos_bias""": True,
          }     )

   UpperCAmelCase_:		Tuple       			=    create_model(
       """beit_large_patch16_224""" if is_large else """beit_base_patch16_224"""		,	**lowerCAmelCase__		,	)
   UpperCAmelCase_:		Any       			=    model.patch_embed.patch_size
   UpperCAmelCase_:		str       			=    (args.input_size // patch_size[0], args.input_size // patch_size[1])
   UpperCAmelCase_:		Optional[int]       			=    torch.load(args.beit_checkpoint		,	map_location="""cpu"""     )

   print(F'Load ckpt from {args.beit_checkpoint}'     )
   UpperCAmelCase_:		List[Any]       			=    None
   for model_key in ("model", "module"):
      if model_key in checkpoint:
         UpperCAmelCase_:		Tuple       			=    checkpoint[model_key]
         print(F'Load state_dict by model_key = {model_key}'     )
         break

   UpperCAmelCase_:		Optional[int]       			=    list(checkpoint_model.keys()     )
   for key in all_keys:
      if "relative_position_index" in key:
         checkpoint_model.pop(lowerCAmelCase__     )

      if "relative_position_bias_table" in key:
         UpperCAmelCase_:		Tuple       			=    checkpoint_model[key]
         UpperCAmelCase_:		Dict       			=    rel_pos_bias.size()
         UpperCAmelCase_:		List[Any]       			=    model.state_dict()[key].size()
         UpperCAmelCase_:		int       			=    model.patch_embed.patch_shape
         if dst_patch_shape[0] != dst_patch_shape[1]:
            raise NotImplementedError()

   load_state_dict(lowerCAmelCase__		,	lowerCAmelCase__		,	prefix=""""""     )

   return model







def    lowerCAmelCase_  ():





   """simple docstring"""


   UpperCAmelCase_:		Any       			=    get_args()

   UpperCAmelCase_:		Dict       			=    """ft1k""" in args.hf_checkpoint_name
   UpperCAmelCase_:		List[str]       			=    """large""" in args.hf_checkpoint_name

   if is_finetuned:
      # To convert Beit's data2vec_vision to HF you need to copy
      # https://github.com/facebookresearch/data2vec_vision/blob/main/beit/modeling_finetune.py
      # into this folder.
      import modeling_finetune  # noqa: F401
   else:
      # To convert Beit's data2vec_vision to HF you need to copy
      # https://github.com/facebookresearch/data2vec_vision/blob/main/beit/modeling_cyclical.py
      # into this folder
      # IMPORTANT: Note that for now we've only converted the down-stream
      # model and not the full pretrained model. This means for the integration
      # test you need to add a `return x` after the following line:
      # https://github.com/facebookresearch/data2vec_vision/blob/af9a36349aaed59ae66e69b5dabeef2d62fdc5da/beit/modeling_cyclical.py#L197
      # to make the integration test pass.
      import modeling_cyclical  # noqa: F401

   # 1. Create model config
   UpperCAmelCase_:		Optional[Any]       			=    DataaVecVisionConfig()
   if is_finetuned:
      UpperCAmelCase_:		Any       			=    True
      UpperCAmelCase_:		str       			=    False
      UpperCAmelCase_:		Tuple       			=    True
      UpperCAmelCase_:		int       			=    1_0_0_0

      UpperCAmelCase_:		List[Any]       			=    """huggingface/label-files"""
      UpperCAmelCase_:		Any       			=    """imagenet-1k-id2label.json"""
      UpperCAmelCase_:		Tuple       			=    json.load(open(hf_hub_download(lowerCAmelCase__		,	lowerCAmelCase__		,	repo_type="""dataset"""     )		,	"""r"""     )     )
      UpperCAmelCase_:		Optional[int]       			=    {int(lowerCAmelCase__     ): v for k, v in idalabel.items()}
      UpperCAmelCase_:		List[Any]       			=    idalabel
      UpperCAmelCase_:		str       			=    {v: k for k, v in idalabel.items()}
   else:
      UpperCAmelCase_:		str       			=    False
      UpperCAmelCase_:		Dict       			=    True
      UpperCAmelCase_:		Tuple       			=    False

   if is_large:
      UpperCAmelCase_:		Union[str, Any]       			=    1_0_2_4
      UpperCAmelCase_:		List[Any]       			=    4_0_9_6
      UpperCAmelCase_:		Union[str, Any]       			=    2_4
      UpperCAmelCase_:		Optional[Any]       			=    1_6

   # 2. Load Beit model
   UpperCAmelCase_:		Optional[Any]       			=    load_beit_model(lowerCAmelCase__		,	lowerCAmelCase__		,	lowerCAmelCase__     )
   orig_model.eval()

   # 3. Forward Beit model
   UpperCAmelCase_:		int       			=    BeitImageProcessor(size=config.image_size		,	do_center_crop=lowerCAmelCase__     )
   UpperCAmelCase_:		int       			=    Image.open("""../../../../tests/fixtures/tests_samples/COCO/000000039769.png"""     )
   UpperCAmelCase_:		Optional[Any]       			=    image_processor(images=lowerCAmelCase__		,	return_tensors="""pt"""     )
   UpperCAmelCase_:		List[str]       			=    encoding["""pixel_values"""]

   UpperCAmelCase_:		str       			=    (pixel_values,) if is_finetuned else (pixel_values, None)
   with torch.no_grad():
      UpperCAmelCase_:		str       			=    orig_model(*lowerCAmelCase__     )

   # 4. Load HF Data2VecVision model
   if is_finetuned:
      UpperCAmelCase_:		Union[str, Any]       			=    DataaVecVisionForImageClassification(lowerCAmelCase__     )
      hf_model.eval()
      UpperCAmelCase_:		List[str]       			=    False
      UpperCAmelCase_:		Optional[int]       			=    """data2vec_vision."""
   else:
      UpperCAmelCase_:		int       			=    DataaVecVisionModel(lowerCAmelCase__     )
      hf_model.eval()
      UpperCAmelCase_:		List[Any]       			=    True
      UpperCAmelCase_:		Any       			=    """"""

   UpperCAmelCase_:		Optional[Any]       			=    create_rename_keys(lowerCAmelCase__		,	hf_prefix=lowerCAmelCase__		,	has_lm_head=lowerCAmelCase__     )
   UpperCAmelCase_:		int       			=    orig_model.state_dict()
   for src, dest in rename_keys:
      UpperCAmelCase_:		Any       			=    state_dict.pop(lowerCAmelCase__     )
      UpperCAmelCase_:		Dict       			=    val

   read_in_q_k_v(lowerCAmelCase__		,	lowerCAmelCase__		,	hf_prefix=lowerCAmelCase__		,	has_lm_head=lowerCAmelCase__     )
   UpperCAmelCase_:		Optional[Any]       			=    hf_model.load_state_dict(lowerCAmelCase__		,	strict=lowerCAmelCase__     )
   print("""HF missing"""		,	lowerCAmelCase__     )
   print("""HF unexpected_keys"""		,	lowerCAmelCase__     )

   # 5. Forward HF Data2VecVision model
   with torch.no_grad():
      UpperCAmelCase_:		Union[str, Any]       			=    hf_model(lowerCAmelCase__     )

   UpperCAmelCase_:		List[str]       			=    hf_model_output.logits if is_finetuned else hf_model_output.last_hidden_state

   # 6. Compare
   UpperCAmelCase_:		Union[str, Any]       			=    torch.max(torch.abs(hf_output - orig_model_output     )     ).item()

   print(F'max_absolute_diff = {max_absolute_diff}'     )
   UpperCAmelCase_:		Optional[Any]       			=    torch.allclose(lowerCAmelCase__		,	lowerCAmelCase__		,	atol=1e-3     )
   print("""Do both models output the same tensors?"""		,	"""🔥""" if success else """💩"""     )
   if not success:
      raise Exception("""Something went wRoNg"""     )

   # 7. Save
   print(F'Saving to {args.hf_checkpoint_name}'     )
   hf_model.save_pretrained(args.hf_checkpoint_name     )
   image_processor.save_pretrained(args.hf_checkpoint_name     )


if __name__ == "__main__":
 main()
 # Run the following to convert checkpoints
 #  python ./convert_data2vec_vision_original_pytorch_checkpoint_to_pytorch.py \
 #          --beit_checkpoint ./pretrained_base.pt \
 #          --hf_checkpoint_name "./data2vec-vision-base"
 #  python ./convert_data2vec_vision_original_pytorch_checkpoint_to_pytorch.py \
 #          --beit_checkpoint ./finetuned_base.pt \
 #          --hf_checkpoint_name "./data2vec-vision-base-ft1k"
 #  python ./convert_data2vec_vision_original_pytorch_checkpoint_to_pytorch.py \
 #          --beit_checkpoint ./pretrained_large.pt \
 #          --hf_checkpoint_name "./data2vec-vision-large"
 #  python ./convert_data2vec_vision_original_pytorch_checkpoint_to_pytorch.py \
 #          --beit_checkpoint ./finetuned_large.pt \
 #          --hf_checkpoint_name "./data2vec-vision-large-ft1k"











import gc
import random
import tempfile
import unittest

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

from diffusers import AutoencoderKL, DDIMScheduler, LMSDiscreteScheduler, PNDMScheduler, UNetaDConditionModel
from diffusers.pipelines.semantic_stable_diffusion import SemanticStableDiffusionPipeline as StableDiffusionPipeline
from diffusers.utils import floats_tensor, nightly, torch_device
from diffusers.utils.testing_utils import enable_full_determinism, require_torch_gpu


enable_full_determinism()







class 							_a							(      unittest.TestCase				):

			def     __snake_case					(self				)					->      Any:
						# clean up the VRAM after each test
						super().tearDown()
						gc.collect()
						torch.cuda.empty_cache()

			@property
			def     __snake_case					(self				)					->      Tuple:
						UpperCAmelCase_:		Tuple       			=    1
						UpperCAmelCase_:		Dict       			=    3
						UpperCAmelCase_:		Optional[int]       			=    (32, 32)

						UpperCAmelCase_:		int       			=    floats_tensor((batch_size, num_channels) + sizes,      rng=random.Random(0				)				).to(SCREAMING_SNAKE_CASE_				)
						return image

			@property
			def     __snake_case					(self				)					->      Optional[Any]:
						torch.manual_seed(0				)
						UpperCAmelCase_:		Tuple       			=    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,      )
						return model

			@property
			def     __snake_case					(self				)					->      str:
						torch.manual_seed(0				)
						UpperCAmelCase_:		int       			=    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,      )
						return model

			@property
			def     __snake_case					(self				)					->      List[str]:
						torch.manual_seed(0				)
						UpperCAmelCase_:		Tuple       			=    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,      )
						return CLIPTextModel(SCREAMING_SNAKE_CASE_				)

			@property
			def     __snake_case					(self				)					->      str:
						def extract(*SCREAMING_SNAKE_CASE_,      **SCREAMING_SNAKE_CASE_				):







									class 							_a							:

												def __init__(self				)					->      Any:
															UpperCAmelCase_:		Optional[Any]       			=    torch.ones([0]				)

												def     __snake_case					(self,      SCREAMING_SNAKE_CASE_				)					->      int:
															self.pixel_values.to(SCREAMING_SNAKE_CASE_				)
															return self






									return Out()

						return extract

			def     __snake_case					(self				)					->      List[Any]:
						UpperCAmelCase_:		Any       			=    """cpu"""  # ensure determinism for the device-dependent torch.Generator
						UpperCAmelCase_:		Dict       			=    self.dummy_cond_unet
						UpperCAmelCase_:		Union[str, Any]       			=    DDIMScheduler(
						    beta_start=0.0_0_0_8_5,      beta_end=0.0_1_2,      beta_schedule="""scaled_linear""",      clip_sample=SCREAMING_SNAKE_CASE_,      set_alpha_to_one=SCREAMING_SNAKE_CASE_,      )

						UpperCAmelCase_:		List[str]       			=    self.dummy_vae
						UpperCAmelCase_:		Union[str, Any]       			=    self.dummy_text_encoder
						UpperCAmelCase_:		Dict       			=    CLIPTokenizer.from_pretrained("""hf-internal-testing/tiny-random-clip"""				)

						# make sure here that pndm scheduler skips prk
						UpperCAmelCase_:		str       			=    StableDiffusionPipeline(
						    unet=SCREAMING_SNAKE_CASE_,      scheduler=SCREAMING_SNAKE_CASE_,      vae=SCREAMING_SNAKE_CASE_,      text_encoder=SCREAMING_SNAKE_CASE_,      tokenizer=SCREAMING_SNAKE_CASE_,      safety_checker=SCREAMING_SNAKE_CASE_,      feature_extractor=self.dummy_extractor,      )
						UpperCAmelCase_:		str       			=    sd_pipe.to(SCREAMING_SNAKE_CASE_				)
						sd_pipe.set_progress_bar_config(disable=SCREAMING_SNAKE_CASE_				)

						UpperCAmelCase_:		int       			=    """A painting of a squirrel eating a burger"""

						UpperCAmelCase_:		Tuple       			=    torch.Generator(device=SCREAMING_SNAKE_CASE_				).manual_seed(0				)
						UpperCAmelCase_:		Dict       			=    sd_pipe([prompt],      generator=SCREAMING_SNAKE_CASE_,      guidance_scale=6.0,      num_inference_steps=2,      output_type="""np"""				)
						UpperCAmelCase_:		Tuple       			=    output.images

						UpperCAmelCase_:		List[Any]       			=    torch.Generator(device=SCREAMING_SNAKE_CASE_				).manual_seed(0				)
						UpperCAmelCase_:		List[str]       			=    sd_pipe(
						    [prompt],      generator=SCREAMING_SNAKE_CASE_,      guidance_scale=6.0,      num_inference_steps=2,      output_type="""np""",      return_dict=SCREAMING_SNAKE_CASE_,      )[0]

						UpperCAmelCase_:		int       			=    image[0, -3:, -3:, -1]
						UpperCAmelCase_:		List[str]       			=    image_from_tuple[0, -3:, -3:, -1]

						assert image.shape == (1, 64, 64, 3)
						UpperCAmelCase_:		Dict       			=    np.array([0.5_7_5_3, 0.6_1_1_4, 0.5_0_0_1, 0.5_0_3_4, 0.5_4_7_0, 0.4_7_2_9, 0.4_9_7_1, 0.4_8_6_7, 0.4_8_6_7]				)

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

			def     __snake_case					(self				)					->      Any:
						UpperCAmelCase_:		List[Any]       			=    """cpu"""  # ensure determinism for the device-dependent torch.Generator
						UpperCAmelCase_:		Tuple       			=    self.dummy_cond_unet
						UpperCAmelCase_:		Optional[int]       			=    PNDMScheduler(skip_prk_steps=SCREAMING_SNAKE_CASE_				)
						UpperCAmelCase_:		Union[str, Any]       			=    self.dummy_vae
						UpperCAmelCase_:		Any       			=    self.dummy_text_encoder
						UpperCAmelCase_:		List[Any]       			=    CLIPTokenizer.from_pretrained("""hf-internal-testing/tiny-random-clip"""				)

						# make sure here that pndm scheduler skips prk
						UpperCAmelCase_:		List[str]       			=    StableDiffusionPipeline(
						    unet=SCREAMING_SNAKE_CASE_,      scheduler=SCREAMING_SNAKE_CASE_,      vae=SCREAMING_SNAKE_CASE_,      text_encoder=SCREAMING_SNAKE_CASE_,      tokenizer=SCREAMING_SNAKE_CASE_,      safety_checker=SCREAMING_SNAKE_CASE_,      feature_extractor=self.dummy_extractor,      )
						UpperCAmelCase_:		List[str]       			=    sd_pipe.to(SCREAMING_SNAKE_CASE_				)
						sd_pipe.set_progress_bar_config(disable=SCREAMING_SNAKE_CASE_				)

						UpperCAmelCase_:		int       			=    """A painting of a squirrel eating a burger"""
						UpperCAmelCase_:		Union[str, Any]       			=    torch.Generator(device=SCREAMING_SNAKE_CASE_				).manual_seed(0				)
						UpperCAmelCase_:		List[Any]       			=    sd_pipe([prompt],      generator=SCREAMING_SNAKE_CASE_,      guidance_scale=6.0,      num_inference_steps=2,      output_type="""np"""				)

						UpperCAmelCase_:		Optional[int]       			=    output.images

						UpperCAmelCase_:		Any       			=    torch.Generator(device=SCREAMING_SNAKE_CASE_				).manual_seed(0				)
						UpperCAmelCase_:		Tuple       			=    sd_pipe(
						    [prompt],      generator=SCREAMING_SNAKE_CASE_,      guidance_scale=6.0,      num_inference_steps=2,      output_type="""np""",      return_dict=SCREAMING_SNAKE_CASE_,      )[0]

						UpperCAmelCase_:		Optional[Any]       			=    image[0, -3:, -3:, -1]
						UpperCAmelCase_:		Tuple       			=    image_from_tuple[0, -3:, -3:, -1]

						assert image.shape == (1, 64, 64, 3)
						UpperCAmelCase_:		Tuple       			=    np.array([0.5_1_2_2, 0.5_7_1_2, 0.4_8_2_5, 0.5_0_5_3, 0.5_6_4_6, 0.4_7_6_9, 0.5_1_7_9, 0.4_8_9_4, 0.4_9_9_4]				)

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

			def     __snake_case					(self				)					->      List[str]:
						UpperCAmelCase_:		Dict       			=    StableDiffusionPipeline.from_pretrained(
						    """hf-internal-testing/tiny-stable-diffusion-lms-pipe""",      safety_checker=SCREAMING_SNAKE_CASE_				)
						assert isinstance(SCREAMING_SNAKE_CASE_,      SCREAMING_SNAKE_CASE_				)
						assert isinstance(pipe.scheduler,      SCREAMING_SNAKE_CASE_				)
						assert pipe.safety_checker is None

						UpperCAmelCase_:		Optional[Any]       			=    pipe("""example prompt""",      num_inference_steps=2				).images[0]
						assert image is not None

						# check that there's no error when saving a pipeline with one of the models being None
						with tempfile.TemporaryDirectory() as tmpdirname:
									pipe.save_pretrained(SCREAMING_SNAKE_CASE_				)
									UpperCAmelCase_:		Optional[int]       			=    StableDiffusionPipeline.from_pretrained(SCREAMING_SNAKE_CASE_				)

						# sanity check that the pipeline still works
						assert pipe.safety_checker is None
						UpperCAmelCase_:		Union[str, Any]       			=    pipe("""example prompt""",      num_inference_steps=2				).images[0]
						assert image is not None

			@unittest.skipIf(torch_device != """cuda""",      """This test requires a GPU"""				)
			def     __snake_case					(self				)					->      Optional[Any]:
						UpperCAmelCase_:		Dict       			=    self.dummy_cond_unet
						UpperCAmelCase_:		Any       			=    PNDMScheduler(skip_prk_steps=SCREAMING_SNAKE_CASE_				)
						UpperCAmelCase_:		List[str]       			=    self.dummy_vae
						UpperCAmelCase_:		Tuple       			=    self.dummy_text_encoder
						UpperCAmelCase_:		Dict       			=    CLIPTokenizer.from_pretrained("""hf-internal-testing/tiny-random-clip"""				)

						# put models in fp16
						UpperCAmelCase_:		Dict       			=    unet.half()
						UpperCAmelCase_:		Optional[Any]       			=    vae.half()
						UpperCAmelCase_:		Union[str, Any]       			=    bert.half()

						# make sure here that pndm scheduler skips prk
						UpperCAmelCase_:		str       			=    StableDiffusionPipeline(
						    unet=SCREAMING_SNAKE_CASE_,      scheduler=SCREAMING_SNAKE_CASE_,      vae=SCREAMING_SNAKE_CASE_,      text_encoder=SCREAMING_SNAKE_CASE_,      tokenizer=SCREAMING_SNAKE_CASE_,      safety_checker=SCREAMING_SNAKE_CASE_,      feature_extractor=self.dummy_extractor,      )
						UpperCAmelCase_:		Optional[Any]       			=    sd_pipe.to(SCREAMING_SNAKE_CASE_				)
						sd_pipe.set_progress_bar_config(disable=SCREAMING_SNAKE_CASE_				)

						UpperCAmelCase_:		str       			=    """A painting of a squirrel eating a burger"""
						UpperCAmelCase_:		Optional[Any]       			=    sd_pipe([prompt],      num_inference_steps=2,      output_type="""np"""				).images

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







@nightly
@require_torch_gpu
class 							_a							(      unittest.TestCase				):

			def     __snake_case					(self				)					->      Dict:
						# clean up the VRAM after each test
						super().tearDown()
						gc.collect()
						torch.cuda.empty_cache()

			def     __snake_case					(self				)					->      str:
						UpperCAmelCase_:		Optional[int]       			=    """cuda"""
						UpperCAmelCase_:		List[str]       			=    StableDiffusionPipeline.from_pretrained("""runwayml/stable-diffusion-v1-5"""				)
						UpperCAmelCase_:		Union[str, Any]       			=    pipe.to(SCREAMING_SNAKE_CASE_				)
						pipe.set_progress_bar_config(disable=SCREAMING_SNAKE_CASE_				)

						UpperCAmelCase_:		int       			=    """a photo of a cat"""
						UpperCAmelCase_:		Tuple       			=    {
						    """editing_prompt""": ["""sunglasses"""],
						    """reverse_editing_direction""": [False],
						    """edit_warmup_steps""": 10,
						    """edit_guidance_scale""": 6,
						    """edit_threshold""": 0.9_5,
						    """edit_momentum_scale""": 0.5,
						    """edit_mom_beta""": 0.6,
						}

						UpperCAmelCase_:		List[str]       			=    3
						UpperCAmelCase_:		str       			=    7

						# no sega enabled
						UpperCAmelCase_:		Tuple       			=    torch.Generator(SCREAMING_SNAKE_CASE_				)
						generator.manual_seed(SCREAMING_SNAKE_CASE_				)
						UpperCAmelCase_:		Any       			=    pipe(
						    [prompt],      generator=SCREAMING_SNAKE_CASE_,      guidance_scale=SCREAMING_SNAKE_CASE_,      num_inference_steps=50,      output_type="""np""",      width=512,      height=512,      )

						UpperCAmelCase_:		Optional[Any]       			=    output.images
						UpperCAmelCase_:		List[Any]       			=    image[0, -3:, -3:, -1]
						UpperCAmelCase_:		Optional[int]       			=    [
						    0.3_4_6_7_3_1_1_3,
						    0.3_8_4_9_2_7_3_3,
						    0.3_7_5_9_7_3_5_2,
						    0.3_4_0_8_6_3_3_5,
						    0.3_5_6_5_0_7_4_8,
						    0.3_5_5_7_9_2_0_5,
						    0.3_3_8_4_7_6_3,
						    0.3_4_3_4_0_2_3_6,
						    0.3_5_7_3_2_7_1,
						]

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

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

						# with sega enabled
						# generator = torch.manual_seed(seed)
						generator.manual_seed(SCREAMING_SNAKE_CASE_				)
						UpperCAmelCase_:		Optional[Any]       			=    pipe(
						    [prompt],      generator=SCREAMING_SNAKE_CASE_,      guidance_scale=SCREAMING_SNAKE_CASE_,      num_inference_steps=50,      output_type="""np""",      width=512,      height=512,      **SCREAMING_SNAKE_CASE_,      )

						UpperCAmelCase_:		Any       			=    output.images
						UpperCAmelCase_:		Dict       			=    image[0, -3:, -3:, -1]
						UpperCAmelCase_:		Optional[int]       			=    [
						    0.4_1_8_8_7_8_2_6,
						    0.3_7_7_2_8_7_6_6,
						    0.3_0_1_3_8_2_7_2,
						    0.4_1_4_1_6_3_3_5,
						    0.4_1_6_6_4_9_8_5,
						    0.3_6_2_8_3_3_9_2,
						    0.3_6_1_9_1_2_4_6,
						    0.4_3_3_6_4_4_6_5,
						    0.4_3_0_0_1_7_3_2,
						]

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

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

			def     __snake_case					(self				)					->      Optional[Any]:
						UpperCAmelCase_:		List[Any]       			=    """cuda"""
						UpperCAmelCase_:		Tuple       			=    StableDiffusionPipeline.from_pretrained("""runwayml/stable-diffusion-v1-5"""				)
						UpperCAmelCase_:		List[Any]       			=    pipe.to(SCREAMING_SNAKE_CASE_				)
						pipe.set_progress_bar_config(disable=SCREAMING_SNAKE_CASE_				)

						UpperCAmelCase_:		Dict       			=    """an image of a crowded boulevard, realistic, 4k"""
						UpperCAmelCase_:		Optional[Any]       			=    {
						    """editing_prompt""": """crowd, crowded, people""",
						    """reverse_editing_direction""": True,
						    """edit_warmup_steps""": 10,
						    """edit_guidance_scale""": 8.3,
						    """edit_threshold""": 0.9,
						    """edit_momentum_scale""": 0.5,
						    """edit_mom_beta""": 0.6,
						}

						UpperCAmelCase_:		Dict       			=    9
						UpperCAmelCase_:		Optional[int]       			=    7

						# no sega enabled
						UpperCAmelCase_:		List[Any]       			=    torch.Generator(SCREAMING_SNAKE_CASE_				)
						generator.manual_seed(SCREAMING_SNAKE_CASE_				)
						UpperCAmelCase_:		str       			=    pipe(
						    [prompt],      generator=SCREAMING_SNAKE_CASE_,      guidance_scale=SCREAMING_SNAKE_CASE_,      num_inference_steps=50,      output_type="""np""",      width=512,      height=512,      )

						UpperCAmelCase_:		Optional[Any]       			=    output.images
						UpperCAmelCase_:		List[Any]       			=    image[0, -3:, -3:, -1]
						UpperCAmelCase_:		int       			=    [
						    0.4_3_4_9_7_9_9_8,
						    0.9_1_8_1_4_0_6_5,
						    0.7_5_4_0_7_3_9,
						    0.5_5_5_8_0_2_0_5,
						    0.8_4_6_7_2_6_5,
						    0.5_3_8_9_6_9_1,
						    0.6_2_5_7_4_5_0_6,
						    0.5_8_8_9_7_7_6_3,
						    0.5_0_9_2_6_7_5_7,
						]

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

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

						# with sega enabled
						# generator = torch.manual_seed(seed)
						generator.manual_seed(SCREAMING_SNAKE_CASE_				)
						UpperCAmelCase_:		Optional[int]       			=    pipe(
						    [prompt],      generator=SCREAMING_SNAKE_CASE_,      guidance_scale=SCREAMING_SNAKE_CASE_,      num_inference_steps=50,      output_type="""np""",      width=512,      height=512,      **SCREAMING_SNAKE_CASE_,      )

						UpperCAmelCase_:		int       			=    output.images
						UpperCAmelCase_:		Any       			=    image[0, -3:, -3:, -1]
						UpperCAmelCase_:		str       			=    [
						    0.3_0_8_9_7_1_9,
						    0.3_0_5_0_0_1_4_4,
						    0.2_9_0_1_6_0_4_2,
						    0.3_0_6_3_0_9_6_4,
						    0.3_2_5_6_8_7,
						    0.2_9_4_1_9_2_2_5,
						    0.2_9_0_8_0_9_1,
						    0.2_8_7_2_3_5_9_8,
						    0.2_7_6_9_6_2_9_4,
						]

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

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

			def     __snake_case					(self				)					->      Tuple:
						UpperCAmelCase_:		Union[str, Any]       			=    """cuda"""
						UpperCAmelCase_:		Optional[Any]       			=    StableDiffusionPipeline.from_pretrained("""runwayml/stable-diffusion-v1-5"""				)
						UpperCAmelCase_:		Optional[Any]       			=    pipe.to(SCREAMING_SNAKE_CASE_				)
						pipe.set_progress_bar_config(disable=SCREAMING_SNAKE_CASE_				)

						UpperCAmelCase_:		Tuple       			=    """a castle next to a river"""
						UpperCAmelCase_:		Optional[Any]       			=    {
						    """editing_prompt""": ["""boat on a river, boat""", """monet, impression, sunrise"""],
						    """reverse_editing_direction""": False,
						    """edit_warmup_steps""": [15, 18],
						    """edit_guidance_scale""": 6,
						    """edit_threshold""": [0.9, 0.8],
						    """edit_momentum_scale""": 0.5,
						    """edit_mom_beta""": 0.6,
						}

						UpperCAmelCase_:		str       			=    48
						UpperCAmelCase_:		Optional[int]       			=    7

						# no sega enabled
						UpperCAmelCase_:		int       			=    torch.Generator(SCREAMING_SNAKE_CASE_				)
						generator.manual_seed(SCREAMING_SNAKE_CASE_				)
						UpperCAmelCase_:		Union[str, Any]       			=    pipe(
						    [prompt],      generator=SCREAMING_SNAKE_CASE_,      guidance_scale=SCREAMING_SNAKE_CASE_,      num_inference_steps=50,      output_type="""np""",      width=512,      height=512,      )

						UpperCAmelCase_:		Tuple       			=    output.images
						UpperCAmelCase_:		Tuple       			=    image[0, -3:, -3:, -1]
						UpperCAmelCase_:		Tuple       			=    [
						    0.7_5_1_6_3_5_5_5,
						    0.7_6_0_3_7_1_4_5,
						    0.6_1_7_8_5,
						    0.9_1_8_9_6_7_3,
						    0.8_6_2_7_7_0_1,
						    0.8_5_1_8_9_6_9_4,
						    0.8_5_1_2_8_1_3,
						    0.8_7_0_1_2_0_7_6,
						    0.8_3_1_2_8_5_7,
						]

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

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

						# with sega enabled
						# generator = torch.manual_seed(seed)
						generator.manual_seed(SCREAMING_SNAKE_CASE_				)
						UpperCAmelCase_:		Dict       			=    pipe(
						    [prompt],      generator=SCREAMING_SNAKE_CASE_,      guidance_scale=SCREAMING_SNAKE_CASE_,      num_inference_steps=50,      output_type="""np""",      width=512,      height=512,      **SCREAMING_SNAKE_CASE_,      )

						UpperCAmelCase_:		List[str]       			=    output.images
						UpperCAmelCase_:		Optional[int]       			=    image[0, -3:, -3:, -1]
						UpperCAmelCase_:		Optional[Any]       			=    [
						    0.7_3_5_5_3_3_6_5,
						    0.7_5_3_7_2_7_1,
						    0.7_4_3_4_1_9_0_5,
						    0.6_6_4_8_0_3_5_6,
						    0.6_4_7_2_9_2_5,
						    0.6_3_0_3_9_4_1_6,
						    0.6_4_8_1_2_9_0_5,
						    0.6_7_4_9_7_1_7,
						    0.6_5_1_7_1_0_2,
						]

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

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

			def     __snake_case					(self				)					->      Optional[int]:
						UpperCAmelCase_:		List[Any]       			=    """cuda"""
						UpperCAmelCase_:		Union[str, Any]       			=    StableDiffusionPipeline.from_pretrained("""runwayml/stable-diffusion-v1-5""",      torch_dtype=torch.floataa				)
						UpperCAmelCase_:		Optional[int]       			=    pipe.to(SCREAMING_SNAKE_CASE_				)
						pipe.set_progress_bar_config(disable=SCREAMING_SNAKE_CASE_				)

						UpperCAmelCase_:		Dict       			=    """a photo of a cat"""
						UpperCAmelCase_:		Dict       			=    {
						    """editing_prompt""": ["""sunglasses"""],
						    """reverse_editing_direction""": [False],
						    """edit_warmup_steps""": 10,
						    """edit_guidance_scale""": 6,
						    """edit_threshold""": 0.9_5,
						    """edit_momentum_scale""": 0.5,
						    """edit_mom_beta""": 0.6,
						}

						UpperCAmelCase_:		List[Any]       			=    3
						UpperCAmelCase_:		Dict       			=    7

						# no sega enabled
						UpperCAmelCase_:		Tuple       			=    torch.Generator(SCREAMING_SNAKE_CASE_				)
						generator.manual_seed(SCREAMING_SNAKE_CASE_				)
						UpperCAmelCase_:		List[Any]       			=    pipe(
						    [prompt],      generator=SCREAMING_SNAKE_CASE_,      guidance_scale=SCREAMING_SNAKE_CASE_,      num_inference_steps=50,      output_type="""np""",      width=512,      height=512,      )

						UpperCAmelCase_:		str       			=    output.images
						UpperCAmelCase_:		Optional[Any]       			=    image[0, -3:, -3:, -1]
						UpperCAmelCase_:		List[Any]       			=    [
						    0.3_4_8_8_7_6_9_5,
						    0.3_8_7_6_9_5_3,
						    0.3_7_5,
						    0.3_4_4_2_3_8_2_8,
						    0.3_5_8_1_5_4_3,
						    0.3_5_7_1_7_7_7_3,
						    0.3_3_8_3_7_8_9,
						    0.3_4_5_7_0_3_1_2,
						    0.3_5_9_3_7_5,
						]

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

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

						# with sega enabled
						# generator = torch.manual_seed(seed)
						generator.manual_seed(SCREAMING_SNAKE_CASE_				)
						UpperCAmelCase_:		Any       			=    pipe(
						    [prompt],      generator=SCREAMING_SNAKE_CASE_,      guidance_scale=SCREAMING_SNAKE_CASE_,      num_inference_steps=50,      output_type="""np""",      width=512,      height=512,      **SCREAMING_SNAKE_CASE_,      )

						UpperCAmelCase_:		List[Any]       			=    output.images
						UpperCAmelCase_:		List[str]       			=    image[0, -3:, -3:, -1]
						UpperCAmelCase_:		List[Any]       			=    [
						    0.4_2_2_8_5_1_5_6,
						    0.3_6_9_1_4_0_6_2,
						    0.2_9_0_7_7_1_4_8,
						    0.4_2_0_4_1_0_1_6,
						    0.4_1_9_1_8_9_4_5,
						    0.3_5_4_9_8_0_4_7,
						    0.3_6_1_8_1_6_4,
						    0.4_4_2_3_8_2_8,
						    0.4_3_1_1_5_2_3_4,
						]

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

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






