




import gc
import unittest

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

from diffusers import (
    AutoencoderKL,
    DDIMScheduler,
    PNDMScheduler,
    StableDiffusionLDMaDPipeline,
    UNetaDConditionModel,
)
from diffusers.utils import nightly, slow, torch_device
from diffusers.utils.testing_utils import enable_full_determinism, require_torch_gpu

from ..pipeline_params import TEXT_TO_IMAGE_BATCH_PARAMS, TEXT_TO_IMAGE_IMAGE_PARAMS, TEXT_TO_IMAGE_PARAMS


enable_full_determinism()





class 	lowerCamelCase_							(				unittest.TestCase ):


   '''simple docstring'''







   lowercase_       =   StableDiffusionLDMaDPipeline
   lowercase_       =   TEXT_TO_IMAGE_PARAMS
   lowercase_       =   TEXT_TO_IMAGE_BATCH_PARAMS
   lowercase_       =   TEXT_TO_IMAGE_IMAGE_PARAMS



   def 				lowerCAmelCase_  (							self  :		str						):
         torch.manual_seed(0						)
         SCREAMING_SNAKE_CASE_     						=      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_     						=      DDIMScheduler(
             beta_start=0.0_0085    ,    beta_end=0.012    ,    beta_schedule='scaled_linear'    ,    clip_sample=_a    ,    set_alpha_to_one=_a    ,    )
         torch.manual_seed(0						)
         SCREAMING_SNAKE_CASE_     						=      AutoencoderKL(
             block_out_channels=[32, 64]    ,    in_channels=6    ,    out_channels=6    ,    down_block_types=['DownEncoderBlock2D', 'DownEncoderBlock2D']    ,    up_block_types=['UpDecoderBlock2D', 'UpDecoderBlock2D']    ,    latent_channels=4    ,    )
         torch.manual_seed(0						)
         SCREAMING_SNAKE_CASE_     						=      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=1_000    ,    )
         SCREAMING_SNAKE_CASE_     						=      CLIPTextModel(_a						)
         SCREAMING_SNAKE_CASE_     						=      CLIPTokenizer.from_pretrained('hf-internal-testing/tiny-random-clip'						)

         SCREAMING_SNAKE_CASE_     						=      {
             "unet": unet,
             "scheduler": scheduler,
             "vae": vae,
             "text_encoder": text_encoder,
             "tokenizer": tokenizer,
             "safety_checker": None,
             "feature_extractor": None,
         }
         return components



   def 				lowerCAmelCase_  (							self  :		List[str]    ,    _lowerCAmelCase  :		Optional[Any]    ,    _lowerCAmelCase  :		Tuple=0						):
         if str(_a						).startswith('mps'						):
               SCREAMING_SNAKE_CASE_     						=      torch.manual_seed(_a						)
         else:
               SCREAMING_SNAKE_CASE_     						=      torch.Generator(device=_a						).manual_seed(_a						)
         SCREAMING_SNAKE_CASE_     						=      {
             "prompt": "A painting of a squirrel eating a burger",
             "generator": generator,
             "num_inference_steps": 2,
             "guidance_scale": 6.0,
             "output_type": "numpy",
         }
         return inputs



   def 				lowerCAmelCase_  (							self  :		Tuple						):
         SCREAMING_SNAKE_CASE_     						=      "cpu"  # ensure determinism for the device-dependent torch.Generator

         SCREAMING_SNAKE_CASE_     						=      self.get_dummy_components()
         SCREAMING_SNAKE_CASE_     						=      StableDiffusionLDMaDPipeline(**_a						)
         SCREAMING_SNAKE_CASE_     						=      ldmad_pipe.to(_a						)
         ldmad_pipe.set_progress_bar_config(disable=_a						)

         SCREAMING_SNAKE_CASE_     						=      self.get_dummy_inputs(_a						)
         SCREAMING_SNAKE_CASE_     						=      ldmad_pipe(**_a						)
         SCREAMING_SNAKE_CASE_     						=      output.rgb, output.depth

         SCREAMING_SNAKE_CASE_     						=      rgb[0, -3:, -3:, -1]
         SCREAMING_SNAKE_CASE_     						=      depth[0, -3:, -1]

         assert rgb.shape == (1, 64, 64, 3)
         assert depth.shape == (1, 64, 64)

         SCREAMING_SNAKE_CASE_     						=      np.array(
             [0.3733_8176, 0.7_0247, 0.7420_3193, 0.5164_3604, 0.5825_6793, 0.6093_2136, 0.418_1095, 0.4835_5877, 0.4653_5262]						)
         SCREAMING_SNAKE_CASE_     						=      np.array([103.46_727, 85.812_004, 87.849_236]						)

         assert np.abs(image_slice_rgb.flatten() - expected_slice_rgb						).max() < 1E-2
         assert np.abs(image_slice_depth.flatten() - expected_slice_depth						).max() < 1E-2



   def 				lowerCAmelCase_  (							self  :		str						):
         SCREAMING_SNAKE_CASE_     						=      self.get_dummy_components()
         SCREAMING_SNAKE_CASE_     						=      StableDiffusionLDMaDPipeline(**_a						)
         SCREAMING_SNAKE_CASE_     						=      ldmad_pipe.to(_a						)
         ldmad_pipe.set_progress_bar_config(disable=_a						)

         SCREAMING_SNAKE_CASE_     						=      self.get_dummy_inputs(_a						)
         SCREAMING_SNAKE_CASE_     						=      3 * [inputs["prompt"]]

         # forward
         SCREAMING_SNAKE_CASE_     						=      ldmad_pipe(**_a						)
         SCREAMING_SNAKE_CASE_     						=      output.rgb, output.depth
         SCREAMING_SNAKE_CASE_     						=      rgb_slice_a[0, -3:, -3:, -1]
         SCREAMING_SNAKE_CASE_     						=      depth_slice_a[0, -3:, -1]

         SCREAMING_SNAKE_CASE_     						=      self.get_dummy_inputs(_a						)
         SCREAMING_SNAKE_CASE_     						=      3 * [inputs.pop('prompt'						)]

         SCREAMING_SNAKE_CASE_     						=      ldmad_pipe.tokenizer(
             _a    ,    padding='max_length'    ,    max_length=ldmad_pipe.tokenizer.model_max_length    ,    truncation=_a    ,    return_tensors='pt'    ,    )
         SCREAMING_SNAKE_CASE_     						=      text_inputs["input_ids"].to(_a						)

         SCREAMING_SNAKE_CASE_     						=      ldmad_pipe.text_encoder(_a						)[0]

         SCREAMING_SNAKE_CASE_     						=      prompt_embeds

         # forward
         SCREAMING_SNAKE_CASE_     						=      ldmad_pipe(**_a						)
         SCREAMING_SNAKE_CASE_     						=      output.rgb, output.depth
         SCREAMING_SNAKE_CASE_     						=      rgb_slice_a[0, -3:, -3:, -1]
         SCREAMING_SNAKE_CASE_     						=      depth_slice_a[0, -3:, -1]

         assert np.abs(rgb_slice_a.flatten() - rgb_slice_a.flatten()						).max() < 1E-4
         assert np.abs(depth_slice_a.flatten() - depth_slice_a.flatten()						).max() < 1E-4




   def 				lowerCAmelCase_  (							self  :		List[str]						):
         SCREAMING_SNAKE_CASE_     						=      "cpu"  # ensure determinism for the device-dependent torch.Generator
         SCREAMING_SNAKE_CASE_     						=      self.get_dummy_components()
         SCREAMING_SNAKE_CASE_     						=      PNDMScheduler(skip_prk_steps=_a						)
         SCREAMING_SNAKE_CASE_     						=      StableDiffusionLDMaDPipeline(**_a						)
         SCREAMING_SNAKE_CASE_     						=      ldmad_pipe.to(_a						)
         ldmad_pipe.set_progress_bar_config(disable=_a						)

         SCREAMING_SNAKE_CASE_     						=      self.get_dummy_inputs(_a						)
         SCREAMING_SNAKE_CASE_     						=      "french fries"
         SCREAMING_SNAKE_CASE_     						=      ldmad_pipe(**_a    ,    negative_prompt=_a						)

         SCREAMING_SNAKE_CASE_     						=      output.rgb, output.depth
         SCREAMING_SNAKE_CASE_     						=      rgb[0, -3:, -3:, -1]
         SCREAMING_SNAKE_CASE_     						=      depth[0, -3:, -1]

         assert rgb.shape == (1, 64, 64, 3)
         assert depth.shape == (1, 64, 64)

         SCREAMING_SNAKE_CASE_     						=      np.array(
             [0.3_7044, 0.7181_1503, 0.722_3251, 0.4860_3675, 0.563_8391, 0.636_4948, 0.4283_3704, 0.490_1315, 0.4792_6217]						)
         SCREAMING_SNAKE_CASE_     						=      np.array([107.84_738, 84.62_802, 89.962_135]						)
         assert np.abs(rgb_slice.flatten() - expected_slice_rgb						).max() < 1E-2
         assert np.abs(depth_slice.flatten() - expected_slice_depth						).max() < 1E-2





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


   '''simple docstring'''







   def 				lowerCAmelCase_  (							self  :		List[str]						):
         super().tearDown()
         gc.collect()
         torch.cuda.empty_cache()



   def 				lowerCAmelCase_  (							self  :		List[str]    ,    _lowerCAmelCase  :		Dict    ,    _lowerCAmelCase  :		List[str]="cpu"    ,    _lowerCAmelCase  :		int=torch.floataa    ,    _lowerCAmelCase  :		Tuple=0						):
         SCREAMING_SNAKE_CASE_     						=      torch.Generator(device=_a						).manual_seed(_a						)
         SCREAMING_SNAKE_CASE_     						=      np.random.RandomState(_a						).standard_normal((1, 4, 64, 64)						)
         SCREAMING_SNAKE_CASE_     						=      torch.from_numpy(_a						).to(device=_a    ,    dtype=_a						)
         SCREAMING_SNAKE_CASE_     						=      {
             "prompt": "a photograph of an astronaut riding a horse",
             "latents": latents,
             "generator": generator,
             "num_inference_steps": 3,
             "guidance_scale": 7.5,
             "output_type": "numpy",
         }
         return inputs




   def 				lowerCAmelCase_  (							self  :		int						):
         SCREAMING_SNAKE_CASE_     						=      StableDiffusionLDMaDPipeline.from_pretrained('Intel/ldm3d'						)
         SCREAMING_SNAKE_CASE_     						=      ldmad_pipe.to(_a						)
         ldmad_pipe.set_progress_bar_config(disable=_a						)

         SCREAMING_SNAKE_CASE_     						=      self.get_inputs(_a						)
         SCREAMING_SNAKE_CASE_     						=      ldmad_pipe(**_a						)
         SCREAMING_SNAKE_CASE_     						=      output.rgb, output.depth
         SCREAMING_SNAKE_CASE_     						=      rgb[0, -3:, -3:, -1].flatten()
         SCREAMING_SNAKE_CASE_     						=      rgb[0, -3:, -1].flatten()

         assert rgb.shape == (1, 512, 512, 3)
         assert depth.shape == (1, 512, 512)

         SCREAMING_SNAKE_CASE_     						=      np.array(
             [0.5380_5465, 0.5670_7305, 0.548_6515, 0.5701_2236, 0.581_4511, 0.5625_3487, 0.5484_3014, 0.5509_2263, 0.645_9706]						)
         SCREAMING_SNAKE_CASE_     						=      np.array(
             [0.926_3781, 0.667_8672, 0.548_6515, 0.9220_2145, 0.6783_1135, 0.5625_3487, 0.924_1694, 0.755_1478, 0.645_9706]						)
         assert np.abs(rgb_slice - expected_slice_rgb						).max() < 3E-3
         assert np.abs(depth_slice - expected_slice_depth						).max() < 3E-3





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


   '''simple docstring'''







   def 				lowerCAmelCase_  (							self  :		str						):
         super().tearDown()
         gc.collect()
         torch.cuda.empty_cache()



   def 				lowerCAmelCase_  (							self  :		int    ,    _lowerCAmelCase  :		str    ,    _lowerCAmelCase  :		Any="cpu"    ,    _lowerCAmelCase  :		Any=torch.floataa    ,    _lowerCAmelCase  :		Optional[int]=0						):
         SCREAMING_SNAKE_CASE_     						=      torch.Generator(device=_a						).manual_seed(_a						)
         SCREAMING_SNAKE_CASE_     						=      np.random.RandomState(_a						).standard_normal((1, 4, 64, 64)						)
         SCREAMING_SNAKE_CASE_     						=      torch.from_numpy(_a						).to(device=_a    ,    dtype=_a						)
         SCREAMING_SNAKE_CASE_     						=      {
             "prompt": "a photograph of an astronaut riding a horse",
             "latents": latents,
             "generator": generator,
             "num_inference_steps": 50,
             "guidance_scale": 7.5,
             "output_type": "numpy",
         }
         return inputs



   def 				lowerCAmelCase_  (							self  :		Any						):
         SCREAMING_SNAKE_CASE_     						=      StableDiffusionLDMaDPipeline.from_pretrained('Intel/ldm3d'						).to(_a						)
         ldmad_pipe.set_progress_bar_config(disable=_a						)

         SCREAMING_SNAKE_CASE_     						=      self.get_inputs(_a						)
         SCREAMING_SNAKE_CASE_     						=      ldmad_pipe(**_a						)
         SCREAMING_SNAKE_CASE_     						=      output.rgb, output.depth

         SCREAMING_SNAKE_CASE_     						=      0.49_5586
         SCREAMING_SNAKE_CASE_     						=      0.3379_5515
         SCREAMING_SNAKE_CASE_     						=      112.48_518
         SCREAMING_SNAKE_CASE_     						=      98.489_746
         assert np.abs(expected_rgb_mean - rgb.mean()						) < 1E-3
         assert np.abs(expected_rgb_std - rgb.std()						) < 1E-3
         assert np.abs(expected_depth_mean - depth.mean()						) < 1E-3
         assert np.abs(expected_depth_std - depth.std()						) < 1E-3




   def 				lowerCAmelCase_  (							self  :		Dict						):
         SCREAMING_SNAKE_CASE_     						=      StableDiffusionLDMaDPipeline.from_pretrained('Intel/ldm3d-4c'						).to(_a						)
         ldmad_pipe.set_progress_bar_config(disable=_a						)

         SCREAMING_SNAKE_CASE_     						=      self.get_inputs(_a						)
         SCREAMING_SNAKE_CASE_     						=      ldmad_pipe(**_a						)
         SCREAMING_SNAKE_CASE_     						=      output.rgb, output.depth

         SCREAMING_SNAKE_CASE_     						=      0.419_4127
         SCREAMING_SNAKE_CASE_     						=      0.3537_5586
         SCREAMING_SNAKE_CASE_     						=      0.563_8502
         SCREAMING_SNAKE_CASE_     						=      0.3468_6103

         assert rgb.shape == (1, 512, 512, 3)
         assert depth.shape == (1, 512, 512, 1)
         assert np.abs(expected_rgb_mean - rgb.mean()						) < 1E-3
         assert np.abs(expected_rgb_std - rgb.std()						) < 1E-3
         assert np.abs(expected_depth_mean - depth.mean()						) < 1E-3
         assert np.abs(expected_depth_std - depth.std()						) < 1E-3





