



'''simple docstring'''





import argparse

import requests
import torch
from PIL import Image

from transformers import CLIPProcessor, GroupViTConfig, GroupViTModel




def lowerCAmelCase    (__A):



     """simple docstring"""



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

     return name




def lowerCAmelCase    (__A    ,      __A):



     """simple docstring"""



     for key in orig_state_dict.copy().keys():
          _a           =    orig_state_dict.pop(__A)

          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
               _a           =    key.split('''.''')
               _a    ,       _a           =    int(key_split[2]), int(key_split[4])
               _a           =    config.vision_config.hidden_size
               if "weight" in key:
                    _a           =    val[:dim, :]
                    _a           =    val[dim : dim * 2, :]
                    _a           =    val[-dim:, :]
               else:
                    _a           =    val[:dim]
                    _a           =    val[dim : dim * 2]
                    _a           =    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
               _a           =    key.split('''.''')
               _a           =    int(key_split[3])
               _a           =    config.text_config.hidden_size
               if "weight" in key:
                    _a           =    val[:dim, :]
                    _a           =    val[
                        dim : dim * 2, :
                    ]
                    _a           =    val[-dim:, :]
               else:
                    _a           =    val[:dim]
                    _a           =    val[dim : dim * 2]
                    _a           =    val[-dim:]
          else:
               _a           =    rename_key(__A)
               # 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
               ):
                    _a           =    val.squeeze_()
               else:
                    _a           =    val

     return orig_state_dict




def lowerCAmelCase    ():



     """simple docstring"""



     _a           =    '''http://images.cocodataset.org/val2017/000000039769.jpg'''
     _a           =    Image.open(requests.get(__A    ,      stream=__A).raw)
     return im




@torch.no_grad()
def lowerCAmelCase    (__A    ,      __A    ,      __A="groupvit-gcc-yfcc"    ,      __A=False):



     """simple docstring"""



     _a           =    GroupViTConfig()
     _a           =    GroupViTModel(__A).eval()

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

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

     with torch.no_grad():
          _a           =    model(**__A)

     if model_name == "groupvit-gcc-yfcc":
          _a           =    torch.tensor([[13.35_23, 6.36_29]])
     elif model_name == "groupvit-gcc-redcaps":
          _a           =    torch.tensor([[16.18_73, 8.62_30]])
     else:
          raise ValueError(F'''Model name {model_name} not supported.''')
     assert torch.allclose(outputs.logits_per_image    ,      __A    ,      atol=1e-3)

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

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


