





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_   (			lowerCAmelCase_				:      Union[str, Any]     ):
   __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							     =     126
      __lowerCAmelCase							     =     7
      __lowerCAmelCase							     =     255.0
      __lowerCAmelCase							     =     ''

   return config

def   a_   (			lowerCAmelCase_				:      Any,     lowerCAmelCase_				:      Any     ):
   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_   (			lowerCAmelCase_				:      List[Any],     lowerCAmelCase_				:      Optional[Any]     ):
   for key in orig_state_dict.copy().keys():
      __lowerCAmelCase							     =     orig_state_dict.pop(lowerCAmelCase_     )

      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_   (			lowerCAmelCase_				:      str,     lowerCAmelCase_				:      List[Any],     lowerCAmelCase_				:      int     ):
   __lowerCAmelCase							     =     get_config(lowerCAmelCase_     )
   __lowerCAmelCase							     =     SwinaSRForImageSuperResolution(lowerCAmelCase_     )
   model.eval()

   __lowerCAmelCase							     =     torch.hub.load_state_dict_from_url(lowerCAmelCase_,     map_location='cpu'     )
   __lowerCAmelCase							     =     convert_state_dict(lowerCAmelCase_,     lowerCAmelCase_     )
   __lowerCAmelCase     ,     __lowerCAmelCase							     =     model.load_state_dict(lowerCAmelCase_,     strict=lowerCAmelCase_     )

   if len(lowerCAmelCase_     ) > 0:
      raise ValueError('Missing keys when converting: {}'.format(lowerCAmelCase_     )     )
   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(lowerCAmelCase_,     stream=lowerCAmelCase_     ).raw     ).convert('RGB'     )
   __lowerCAmelCase							     =     SwinaSRImageProcessor()
   # pixel_values = processor(image, return_tensors="pt").pixel_values

   __lowerCAmelCase							     =     126 if 'Jpeg' in checkpoint_url else 256
   __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(lowerCAmelCase_     ).unsqueeze(0     )

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

   __lowerCAmelCase							     =     model(lowerCAmelCase_     )

   # assert values
   if "Swin2SR_ClassicalSR_X2_64" in checkpoint_url:
      __lowerCAmelCase							     =     torch.Size([1, 3, 512, 512]     )
      __lowerCAmelCase							     =     torch.tensor(
          [[-0.7087, -0.7138, -0.6721], [-0.8340, -0.8095, -0.7298], [-0.9149, -0.8414, -0.7940]]     )
   elif "Swin2SR_ClassicalSR_X4_64" in checkpoint_url:
      __lowerCAmelCase							     =     torch.Size([1, 3, 1024, 1024]     )
      __lowerCAmelCase							     =     torch.tensor(
          [[-0.7775, -0.8105, -0.8933], [-0.7764, -0.8356, -0.9225], [-0.7976, -0.8686, -0.9579]]     )
   elif "Swin2SR_CompressedSR_X4_48" in checkpoint_url:
      # TODO values didn't match exactly here
      __lowerCAmelCase							     =     torch.Size([1, 3, 1024, 1024]     )
      __lowerCAmelCase							     =     torch.tensor(
          [[-0.8035, -0.7504, -0.7491], [-0.8538, -0.8124, -0.7782], [-0.8804, -0.8651, -0.8493]]     )
   elif "Swin2SR_Lightweight_X2_64" in checkpoint_url:
      __lowerCAmelCase							     =     torch.Size([1, 3, 512, 512]     )
      __lowerCAmelCase							     =     torch.tensor(
          [[-0.7669, -0.8662, -0.8767], [-0.8810, -0.9962, -0.9820], [-0.9340, -1.0322, -1.1149]]     )
   elif "Swin2SR_RealworldSR_X4_64_BSRGAN_PSNR" in checkpoint_url:
      __lowerCAmelCase							     =     torch.Size([1, 3, 1024, 1024]     )
      __lowerCAmelCase							     =     torch.tensor(
          [[-0.5238, -0.5557, -0.6321], [-0.6016, -0.5903, -0.6391], [-0.6244, -0.6334, -0.6889]]     )

   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],     lowerCAmelCase_,     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(lowerCAmelCase_     )
      print(F"""Saving image processor to {pytorch_dump_folder_path}"""     )
      processor.save_pretrained(lowerCAmelCase_     )

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


if __name__ == "__main__":
     _snake_case  :   Dict						   =  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.')

     _snake_case  :   Any						   =  parser.parse_args()
     convert_swinasr_checkpoint(args.checkpoint_url, args.pytorch_dump_folder_path, args.push_to_hub)








import gc
import unittest

import numpy as np
import torch

from diffusers import (
    AudioDiffusionPipeline,
    AutoencoderKL,
    DDIMScheduler,
    DDPMScheduler,
    DiffusionPipeline,
    Mel,
    UNetaDConditionModel,
    UNetaDModel,
)
from diffusers.utils import slow, torch_device
from diffusers.utils.testing_utils import enable_full_determinism, require_torch_gpu