import argparse

import requests
import torch
from PIL import Image

from transformers import CLIPProcessor, GroupViTConfig, GroupViTModel

def  UpperCAmelCase_							(				__UpperCAmelCase						:     List[str]							) ->			List[Any]:
      # vision encoder
      if "img_encoder.pos_embed" in name:
            SCREAMING_SNAKE_CASE_     						=      name.replace('img_encoder.pos_embed'      ,      'vision_model.embeddings.position_embeddings'							)
      if "img_encoder.patch_embed.proj" in name:
            SCREAMING_SNAKE_CASE_     						=      name.replace('img_encoder.patch_embed.proj'      ,      'vision_model.embeddings.patch_embeddings.projection'							)
      if "img_encoder.patch_embed.norm" in name:
            SCREAMING_SNAKE_CASE_     						=      name.replace('img_encoder.patch_embed.norm'      ,      'vision_model.embeddings.layernorm'							)
      if "img_encoder.layers" in name:
            SCREAMING_SNAKE_CASE_     						=      name.replace('img_encoder.layers'      ,      'vision_model.encoder.stages'							)
      if "blocks" in name and "res" not in name:
            SCREAMING_SNAKE_CASE_     						=      name.replace('blocks'      ,      'layers'							)
      if "attn" in name and "pre_assign" not in name:
            SCREAMING_SNAKE_CASE_     						=      name.replace('attn'      ,      'self_attn'							)
      if "proj" in name and "self_attn" in name and "text" not in name:
            SCREAMING_SNAKE_CASE_     						=      name.replace('proj'      ,      'out_proj'							)
      if "pre_assign_attn.attn.proj" in name:
            SCREAMING_SNAKE_CASE_     						=      name.replace('pre_assign_attn.attn.proj'      ,      'pre_assign_attn.attn.out_proj'							)
      if "norm1" in name:
            SCREAMING_SNAKE_CASE_     						=      name.replace('norm1'      ,      'layer_norm1'							)
      if "norm2" in name and "pre_assign" not in name:
            SCREAMING_SNAKE_CASE_     						=      name.replace('norm2'      ,      'layer_norm2'							)
      if "img_encoder.norm" in name:
            SCREAMING_SNAKE_CASE_     						=      name.replace('img_encoder.norm'      ,      'vision_model.layernorm'							)
      # text encoder
      if "text_encoder.token_embedding" in name:
            SCREAMING_SNAKE_CASE_     						=      name.replace('text_encoder.token_embedding'      ,      'text_model.embeddings.token_embedding'							)
      if "text_encoder.positional_embedding" in name:
            SCREAMING_SNAKE_CASE_     						=      name.replace('text_encoder.positional_embedding'      ,      'text_model.embeddings.position_embedding.weight'							)
      if "text_encoder.transformer.resblocks." in name:
            SCREAMING_SNAKE_CASE_     						=      name.replace('text_encoder.transformer.resblocks.'      ,      'text_model.encoder.layers.'							)
      if "ln_1" in name:
            SCREAMING_SNAKE_CASE_     						=      name.replace('ln_1'      ,      'layer_norm1'							)
      if "ln_2" in name:
            SCREAMING_SNAKE_CASE_     						=      name.replace('ln_2'      ,      'layer_norm2'							)
      if "c_fc" in name:
            SCREAMING_SNAKE_CASE_     						=      name.replace('c_fc'      ,      'fc1'							)
      if "c_proj" in name:
            SCREAMING_SNAKE_CASE_     						=      name.replace('c_proj'      ,      'fc2'							)
      if "text_encoder" in name:
            SCREAMING_SNAKE_CASE_     						=      name.replace('text_encoder'      ,      'text_model'							)
      if "ln_final" in name:
            SCREAMING_SNAKE_CASE_     						=      name.replace('ln_final'      ,      'final_layer_norm'							)
      # projection layers
      if "img_projector.linear_hidden." in name:
            SCREAMING_SNAKE_CASE_     						=      name.replace('img_projector.linear_hidden.'      ,      'visual_projection.'							)
      if "img_projector.linear_out." in name:
            SCREAMING_SNAKE_CASE_     						=      name.replace('img_projector.linear_out.'      ,      'visual_projection.3.'							)
      if "text_projector.linear_hidden" in name:
            SCREAMING_SNAKE_CASE_     						=      name.replace('text_projector.linear_hidden'      ,      'text_projection'							)
      if "text_projector.linear_out" in name:
            SCREAMING_SNAKE_CASE_     						=      name.replace('text_projector.linear_out'      ,      'text_projection.3'							)

      return name

