






import gc
import random
import unittest

import numpy as np
import torch
from transformers import CLIPImageProcessor, CLIPVisionConfig, CLIPVisionModel

from diffusers import HeunDiscreteScheduler, PriorTransformer, ShapEImgaImgPipeline
from diffusers.pipelines.shap_e import ShapERenderer
from diffusers.utils import floats_tensor, load_image, load_numpy, slow
from diffusers.utils.testing_utils import require_torch_gpu, torch_device

from ..test_pipelines_common import PipelineTesterMixin, assert_mean_pixel_difference
class     lowerCamelCase_					(   lowerCAmelCase__      ,					unittest.TestCase			):
				a__       							=	ShapEImgaImgPipeline
				a__       							=	['''image''']
				a__       							=	['''image''']
				a__       							=	[
				    '''num_images_per_prompt''',
				    '''num_inference_steps''',
				    '''generator''',
				    '''latents''',
				    '''guidance_scale''',
				    '''frame_size''',
				    '''output_type''',
				    '''return_dict''',
				]
				a__       							=	False







				@property
				def       A						(    self  ):
								"""simple docstring"""
								return 3_2







				@property
				def       A						(    self  ):
								"""simple docstring"""
								return 3_2







				@property
				def       A						(    self  ):
								"""simple docstring"""
								return self.time_input_dim * 4







				@property
				def       A						(    self  ):
								"""simple docstring"""
								return 8







				@property
				def       A						(    self  ):
								"""simple docstring"""
								torch.manual_seed(0  )
								__magic_name__		:Union[str, Any]   =   CLIPVisionConfig(
								    hidden_size=self.text_embedder_hidden_size							,       image_size=6_4							,       projection_dim=self.text_embedder_hidden_size							,       intermediate_size=3_7							,       num_attention_heads=4							,       num_channels=3							,       num_hidden_layers=5							,       patch_size=1							,       )

								__magic_name__		:Union[str, Any]   =   CLIPVisionModel(__UpperCAmelCase  )
								return model







				@property
				def       A						(    self  ):
								"""simple docstring"""
								__magic_name__		:Dict   =   CLIPImageProcessor(
								    crop_size=2_2_4							,       do_center_crop=__UpperCAmelCase							,       do_normalize=__UpperCAmelCase							,       do_resize=__UpperCAmelCase							,       image_mean=[0.48145466, 0.4578275, 0.40821073]							,       image_std=[0.26862954, 0.26130258, 0.27577711]							,       resample=3							,       size=2_2_4							,       )

								return image_processor







				@property
				def       A						(    self  ):
								"""simple docstring"""
								torch.manual_seed(0  )

								__magic_name__		:Optional[int]   =   {
								    '''num_attention_heads''': 2,
								    '''attention_head_dim''': 1_6,
								    '''embedding_dim''': self.time_input_dim,
								    '''num_embeddings''': 3_2,
								    '''embedding_proj_dim''': self.text_embedder_hidden_size,
								    '''time_embed_dim''': self.time_embed_dim,
								    '''num_layers''': 1,
								    '''clip_embed_dim''': self.time_input_dim * 2,
								    '''additional_embeddings''': 0,
								    '''time_embed_act_fn''': '''gelu''',
								    '''norm_in_type''': '''layer''',
								    '''embedding_proj_norm_type''': '''layer''',
								    '''encoder_hid_proj_type''': None,
								    '''added_emb_type''': None,
								}

								__magic_name__		:str   =   PriorTransformer(**__UpperCAmelCase  )
								return model







				@property
				def       A						(    self  ):
								"""simple docstring"""
								torch.manual_seed(0  )

								__magic_name__		:int   =   {
								    '''param_shapes''': (
								        (self.renderer_dim, 9_3),
								        (self.renderer_dim, 8),
								        (self.renderer_dim, 8),
								        (self.renderer_dim, 8),
								    ),
								    '''d_latent''': self.time_input_dim,
								    '''d_hidden''': self.renderer_dim,
								    '''n_output''': 1_2,
								    '''background''': (
								        0.1,
								        0.1,
								        0.1,
								    ),
								}
								__magic_name__		:List[str]   =   ShapERenderer(**__UpperCAmelCase  )
								return model







				def       A						(    self  ):
								"""simple docstring"""
								__magic_name__		:Optional[int]   =   self.dummy_prior
								__magic_name__		:List[Any]   =   self.dummy_image_encoder
								__magic_name__		:str   =   self.dummy_image_processor
								__magic_name__		:Dict   =   self.dummy_renderer

								__magic_name__		:int   =   HeunDiscreteScheduler(
								    beta_schedule='''exp'''							,       num_train_timesteps=1_0_2_4							,       prediction_type='''sample'''							,       use_karras_sigmas=__UpperCAmelCase							,       clip_sample=__UpperCAmelCase							,       clip_sample_range=1.0							,       )
								__magic_name__		:Dict   =   {
								    '''prior''': prior,
								    '''image_encoder''': image_encoder,
								    '''image_processor''': image_processor,
								    '''renderer''': renderer,
								    '''scheduler''': scheduler,
								}

								return components







				def       A						(    self							,       __lowerCAmelCase							,       __lowerCAmelCase=0  ):
								"""simple docstring"""
								__magic_name__		:List[Any]   =   floats_tensor((1, 3, 6_4, 6_4)							,       rng=random.Random(__UpperCAmelCase  )  ).to(__UpperCAmelCase  )

								if str(__UpperCAmelCase  ).startswith('''mps'''  ):
												__magic_name__		:List[Any]   =   torch.manual_seed(__UpperCAmelCase  )
								else:
												__magic_name__		:Tuple   =   torch.Generator(device=__UpperCAmelCase  ).manual_seed(__UpperCAmelCase  )
								__magic_name__		:List[str]   =   {
								    '''image''': input_image,
								    '''generator''': generator,
								    '''num_inference_steps''': 1,
								    '''frame_size''': 3_2,
								    '''output_type''': '''np''',
								}
								return inputs







				def       A						(    self  ):
								"""simple docstring"""
								__magic_name__		:List[Any]   =   '''cpu'''

								__magic_name__		:Dict   =   self.get_dummy_components()

								__magic_name__		:int   =   self.pipeline_class(**__UpperCAmelCase  )
								__magic_name__		:Union[str, Any]   =   pipe.to(__UpperCAmelCase  )

								pipe.set_progress_bar_config(disable=__UpperCAmelCase  )

								__magic_name__		:Any   =   pipe(**self.get_dummy_inputs(__UpperCAmelCase  )  )
								__magic_name__		:str   =   output.images[0]
								__magic_name__		:Union[str, Any]   =   image[0, -3:, -3:, -1]

								assert image.shape == (2_0, 3_2, 3_2, 3)

								__magic_name__		:List[str]   =   np.array(
								    [
								        0.00039216,
								        0.00039216,
								        0.00039216,
								        0.00039216,
								        0.00039216,
								        0.00039216,
								        0.00039216,
								        0.00039216,
								        0.00039216,
								    ]  )

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







				def       A						(    self  ):
								"""simple docstring"""
								self._test_inference_batch_consistent(batch_sizes=[1, 2]  )







				def       A						(    self  ):
								"""simple docstring"""
								__magic_name__		:str   =   torch_device == '''cpu'''
								__magic_name__		:Union[str, Any]   =   True
								self._test_inference_batch_single_identical(
								    batch_size=2							,       test_max_difference=__UpperCAmelCase							,       relax_max_difference=__UpperCAmelCase							,       )







				def       A						(    self  ):
								"""simple docstring"""
								__magic_name__		:Dict   =   self.get_dummy_components()
								__magic_name__		:Dict   =   self.pipeline_class(**__UpperCAmelCase  )
								__magic_name__		:Any   =   pipe.to(__UpperCAmelCase  )
								pipe.set_progress_bar_config(disable=__UpperCAmelCase  )

								__magic_name__		:Tuple   =   1
								__magic_name__		:str   =   2

								__magic_name__		:str   =   self.get_dummy_inputs(__UpperCAmelCase  )

								for key in inputs.keys():
												if key in self.batch_params:
																__magic_name__		:Tuple   =   batch_size * [inputs[key]]

								__magic_name__		:Tuple   =   pipe(**__UpperCAmelCase							,       num_images_per_prompt=__UpperCAmelCase  )[0]

								assert images.shape[0] == batch_size * num_images_per_prompt







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







				def       A						(    self  ):
								"""simple docstring"""
								super().tearDown()
								gc.collect()
								torch.cuda.empty_cache()







				def       A						(    self  ):
								"""simple docstring"""
								__magic_name__		:Tuple   =   load_image(
								    '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main''' '''/shap_e/corgi.png'''  )
								__magic_name__		:Dict   =   load_numpy(
								    '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'''
								    '''/shap_e/test_shap_e_img2img_out.npy'''  )
								__magic_name__		:Any   =   ShapEImgaImgPipeline.from_pretrained('''openai/shap-e-img2img'''  )
								__magic_name__		:Tuple   =   pipe.to(__UpperCAmelCase  )
								pipe.set_progress_bar_config(disable=__UpperCAmelCase  )

								__magic_name__		:Dict   =   torch.Generator(device=__UpperCAmelCase  ).manual_seed(0  )

								__magic_name__		:List[Any]   =   pipe(
								    __UpperCAmelCase							,       generator=__UpperCAmelCase							,       guidance_scale=3.0							,       num_inference_steps=6_4							,       frame_size=6_4							,       output_type='''np'''							,       ).images[0]

								assert images.shape == (2_0, 6_4, 6_4, 3)

								assert_mean_pixel_difference(__UpperCAmelCase							,       __UpperCAmelCase  )