if __name__ == "__main__":
  lowercase_									=    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`.",
  )
  lowercase_									=    parser.parse_args()

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










"""simple docstring"""





import gc
import unittest

import torch
from transformers import CLIPTextConfig, CLIPTextModel, CLIPTextModelWithProjection, CLIPTokenizer

from diffusers import (
    AutoencoderKL,
    DDIMScheduler,
    DDPMScheduler,
    PriorTransformer,
    StableUnCLIPPipeline,
    UNetaDConditionModel,
)
from diffusers.pipelines.stable_diffusion.stable_unclip_image_normalizer import StableUnCLIPImageNormalizer
from diffusers.utils.testing_utils import enable_full_determinism, load_numpy, require_torch_gpu, slow, torch_device

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,
    assert_mean_pixel_difference,
)


enable_full_determinism()






class    _lowerCAmelCase  (	a    ,  a    ,  a    ,  unittest.TestCase			):
       """simple docstring"""




       __magic_name__	:int						    =   StableUnCLIPPipeline
       __magic_name__	:int						    =   TEXT_TO_IMAGE_PARAMS
       __magic_name__	:List[str]						    =   TEXT_TO_IMAGE_BATCH_PARAMS
       __magic_name__	:Optional[int]						    =   TEXT_TO_IMAGE_IMAGE_PARAMS
       __magic_name__	:Union[str, Any]						    =   TEXT_TO_IMAGE_IMAGE_PARAMS

       # TODO(will) Expected attn_bias.stride(1) == 0 to be true, but got false
       __magic_name__	:List[str]						    =   False






       def 	snake_case		(       self ):



         '''simple docstring'''

         lowerCAmelCase__     :Dict											=				3_2
         lowerCAmelCase__     :List[Any]											=				embedder_hidden_size

         # prior components

         torch.manual_seed(0 )
         lowerCAmelCase__     :Any											=				CLIPTokenizer.from_pretrained('hf-internal-testing/tiny-random-clip' )

         torch.manual_seed(0 )
         lowerCAmelCase__     :Any											=				CLIPTextModelWithProjection(
             CLIPTextConfig(
                 bos_token_id=0       ,					eos_token_id=2       ,					hidden_size=__UpperCAmelCase       ,					projection_dim=__UpperCAmelCase       ,					intermediate_size=3_7       ,					layer_norm_eps=1E-05       ,					num_attention_heads=4       ,					num_hidden_layers=5       ,					pad_token_id=1       ,					vocab_size=1_0_0_0       ,					) )

         torch.manual_seed(0 )
         lowerCAmelCase__     :int											=				PriorTransformer(
             num_attention_heads=2       ,					attention_head_dim=1_2       ,					embedding_dim=__UpperCAmelCase       ,					num_layers=1       ,					)

         torch.manual_seed(0 )
         lowerCAmelCase__     :Optional[int]											=				DDPMScheduler(
             variance_type='fixed_small_log'       ,					prediction_type='sample'       ,					num_train_timesteps=1_0_0_0       ,					clip_sample=__UpperCAmelCase       ,					clip_sample_range=5.0       ,					beta_schedule='squaredcos_cap_v2'       ,					)

         # regular denoising components

         torch.manual_seed(0 )
         lowerCAmelCase__     :Dict											=				StableUnCLIPImageNormalizer(embedding_dim=__UpperCAmelCase )
         lowerCAmelCase__     :List[Any]											=				DDPMScheduler(beta_schedule='squaredcos_cap_v2' )

         torch.manual_seed(0 )
         lowerCAmelCase__     :List[Any]											=				CLIPTokenizer.from_pretrained('hf-internal-testing/tiny-random-clip' )

         torch.manual_seed(0 )
         lowerCAmelCase__     :str											=				CLIPTextModel(
             CLIPTextConfig(
                 bos_token_id=0       ,					eos_token_id=2       ,					hidden_size=__UpperCAmelCase       ,					projection_dim=3_2       ,					intermediate_size=3_7       ,					layer_norm_eps=1E-05       ,					num_attention_heads=4       ,					num_hidden_layers=5       ,					pad_token_id=1       ,					vocab_size=1_0_0_0       ,					) )

         torch.manual_seed(0 )
         lowerCAmelCase__     :Optional[int]											=				UNetaDConditionModel(
             sample_size=3_2       ,					in_channels=4       ,					out_channels=4       ,					down_block_types=('CrossAttnDownBlock2D', 'DownBlock2D')       ,					up_block_types=('UpBlock2D', 'CrossAttnUpBlock2D')       ,					block_out_channels=(3_2, 6_4)       ,					attention_head_dim=(2, 4)       ,					class_embed_type='projection'       ,					projection_class_embeddings_input_dim=embedder_projection_dim * 2       ,					cross_attention_dim=__UpperCAmelCase       ,					layers_per_block=1       ,					upcast_attention=__UpperCAmelCase       ,					use_linear_projection=__UpperCAmelCase       ,					)

         torch.manual_seed(0 )
         lowerCAmelCase__     :Tuple											=				DDIMScheduler(
             beta_schedule='scaled_linear'       ,					beta_start=0.0_00_85       ,					beta_end=0.0_12       ,					prediction_type='v_prediction'       ,					set_alpha_to_one=__UpperCAmelCase       ,					steps_offset=1       ,					)

         torch.manual_seed(0 )
         lowerCAmelCase__     :Optional[int]											=				AutoencoderKL()

         lowerCAmelCase__     :Optional[int]											=				{
             # prior components
             'prior_tokenizer': prior_tokenizer,
             'prior_text_encoder': prior_text_encoder,
             'prior': prior,
             'prior_scheduler': prior_scheduler,
             # image noising components
             'image_normalizer': image_normalizer,
             'image_noising_scheduler': image_noising_scheduler,
             # regular denoising components
             'tokenizer': tokenizer,
             'text_encoder': text_encoder,
             'unet': unet,
             'scheduler': scheduler,
             'vae': vae,
         }

         return components






       def 	snake_case		(       self       ,					__UpperCAmelCase       ,					__UpperCAmelCase=0 ):



         '''simple docstring'''

         if str(__UpperCAmelCase ).startswith('mps' ):
           lowerCAmelCase__     :str											=				torch.manual_seed(__UpperCAmelCase )
         else:
           lowerCAmelCase__     :Optional[int]											=				torch.Generator(device=__UpperCAmelCase ).manual_seed(__UpperCAmelCase )
         lowerCAmelCase__     :Tuple											=				{
             'prompt': 'A painting of a squirrel eating a burger',
             'generator': generator,
             'num_inference_steps': 2,
             'prior_num_inference_steps': 2,
             'output_type': 'numpy',
         }
         return inputs






       def 	snake_case		(       self ):



         '''simple docstring'''

         lowerCAmelCase__     :Optional[int]											=				torch_device == 'cpu'

         self._test_attention_slicing_forward_pass(test_max_difference=__UpperCAmelCase )






       def 	snake_case		(       self ):



         '''simple docstring'''

         lowerCAmelCase__     :str											=				torch_device in ['cpu', 'mps']

         self._test_inference_batch_single_identical(test_max_difference=__UpperCAmelCase )






@slow
@require_torch_gpu
class    _lowerCAmelCase  (	unittest.TestCase			):
       """simple docstring"""




       def 	snake_case		(       self ):



         '''simple docstring'''

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






       def 	snake_case		(       self ):



         '''simple docstring'''

         lowerCAmelCase__     :Tuple											=				load_numpy(
             'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/stable_unclip/stable_unclip_2_1_l_anime_turtle_fp16.npy' )

         lowerCAmelCase__     :List[Any]											=				StableUnCLIPPipeline.from_pretrained('fusing/stable-unclip-2-1-l'       ,					torch_dtype=torch.floataa )
         pipe.to(__UpperCAmelCase )
         pipe.set_progress_bar_config(disable=__UpperCAmelCase )
         # stable unclip will oom when integration tests are run on a V100,
         # so turn on memory savings
         pipe.enable_attention_slicing()
         pipe.enable_sequential_cpu_offload()

         lowerCAmelCase__     :List[str]											=				torch.Generator(device='cpu' ).manual_seed(0 )
         lowerCAmelCase__     :List[str]											=				pipe('anime turle'       ,					generator=__UpperCAmelCase       ,					output_type='np' )

         lowerCAmelCase__     :Union[str, Any]											=				output.images[0]

         assert image.shape == (7_6_8, 7_6_8, 3)

         assert_mean_pixel_difference(__UpperCAmelCase       ,					__UpperCAmelCase )






       def 	snake_case		(       self ):



         '''simple docstring'''

         torch.cuda.empty_cache()
         torch.cuda.reset_max_memory_allocated()
         torch.cuda.reset_peak_memory_stats()

         lowerCAmelCase__     :int											=				StableUnCLIPPipeline.from_pretrained('fusing/stable-unclip-2-1-l'       ,					torch_dtype=torch.floataa )
         lowerCAmelCase__     :List[str]											=				pipe.to(__UpperCAmelCase )
         pipe.set_progress_bar_config(disable=__UpperCAmelCase )
         pipe.enable_attention_slicing()
         pipe.enable_sequential_cpu_offload()

         lowerCAmelCase__     :Tuple											=				pipe(
             'anime turtle'       ,					prior_num_inference_steps=2       ,					num_inference_steps=2       ,					output_type='np'       ,					)

         lowerCAmelCase__     :str											=				torch.cuda.max_memory_allocated()
         # make sure that less than 7 GB is allocated
         assert mem_bytes < 7 * 1_0**9