def  UpperCAmelCase_							(				__UpperCAmelCase						:     Union[str, Any]      ,      __UpperCAmelCase						:     int							) ->			List[str]:
      for key in orig_state_dict.copy().keys():
            SCREAMING_SNAKE_CASE_     						=      orig_state_dict.pop(__UpperCAmelCase							)

            if "qkv" in key:
                  # weights and biases of the key, value and query projections of vision encoder's attention layers require special treatment:
                  # we need to split them up into separate matrices/vectors
                  SCREAMING_SNAKE_CASE_     						=      key.split('.'							)
                  SCREAMING_SNAKE_CASE_						,			SCREAMING_SNAKE_CASE_     						=      int(key_split[2]							), int(key_split[4]							)
                  SCREAMING_SNAKE_CASE_     						=      config.vision_config.hidden_size
                  if "weight" in key:
                        SCREAMING_SNAKE_CASE_     						=      val[:dim, :]
                        SCREAMING_SNAKE_CASE_     						=      val[dim : dim * 2, :]
                        SCREAMING_SNAKE_CASE_     						=      val[-dim:, :]
                  else:
                        SCREAMING_SNAKE_CASE_     						=      val[:dim]
                        SCREAMING_SNAKE_CASE_     						=      val[dim : dim * 2]
                        SCREAMING_SNAKE_CASE_     						=      val[-dim:]
            elif "in_proj" in key:
                  # weights and biases of the key, value and query projections of text encoder's attention layers require special treatment:
                  # we need to split them up into separate matrices/vectors
                  SCREAMING_SNAKE_CASE_     						=      key.split('.'							)
                  SCREAMING_SNAKE_CASE_     						=      int(key_split[3]							)
                  SCREAMING_SNAKE_CASE_     						=      config.text_config.hidden_size
                  if "weight" in key:
                        SCREAMING_SNAKE_CASE_     						=      val[:dim, :]
                        SCREAMING_SNAKE_CASE_     						=      val[
                            dim : dim * 2, :
                        ]
                        SCREAMING_SNAKE_CASE_     						=      val[-dim:, :]
                  else:
                        SCREAMING_SNAKE_CASE_     						=      val[:dim]
                        SCREAMING_SNAKE_CASE_     						=      val[dim : dim * 2]
                        SCREAMING_SNAKE_CASE_     						=      val[-dim:]
            else:
                  SCREAMING_SNAKE_CASE_     						=      rename_key(__UpperCAmelCase							)
                  # squeeze if necessary
                  if (
                      "text_projection.0" in new_name
                      or "text_projection.3" in new_name
                      or "visual_projection.0" in new_name
                      or "visual_projection.3" in new_name
                  ):
                        SCREAMING_SNAKE_CASE_     						=      val.squeeze_()
                  else:
                        SCREAMING_SNAKE_CASE_     						=      val

      return orig_state_dict