enable_full_determinism()



class     _UpperCAmelCase   (   unittest.TestCase   ):

  """simple docstring"""


  def     lowercase		(    self			:		List[str]					)							->						int:
     # clean up the VRAM after each test
     super().tearDown()
     gc.collect()
     torch.cuda.empty_cache()






  @property
  def     lowercase		(    self			:		Union[str, Any]					)							->						Any:
     torch.manual_seed(0					)
     __lowerCAmelCase							     =     UNetaDModel(
         sample_size=(3_2, 6_4) ,  in_channels=1 ,  out_channels=1 ,  layers_per_block=2 ,  block_out_channels=(1_2_8, 1_2_8) ,  down_block_types=('AttnDownBlock2D', 'DownBlock2D') ,  up_block_types=('UpBlock2D', 'AttnUpBlock2D') ,  )
     return model






  @property
  def     lowercase		(    self			:		Optional[Any]					)							->						List[Any]:
     torch.manual_seed(0					)
     __lowerCAmelCase							     =     UNetaDConditionModel(
         sample_size=(6_4, 3_2) ,  in_channels=1 ,  out_channels=1 ,  layers_per_block=2 ,  block_out_channels=(1_2_8, 1_2_8) ,  down_block_types=('CrossAttnDownBlock2D', 'DownBlock2D') ,  up_block_types=('UpBlock2D', 'CrossAttnUpBlock2D') ,  cross_attention_dim=1_0 ,  )
     return model






  @property
  def     lowercase		(    self			:		Dict					)							->						Optional[Any]:
     torch.manual_seed(0					)
     __lowerCAmelCase							     =     AutoencoderKL(
         sample_size=(1_2_8, 6_4) ,  in_channels=1 ,  out_channels=1 ,  latent_channels=1 ,  layers_per_block=2 ,  block_out_channels=(1_2_8, 1_2_8) ,  down_block_types=('DownEncoderBlock2D', 'DownEncoderBlock2D') ,  up_block_types=('UpDecoderBlock2D', 'UpDecoderBlock2D') ,  )
     __lowerCAmelCase							     =     UNetaDModel(
         sample_size=(6_4, 3_2) ,  in_channels=1 ,  out_channels=1 ,  layers_per_block=2 ,  block_out_channels=(1_2_8, 1_2_8) ,  down_block_types=('AttnDownBlock2D', 'DownBlock2D') ,  up_block_types=('UpBlock2D', 'AttnUpBlock2D') ,  )
     return vqvae, unet






  @slow
  def     lowercase		(    self			:		Dict					)							->						Optional[Any]:
     __lowerCAmelCase							     =     'cpu'  # ensure determinism for the device-dependent torch.Generator
     __lowerCAmelCase							     =     Mel(
         x_res=self.dummy_unet.config.sample_size[1] ,  y_res=self.dummy_unet.config.sample_size[0] ,  )

     __lowerCAmelCase							     =     DDPMScheduler()
     __lowerCAmelCase							     =     AudioDiffusionPipeline(vqvae=lowerCAmelCase_ ,  unet=self.dummy_unet ,  mel=lowerCAmelCase_ ,  scheduler=lowerCAmelCase_					)
     __lowerCAmelCase							     =     pipe.to(lowerCAmelCase_					)
     pipe.set_progress_bar_config(disable=lowerCAmelCase_					)

     __lowerCAmelCase							     =     torch.Generator(device=lowerCAmelCase_					).manual_seed(4_2					)
     __lowerCAmelCase							     =     pipe(generator=lowerCAmelCase_ ,  steps=4					)
     __lowerCAmelCase							     =     output.audios[0]
     __lowerCAmelCase							     =     output.images[0]

     __lowerCAmelCase							     =     torch.Generator(device=lowerCAmelCase_					).manual_seed(4_2					)
     __lowerCAmelCase							     =     pipe(generator=lowerCAmelCase_ ,  steps=4 ,  return_dict=lowerCAmelCase_					)
     __lowerCAmelCase							     =     output[0][0]

     assert audio.shape == (1, (self.dummy_unet.config.sample_size[1] - 1) * mel.hop_length)
     assert (
         image.height == self.dummy_unet.config.sample_size[0]
         and image.width == self.dummy_unet.config.sample_size[1]
     )
     __lowerCAmelCase							     =     np.frombuffer(image.tobytes() ,  dtype='uint8'					)[:1_0]
     __lowerCAmelCase							     =     np.frombuffer(image_from_tuple.tobytes() ,  dtype='uint8'					)[:1_0]
     __lowerCAmelCase							     =     np.array([6_9, 2_5_5, 2_5_5, 2_5_5, 0, 0, 7_7, 1_8_1, 1_2, 1_2_7]					)

     assert np.abs(image_slice.flatten() - expected_slice					).max() == 0
     assert np.abs(image_from_tuple_slice.flatten() - expected_slice					).max() == 0

     __lowerCAmelCase							     =     Mel(
         x_res=self.dummy_vqvae_and_unet[0].config.sample_size[1] ,  y_res=self.dummy_vqvae_and_unet[0].config.sample_size[0] ,  )

     __lowerCAmelCase							     =     DDIMScheduler()
     __lowerCAmelCase							     =     self.dummy_vqvae_and_unet
     __lowerCAmelCase							     =     AudioDiffusionPipeline(
         vqvae=self.dummy_vqvae_and_unet[0] ,  unet=dummy_vqvae_and_unet[1] ,  mel=lowerCAmelCase_ ,  scheduler=lowerCAmelCase_					)
     __lowerCAmelCase							     =     pipe.to(lowerCAmelCase_					)
     pipe.set_progress_bar_config(disable=lowerCAmelCase_					)

     np.random.seed(0					)
     __lowerCAmelCase							     =     np.random.uniform(-1 ,  1 ,  ((dummy_vqvae_and_unet[0].config.sample_size[1] - 1) * mel.hop_length,)					)
     __lowerCAmelCase							     =     torch.Generator(device=lowerCAmelCase_					).manual_seed(4_2					)
     __lowerCAmelCase							     =     pipe(raw_audio=lowerCAmelCase_ ,  generator=lowerCAmelCase_ ,  start_step=5 ,  steps=1_0					)
     __lowerCAmelCase							     =     output.images[0]

     assert (
         image.height == self.dummy_vqvae_and_unet[0].config.sample_size[0]
         and image.width == self.dummy_vqvae_and_unet[0].config.sample_size[1]
     )
     __lowerCAmelCase							     =     np.frombuffer(image.tobytes() ,  dtype='uint8'					)[:1_0]
     __lowerCAmelCase							     =     np.array([1_2_0, 1_1_7, 1_1_0, 1_0_9, 1_3_8, 1_6_7, 1_3_8, 1_4_8, 1_3_2, 1_2_1]					)

     assert np.abs(image_slice.flatten() - expected_slice					).max() == 0

     __lowerCAmelCase							     =     self.dummy_unet_condition
     __lowerCAmelCase							     =     AudioDiffusionPipeline(
         vqvae=self.dummy_vqvae_and_unet[0] ,  unet=lowerCAmelCase_ ,  mel=lowerCAmelCase_ ,  scheduler=lowerCAmelCase_					)
     __lowerCAmelCase							     =     pipe.to(lowerCAmelCase_					)
     pipe.set_progress_bar_config(disable=lowerCAmelCase_					)

     np.random.seed(0					)
     __lowerCAmelCase							     =     torch.rand((1, 1, 1_0)					)
     __lowerCAmelCase							     =     pipe(generator=lowerCAmelCase_ ,  encoding=lowerCAmelCase_					)
     __lowerCAmelCase							     =     output.images[0]
     __lowerCAmelCase							     =     np.frombuffer(image.tobytes() ,  dtype='uint8'					)[:1_0]
     __lowerCAmelCase							     =     np.array([1_0_7, 1_0_3, 1_2_0, 1_2_7, 1_4_2, 1_2_2, 1_1_3, 1_2_2, 9_7, 1_1_1]					)

     assert np.abs(image_slice.flatten() - expected_slice					).max() == 0





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

  """simple docstring"""


  def     lowercase		(    self			:		List[Any]					)							->						Optional[int]:
     # clean up the VRAM after each test
     super().tearDown()
     gc.collect()
     torch.cuda.empty_cache()






  def     lowercase		(    self			:		Union[str, Any]					)							->						Optional[int]:
     __lowerCAmelCase							     =     torch_device

     __lowerCAmelCase							     =     DiffusionPipeline.from_pretrained('teticio/audio-diffusion-ddim-256'					)
     __lowerCAmelCase							     =     pipe.to(lowerCAmelCase_					)
     pipe.set_progress_bar_config(disable=lowerCAmelCase_					)

     __lowerCAmelCase							     =     torch.Generator(device=lowerCAmelCase_					).manual_seed(4_2					)
     __lowerCAmelCase							     =     pipe(generator=lowerCAmelCase_					)
     __lowerCAmelCase							     =     output.audios[0]
     __lowerCAmelCase							     =     output.images[0]

     assert audio.shape == (1, (pipe.unet.config.sample_size[1] - 1) * pipe.mel.hop_length)
     assert image.height == pipe.unet.config.sample_size[0] and image.width == pipe.unet.config.sample_size[1]
     __lowerCAmelCase							     =     np.frombuffer(image.tobytes() ,  dtype='uint8'					)[:1_0]
     __lowerCAmelCase							     =     np.array([1_5_1, 1_6_7, 1_5_4, 1_4_4, 1_2_2, 1_3_4, 1_2_1, 1_0_5, 7_0, 2_6]					)

     assert np.abs(image_slice.flatten() - expected_slice					).max() == 0