import argparse

import requests
import torch
from PIL import Image
from torchvision.transforms import Compose, Normalize, Resize, ToTensor

from transformers import SwinaSRConfig, SwinaSRForImageSuperResolution, SwinaSRImageProcessor





def 		a__      (		_UpperCamelCase      :			List[str]				):
   __lowerCamelCase    				=     SwinaSRConfig()

   if "Swin2SR_ClassicalSR_X4_64" in checkpoint_url:
      __lowerCamelCase    				=     4
   elif "Swin2SR_CompressedSR_X4_48" in checkpoint_url:
      __lowerCamelCase    				=     4
      __lowerCamelCase    				=     48
      __lowerCamelCase    				=     '''pixelshuffle_aux'''
   elif "Swin2SR_Lightweight_X2_64" in checkpoint_url:
      __lowerCamelCase    				=     [6, 6, 6, 6]
      __lowerCamelCase    				=     60
      __lowerCamelCase    				=     [6, 6, 6, 6]
      __lowerCamelCase    				=     '''pixelshuffledirect'''
   elif "Swin2SR_RealworldSR_X4_64_BSRGAN_PSNR" in checkpoint_url:
      __lowerCamelCase    				=     4
      __lowerCamelCase    				=     '''nearest+conv'''
   elif "Swin2SR_Jpeg_dynamic" in checkpoint_url:
      __lowerCamelCase    				=     1
      __lowerCamelCase    				=     1
      __lowerCamelCase    				=     1_26
      __lowerCamelCase    				=     7
      __lowerCamelCase    				=     255.0
      __lowerCamelCase    				=     ''''''

   return config