def  UpperCAmelCase_							(				) ->			List[Any]:
      SCREAMING_SNAKE_CASE_     						=      'http://images.cocodataset.org/val2017/000000039769.jpg'
      SCREAMING_SNAKE_CASE_     						=      Image.open(requests.get(__UpperCAmelCase      ,      stream=__UpperCAmelCase							).raw							)
      return im

@torch.no_grad()
def  UpperCAmelCase_							(				__UpperCAmelCase						:     int      ,      __UpperCAmelCase						:     str      ,      __UpperCAmelCase						:     Dict="groupvit-gcc-yfcc"      ,      __UpperCAmelCase						:     Any=False							) ->			Optional[int]:
      SCREAMING_SNAKE_CASE_     						=      GroupViTConfig()
      SCREAMING_SNAKE_CASE_     						=      GroupViTModel(__UpperCAmelCase							).eval()

      SCREAMING_SNAKE_CASE_     						=      torch.load(__UpperCAmelCase      ,      map_location='cpu'							)['model']
      SCREAMING_SNAKE_CASE_     						=      convert_state_dict(__UpperCAmelCase      ,      __UpperCAmelCase							)
      SCREAMING_SNAKE_CASE_						,			SCREAMING_SNAKE_CASE_     						=      model.load_state_dict(__UpperCAmelCase      ,      strict=__UpperCAmelCase							)
      assert missing_keys == ["text_model.embeddings.position_ids"]
      assert (unexpected_keys == ["multi_label_logit_scale"]) or (len(__UpperCAmelCase							) == 0)

      # verify result
      SCREAMING_SNAKE_CASE_     						=      CLIPProcessor.from_pretrained('openai/clip-vit-base-patch32'							)
      SCREAMING_SNAKE_CASE_     						=      prepare_img()
      SCREAMING_SNAKE_CASE_     						=      processor(text=['a photo of a cat', 'a photo of a dog']      ,      images=__UpperCAmelCase      ,      padding=__UpperCAmelCase      ,      return_tensors='pt'							)

      with torch.no_grad():
            SCREAMING_SNAKE_CASE_     						=      model(**__UpperCAmelCase							)

      if model_name == "groupvit-gcc-yfcc":
            SCREAMING_SNAKE_CASE_     						=      torch.tensor([[1_3.3_5_2_3, 6.3_6_2_9]]							)
      elif model_name == "groupvit-gcc-redcaps":
            SCREAMING_SNAKE_CASE_     						=      torch.tensor([[1_6.1_8_7_3, 8.6_2_3_0]]							)
      else:
            raise ValueError(f"Model name {model_name} not supported."							)
      assert torch.allclose(outputs.logits_per_image      ,      __UpperCAmelCase      ,      atol=1E-3							)

      processor.save_pretrained(__UpperCAmelCase							)
      model.save_pretrained(__UpperCAmelCase							)
      print('Successfully saved processor and model to'      ,      __UpperCAmelCase							)

      if push_to_hub:
            print('Pushing to the hub...'							)
            processor.push_to_hub(__UpperCAmelCase      ,      organization='nielsr'							)
            model.push_to_hub(__UpperCAmelCase      ,      organization='nielsr'							)


if __name__ == "__main__":
   lowerCamelCase__			:       int         =							argparse.ArgumentParser()
   parser.add_argument(
       '--pytorch_dump_folder_path', default=None, type=str, help='Path to dump the processor and PyTorch model.'
   )
   parser.add_argument('--checkpoint_path', default=None, type=str, help='Path to GroupViT checkpoint')
   parser.add_argument(
       '--model_name',
       default='groupvit-gccy-fcc',
       type=str,
       help='Name of the model. Expecting either \'groupvit-gcc-yfcc\' or \'groupvit-gcc-redcaps\'',
   )
   parser.add_argument(
       '--push_to_hub',
       action='store_true',
       help='Whether or not to push the converted model and processor to the 🤗 hub using the provided `model_name`.',
   )
   lowerCamelCase__			:       List[str]         =							parser.parse_args()

   convert_groupvit_checkpoint(args.checkpoint_path, args.pytorch_dump_folder_path, args.model_name, args.push_to_hub)