def 		a__      (		_UpperCamelCase      :			List[Any]			,_UpperCamelCase      :			List[str]				):
   if "patch_embed.proj" in name and "layers" not in name:
      __lowerCamelCase    				=     name.replace('''patch_embed.proj'''			,'''embeddings.patch_embeddings.projection'''				)
   if "patch_embed.norm" in name:
      __lowerCamelCase    				=     name.replace('''patch_embed.norm'''			,'''embeddings.patch_embeddings.layernorm'''				)
   if "layers" in name:
      __lowerCamelCase    				=     name.replace('''layers'''			,'''encoder.stages'''				)
   if "residual_group.blocks" in name:
      __lowerCamelCase    				=     name.replace('''residual_group.blocks'''			,'''layers'''				)
   if "attn.proj" in name:
      __lowerCamelCase    				=     name.replace('''attn.proj'''			,'''attention.output.dense'''				)
   if "attn" in name:
      __lowerCamelCase    				=     name.replace('''attn'''			,'''attention.self'''				)
   if "norm1" in name:
      __lowerCamelCase    				=     name.replace('''norm1'''			,'''layernorm_before'''				)
   if "norm2" in name:
      __lowerCamelCase    				=     name.replace('''norm2'''			,'''layernorm_after'''				)
   if "mlp.fc1" in name:
      __lowerCamelCase    				=     name.replace('''mlp.fc1'''			,'''intermediate.dense'''				)
   if "mlp.fc2" in name:
      __lowerCamelCase    				=     name.replace('''mlp.fc2'''			,'''output.dense'''				)
   if "q_bias" in name:
      __lowerCamelCase    				=     name.replace('''q_bias'''			,'''query.bias'''				)
   if "k_bias" in name:
      __lowerCamelCase    				=     name.replace('''k_bias'''			,'''key.bias'''				)
   if "v_bias" in name:
      __lowerCamelCase    				=     name.replace('''v_bias'''			,'''value.bias'''				)
   if "cpb_mlp" in name:
      __lowerCamelCase    				=     name.replace('''cpb_mlp'''			,'''continuous_position_bias_mlp'''				)
   if "patch_embed.proj" in name:
      __lowerCamelCase    				=     name.replace('''patch_embed.proj'''			,'''patch_embed.projection'''				)

   if name == "norm.weight":
      __lowerCamelCase    				=     '''layernorm.weight'''
   if name == "norm.bias":
      __lowerCamelCase    				=     '''layernorm.bias'''

   if "conv_first" in name:
      __lowerCamelCase    				=     name.replace('''conv_first'''			,'''first_convolution'''				)

   if (
       "upsample" in name
       or "conv_before_upsample" in name
       or "conv_bicubic" in name
       or "conv_up" in name
       or "conv_hr" in name
       or "conv_last" in name
       or "aux" in name
   ):
      # heads
      if "conv_last" in name:
         __lowerCamelCase    				=     name.replace('''conv_last'''			,'''final_convolution'''				)
      if config.upsampler in ["pixelshuffle", "pixelshuffle_aux", "nearest+conv"]:
         if "conv_before_upsample.0" in name:
            __lowerCamelCase    				=     name.replace('''conv_before_upsample.0'''			,'''conv_before_upsample'''				)
         if "upsample.0" in name:
            __lowerCamelCase    				=     name.replace('''upsample.0'''			,'''upsample.convolution_0'''				)
         if "upsample.2" in name:
            __lowerCamelCase    				=     name.replace('''upsample.2'''			,'''upsample.convolution_1'''				)
         __lowerCamelCase    				=     '''upsample.''' + name
      elif config.upsampler == "pixelshuffledirect":
         __lowerCamelCase    				=     name.replace('''upsample.0.weight'''			,'''upsample.conv.weight'''				)
         __lowerCamelCase    				=     name.replace('''upsample.0.bias'''			,'''upsample.conv.bias'''				)
      else:
         pass
   else:
      __lowerCamelCase    				=     '''swin2sr.''' + name

   return name





def 		a__      (		_UpperCamelCase      :			List[Any]			,_UpperCamelCase      :			Union[str, Any]				):
   for key in orig_state_dict.copy().keys():
      __lowerCamelCase    				=     orig_state_dict.pop(_UpperCamelCase				)

      if "qkv" in key:
         __lowerCamelCase    				=     key.split('''.'''				)
         __lowerCamelCase    				=     int(key_split[1]				)
         __lowerCamelCase    				=     int(key_split[4]				)
         __lowerCamelCase    				=     config.embed_dim

         if "weight" in key:
            __lowerCamelCase    				=     val[:dim, :]
            __lowerCamelCase    				=     val[dim : dim * 2, :]
            __lowerCamelCase    				=     val[-dim:, :]
         else:
            __lowerCamelCase    				=     val[:dim]
            __lowerCamelCase    				=     val[dim : dim * 2]
            __lowerCamelCase    				=     val[-dim:]
         pass
      else:
         __lowerCamelCase    				=     val

   return orig_state_dict





def 		a__      (		_UpperCamelCase      :			str			,_UpperCamelCase      :			int			,_UpperCamelCase      :			Any				):
   __lowerCamelCase    				=     get_config(_UpperCamelCase				)
   __lowerCamelCase    				=     SwinaSRForImageSuperResolution(_UpperCamelCase				)
   model.eval()

   __lowerCamelCase    				=     torch.hub.load_state_dict_from_url(_UpperCamelCase			,map_location='''cpu'''				)
   __lowerCamelCase    				=     convert_state_dict(_UpperCamelCase			,_UpperCamelCase				)
   __lowerCamelCase				,__lowerCamelCase    				=     model.load_state_dict(_UpperCamelCase			,strict=_UpperCamelCase				)

   if len(_UpperCamelCase				) > 0:
      raise ValueError('''Missing keys when converting: {}'''.format(_UpperCamelCase				)				)
   for key in unexpected_keys:
      if not ("relative_position_index" in key or "relative_coords_table" in key or "self_mask" in key):
         raise ValueError(F"""Unexpected key {key} in state_dict"""				)

    # verify values
   __lowerCamelCase    				=     '''https://github.com/mv-lab/swin2sr/blob/main/testsets/real-inputs/shanghai.jpg?raw=true'''
   __lowerCamelCase    				=     Image.open(requests.get(_UpperCamelCase			,stream=_UpperCamelCase				).raw				).convert('''RGB'''				)
   __lowerCamelCase    				=     SwinaSRImageProcessor()
   # pixel_values = processor(image, return_tensors="pt").pixel_values

   __lowerCamelCase    				=     1_26 if '''Jpeg''' in checkpoint_url else 2_56
   __lowerCamelCase    				=     Compose(
       [
           Resize((image_size, image_size)				),
           ToTensor(),
           Normalize(mean=[0.485, 0.456, 0.406]			,std=[0.229, 0.224, 0.225]				),
       ]				)
   __lowerCamelCase    				=     transforms(_UpperCamelCase				).unsqueeze(0				)

   if config.num_channels == 1:
      __lowerCamelCase    				=     pixel_values[:, 0, :, :].unsqueeze(1				)

   __lowerCamelCase    				=     model(_UpperCamelCase				)

   # assert values
   if "Swin2SR_ClassicalSR_X2_64" in checkpoint_url:
      __lowerCamelCase    				=     torch.Size([1, 3, 5_12, 5_12]				)
      __lowerCamelCase    				=     torch.tensor(
          [[-0.7_087, -0.7_138, -0.6_721], [-0.8_340, -0.8_095, -0.7_298], [-0.9_149, -0.8_414, -0.7_940]]				)
   elif "Swin2SR_ClassicalSR_X4_64" in checkpoint_url:
      __lowerCamelCase    				=     torch.Size([1, 3, 10_24, 10_24]				)
      __lowerCamelCase    				=     torch.tensor(
          [[-0.7_775, -0.8_105, -0.8_933], [-0.7_764, -0.8_356, -0.9_225], [-0.7_976, -0.8_686, -0.9_579]]				)
   elif "Swin2SR_CompressedSR_X4_48" in checkpoint_url:
      # TODO values didn't match exactly here
      __lowerCamelCase    				=     torch.Size([1, 3, 10_24, 10_24]				)
      __lowerCamelCase    				=     torch.tensor(
          [[-0.8_035, -0.7_504, -0.7_491], [-0.8_538, -0.8_124, -0.7_782], [-0.8_804, -0.8_651, -0.8_493]]				)
   elif "Swin2SR_Lightweight_X2_64" in checkpoint_url:
      __lowerCamelCase    				=     torch.Size([1, 3, 5_12, 5_12]				)
      __lowerCamelCase    				=     torch.tensor(
          [[-0.7_669, -0.8_662, -0.8_767], [-0.8_810, -0.9_962, -0.9_820], [-0.9_340, -1.0_322, -1.1_149]]				)
   elif "Swin2SR_RealworldSR_X4_64_BSRGAN_PSNR" in checkpoint_url:
      __lowerCamelCase    				=     torch.Size([1, 3, 10_24, 10_24]				)
      __lowerCamelCase    				=     torch.tensor(
          [[-0.5_238, -0.5_557, -0.6_321], [-0.6_016, -0.5_903, -0.6_391], [-0.6_244, -0.6_334, -0.6_889]]				)

   assert (
       outputs.reconstruction.shape == expected_shape
   ), F"""Shape of reconstruction should be {expected_shape}, but is {outputs.reconstruction.shape}"""
   assert torch.allclose(outputs.reconstruction[0, 0, :3, :3]			,_UpperCamelCase			,atol=1e-3				)
   print('''Looks ok!'''				)

   __lowerCamelCase    				=     {
       '''https://github.com/mv-lab/swin2sr/releases/download/v0.0.1/Swin2SR_ClassicalSR_X2_64.pth''': (
           '''swin2SR-classical-sr-x2-64'''
       ),
       '''https://github.com/mv-lab/swin2sr/releases/download/v0.0.1/Swin2SR_ClassicalSR_X4_64.pth''': (
           '''swin2SR-classical-sr-x4-64'''
       ),
       '''https://github.com/mv-lab/swin2sr/releases/download/v0.0.1/Swin2SR_CompressedSR_X4_48.pth''': (
           '''swin2SR-compressed-sr-x4-48'''
       ),
       '''https://github.com/mv-lab/swin2sr/releases/download/v0.0.1/Swin2SR_Lightweight_X2_64.pth''': (
           '''swin2SR-lightweight-x2-64'''
       ),
       '''https://github.com/mv-lab/swin2sr/releases/download/v0.0.1/Swin2SR_RealworldSR_X4_64_BSRGAN_PSNR.pth''': (
           '''swin2SR-realworld-sr-x4-64-bsrgan-psnr'''
       ),
   }
   __lowerCamelCase    				=     url_to_name[checkpoint_url]

   if pytorch_dump_folder_path is not None:
      print(F"""Saving model {model_name} to {pytorch_dump_folder_path}"""				)
      model.save_pretrained(_UpperCamelCase				)
      print(F"""Saving image processor to {pytorch_dump_folder_path}"""				)
      processor.save_pretrained(_UpperCamelCase				)

   if push_to_hub:
      model.push_to_hub(F"""caidas/{model_name}"""				)
      processor.push_to_hub(F"""caidas/{model_name}"""				)


if __name__ == "__main__":
  a_								=       argparse.ArgumentParser()
  # Required parameters
  parser.add_argument(
      """--checkpoint_url""",
      default="""https://github.com/mv-lab/swin2sr/releases/download/v0.0.1/Swin2SR_ClassicalSR_X2_64.pth""",
      type=str,
      help="""URL of the original Swin2SR checkpoint you'd like to convert.""",
  )
  parser.add_argument(
      """--pytorch_dump_folder_path""", default=None, type=str, help="""Path to the output PyTorch model directory."""
  )
  parser.add_argument("""--push_to_hub""", action="""store_true""", help="""Whether to push the converted model to the hub.""")

  a_								=       parser.parse_args()
  convert_swinasr_checkpoint(args.checkpoint_url, args.pytorch_dump_folder_path, args.push_to_hub)






