




'''simple docstring'''







import argparse
import json
from collections import OrderedDict

import torch
from huggingface_hub import cached_download, hf_hub_url

from transformers import AutoImageProcessor, CvtConfig, CvtForImageClassification





def 						snake_case_							(_a	:							str     ):
    UpperCAmelCase			       =		[]
    embed.append(
        (
            F"cvt.encoder.stages.{idx}.embedding.convolution_embeddings.projection.weight",
            F"stage{idx}.patch_embed.proj.weight",
        )     )
    embed.append(
        (
            F"cvt.encoder.stages.{idx}.embedding.convolution_embeddings.projection.bias",
            F"stage{idx}.patch_embed.proj.bias",
        )     )
    embed.append(
        (
            F"cvt.encoder.stages.{idx}.embedding.convolution_embeddings.normalization.weight",
            F"stage{idx}.patch_embed.norm.weight",
        )     )
    embed.append(
        (
            F"cvt.encoder.stages.{idx}.embedding.convolution_embeddings.normalization.bias",
            F"stage{idx}.patch_embed.norm.bias",
        )     )
    return embed





def 						snake_case_							(_a	:							Union[str, Any]    ,   _a	:							str     ):
    UpperCAmelCase			       =		[]
    attention_weights.append(
        (
            F"cvt.encoder.stages.{idx}.layers.{cnt}.attention.attention.convolution_projection_query.convolution_projection.convolution.weight",
            F"stage{idx}.blocks.{cnt}.attn.conv_proj_q.conv.weight",
        )     )
    attention_weights.append(
        (
            F"cvt.encoder.stages.{idx}.layers.{cnt}.attention.attention.convolution_projection_query.convolution_projection.normalization.weight",
            F"stage{idx}.blocks.{cnt}.attn.conv_proj_q.bn.weight",
        )     )
    attention_weights.append(
        (
            F"cvt.encoder.stages.{idx}.layers.{cnt}.attention.attention.convolution_projection_query.convolution_projection.normalization.bias",
            F"stage{idx}.blocks.{cnt}.attn.conv_proj_q.bn.bias",
        )     )
    attention_weights.append(
        (
            F"cvt.encoder.stages.{idx}.layers.{cnt}.attention.attention.convolution_projection_query.convolution_projection.normalization.running_mean",
            F"stage{idx}.blocks.{cnt}.attn.conv_proj_q.bn.running_mean",
        )     )
    attention_weights.append(
        (
            F"cvt.encoder.stages.{idx}.layers.{cnt}.attention.attention.convolution_projection_query.convolution_projection.normalization.running_var",
            F"stage{idx}.blocks.{cnt}.attn.conv_proj_q.bn.running_var",
        )     )
    attention_weights.append(
        (
            F"cvt.encoder.stages.{idx}.layers.{cnt}.attention.attention.convolution_projection_query.convolution_projection.normalization.num_batches_tracked",
            F"stage{idx}.blocks.{cnt}.attn.conv_proj_q.bn.num_batches_tracked",
        )     )
    attention_weights.append(
        (
            F"cvt.encoder.stages.{idx}.layers.{cnt}.attention.attention.convolution_projection_key.convolution_projection.convolution.weight",
            F"stage{idx}.blocks.{cnt}.attn.conv_proj_k.conv.weight",
        )     )
    attention_weights.append(
        (
            F"cvt.encoder.stages.{idx}.layers.{cnt}.attention.attention.convolution_projection_key.convolution_projection.normalization.weight",
            F"stage{idx}.blocks.{cnt}.attn.conv_proj_k.bn.weight",
        )     )
    attention_weights.append(
        (
            F"cvt.encoder.stages.{idx}.layers.{cnt}.attention.attention.convolution_projection_key.convolution_projection.normalization.bias",
            F"stage{idx}.blocks.{cnt}.attn.conv_proj_k.bn.bias",
        )     )
    attention_weights.append(
        (
            F"cvt.encoder.stages.{idx}.layers.{cnt}.attention.attention.convolution_projection_key.convolution_projection.normalization.running_mean",
            F"stage{idx}.blocks.{cnt}.attn.conv_proj_k.bn.running_mean",
        )     )
    attention_weights.append(
        (
            F"cvt.encoder.stages.{idx}.layers.{cnt}.attention.attention.convolution_projection_key.convolution_projection.normalization.running_var",
            F"stage{idx}.blocks.{cnt}.attn.conv_proj_k.bn.running_var",
        )     )
    attention_weights.append(
        (
            F"cvt.encoder.stages.{idx}.layers.{cnt}.attention.attention.convolution_projection_key.convolution_projection.normalization.num_batches_tracked",
            F"stage{idx}.blocks.{cnt}.attn.conv_proj_k.bn.num_batches_tracked",
        )     )
    attention_weights.append(
        (
            F"cvt.encoder.stages.{idx}.layers.{cnt}.attention.attention.convolution_projection_value.convolution_projection.convolution.weight",
            F"stage{idx}.blocks.{cnt}.attn.conv_proj_v.conv.weight",
        )     )
    attention_weights.append(
        (
            F"cvt.encoder.stages.{idx}.layers.{cnt}.attention.attention.convolution_projection_value.convolution_projection.normalization.weight",
            F"stage{idx}.blocks.{cnt}.attn.conv_proj_v.bn.weight",
        )     )
    attention_weights.append(
        (
            F"cvt.encoder.stages.{idx}.layers.{cnt}.attention.attention.convolution_projection_value.convolution_projection.normalization.bias",
            F"stage{idx}.blocks.{cnt}.attn.conv_proj_v.bn.bias",
        )     )
    attention_weights.append(
        (
            F"cvt.encoder.stages.{idx}.layers.{cnt}.attention.attention.convolution_projection_value.convolution_projection.normalization.running_mean",
            F"stage{idx}.blocks.{cnt}.attn.conv_proj_v.bn.running_mean",
        )     )
    attention_weights.append(
        (
            F"cvt.encoder.stages.{idx}.layers.{cnt}.attention.attention.convolution_projection_value.convolution_projection.normalization.running_var",
            F"stage{idx}.blocks.{cnt}.attn.conv_proj_v.bn.running_var",
        )     )
    attention_weights.append(
        (
            F"cvt.encoder.stages.{idx}.layers.{cnt}.attention.attention.convolution_projection_value.convolution_projection.normalization.num_batches_tracked",
            F"stage{idx}.blocks.{cnt}.attn.conv_proj_v.bn.num_batches_tracked",
        )     )
    attention_weights.append(
        (
            F"cvt.encoder.stages.{idx}.layers.{cnt}.attention.attention.projection_query.weight",
            F"stage{idx}.blocks.{cnt}.attn.proj_q.weight",
        )     )
    attention_weights.append(
        (
            F"cvt.encoder.stages.{idx}.layers.{cnt}.attention.attention.projection_query.bias",
            F"stage{idx}.blocks.{cnt}.attn.proj_q.bias",
        )     )
    attention_weights.append(
        (
            F"cvt.encoder.stages.{idx}.layers.{cnt}.attention.attention.projection_key.weight",
            F"stage{idx}.blocks.{cnt}.attn.proj_k.weight",
        )     )
    attention_weights.append(
        (
            F"cvt.encoder.stages.{idx}.layers.{cnt}.attention.attention.projection_key.bias",
            F"stage{idx}.blocks.{cnt}.attn.proj_k.bias",
        )     )
    attention_weights.append(
        (
            F"cvt.encoder.stages.{idx}.layers.{cnt}.attention.attention.projection_value.weight",
            F"stage{idx}.blocks.{cnt}.attn.proj_v.weight",
        )     )
    attention_weights.append(
        (
            F"cvt.encoder.stages.{idx}.layers.{cnt}.attention.attention.projection_value.bias",
            F"stage{idx}.blocks.{cnt}.attn.proj_v.bias",
        )     )
    attention_weights.append(
        (
            F"cvt.encoder.stages.{idx}.layers.{cnt}.attention.output.dense.weight",
            F"stage{idx}.blocks.{cnt}.attn.proj.weight",
        )     )
    attention_weights.append(
        (
            F"cvt.encoder.stages.{idx}.layers.{cnt}.attention.output.dense.bias",
            F"stage{idx}.blocks.{cnt}.attn.proj.bias",
        )     )
    attention_weights.append(
        (F"cvt.encoder.stages.{idx}.layers.{cnt}.intermediate.dense.weight", F"stage{idx}.blocks.{cnt}.mlp.fc1.weight")     )
    attention_weights.append(
        (F"cvt.encoder.stages.{idx}.layers.{cnt}.intermediate.dense.bias", F"stage{idx}.blocks.{cnt}.mlp.fc1.bias")     )
    attention_weights.append(
        (F"cvt.encoder.stages.{idx}.layers.{cnt}.output.dense.weight", F"stage{idx}.blocks.{cnt}.mlp.fc2.weight")     )
    attention_weights.append(
        (F"cvt.encoder.stages.{idx}.layers.{cnt}.output.dense.bias", F"stage{idx}.blocks.{cnt}.mlp.fc2.bias")     )
    attention_weights.append(
        (F"cvt.encoder.stages.{idx}.layers.{cnt}.layernorm_before.weight", F"stage{idx}.blocks.{cnt}.norm1.weight")     )
    attention_weights.append(
        (F"cvt.encoder.stages.{idx}.layers.{cnt}.layernorm_before.bias", F"stage{idx}.blocks.{cnt}.norm1.bias")     )
    attention_weights.append(
        (F"cvt.encoder.stages.{idx}.layers.{cnt}.layernorm_after.weight", F"stage{idx}.blocks.{cnt}.norm2.weight")     )
    attention_weights.append(
        (F"cvt.encoder.stages.{idx}.layers.{cnt}.layernorm_after.bias", F"stage{idx}.blocks.{cnt}.norm2.bias")     )
    return attention_weights





def 						snake_case_							(_a	:							Union[str, Any]     ):
    UpperCAmelCase			       =		[]
    token.append((F"cvt.encoder.stages.{idx}.cls_token", '''stage2.cls_token''')     )
    return token





def 						snake_case_							():
    UpperCAmelCase			       =		[]
    head.append(('''layernorm.weight''', '''norm.weight''')     )
    head.append(('''layernorm.bias''', '''norm.bias''')     )
    head.append(('''classifier.weight''', '''head.weight''')     )
    head.append(('''classifier.bias''', '''head.bias''')     )
    return head





def 						snake_case_							(_a	:							Optional[Any]    ,   _a	:							str    ,   _a	:							Dict    ,   _a	:							Optional[Any]     ):
    UpperCAmelCase			       =		'''imagenet-1k-id2label.json'''
    UpperCAmelCase			       =		1_0_0_0

    UpperCAmelCase			       =		'''huggingface/label-files'''
    UpperCAmelCase			       =		num_labels
    UpperCAmelCase			       =		json.load(open(cached_download(hf_hub_url(_a    ,   _a    ,   repo_type='''dataset'''     )     )    ,   '''r'''     )     )
    UpperCAmelCase			       =		{int(_a     ): v for k, v in idalabel.items()}

    UpperCAmelCase			       =		idalabel
    UpperCAmelCase			       =		{v: k for k, v in idalabel.items()}

    UpperCAmelCase			       =		UpperCAmelCase			       =		CvtConfig(num_labels=_a    ,   idalabel=_a    ,   labelaid=_a     )

    # For depth size 13 (13 = 1+2+10)
    if cvt_model.rsplit('''/'''    ,   1     )[-1][4:6] == "13":
        UpperCAmelCase			       =		[1, 2, 1_0]

    # For depth size 21 (21 = 1+4+16)
    elif cvt_model.rsplit('''/'''    ,   1     )[-1][4:6] == "21":
        UpperCAmelCase			       =		[1, 4, 1_6]

    # For wide cvt (similar to wide-resnet) depth size 24 (w24 = 2 + 2 20)
    else:
        UpperCAmelCase			       =		[2, 2, 2_0]
        UpperCAmelCase			       =		[3, 1_2, 1_6]
        UpperCAmelCase			       =		[1_9_2, 7_6_8, 1_0_2_4]

    UpperCAmelCase			       =		CvtForImageClassification(_a     )
    UpperCAmelCase			       =		AutoImageProcessor.from_pretrained('''facebook/convnext-base-224-22k-1k'''     )
    UpperCAmelCase			       =		image_size
    UpperCAmelCase			       =		torch.load(_a    ,   map_location=torch.device('''cpu'''     )     )

    UpperCAmelCase			       =		OrderedDict()
    UpperCAmelCase			       =		[]

    for idx in range(len(config.depth     )     ):
        if config.cls_token[idx]:
            UpperCAmelCase			       =		list_of_state_dict + cls_token(_a     )
        UpperCAmelCase			       =		list_of_state_dict + embeddings(_a     )
        for cnt in range(config.depth[idx]     ):
            UpperCAmelCase			       =		list_of_state_dict + attention(_a    ,   _a     )

    UpperCAmelCase			       =		list_of_state_dict + final()
    for gg in list_of_state_dict:
        print(_a     )
    for i in range(len(_a     )     ):
        UpperCAmelCase			       =		original_weights[list_of_state_dict[i][1]]

    model.load_state_dict(_a     )
    model.save_pretrained(_a     )
    image_processor.save_pretrained(_a     )


# Download the weights from zoo: https://1drv.ms/u/s!AhIXJn_J-blW9RzF3rMW7SsLHa8h?e=blQ0Al

if __name__ == "__main__":
   A     =argparse.ArgumentParser()
   parser.add_argument(
       '--cvt_model',
       default='cvt-w24',
       type=str,
       help='Name of the cvt model you\'d like to convert.',
   )
   parser.add_argument(
       '--image_size',
       default=3_84,
       type=int,
       help='Input Image Size',
   )
   parser.add_argument(
       '--cvt_file_name',
       default=r'cvtmodels\CvT-w24-384x384-IN-22k.pth',
       type=str,
       help='Input Image Size',
   )
   parser.add_argument(
       '--pytorch_dump_folder_path', default=None, type=str, help='Path to the output PyTorch model directory.'
   )

   A     =parser.parse_args()
   convert_cvt_checkpoint(args.cvt_model, args.image_size, args.cvt_file_name, args.pytorch_dump_folder_path)







"""simple docstring"""


import gc
import random
import unittest

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

from diffusers import AutoencoderKL, DDIMScheduler, DDPMScheduler, StableDiffusionUpscalePipeline, UNetaDConditionModel
from diffusers.utils import floats_tensor, load_image, load_numpy, slow, torch_device
from diffusers.utils.testing_utils import enable_full_determinism, require_torch_gpu


enable_full_determinism()





class     snake_case	(				unittest.TestCase  ):

				def 							lowercase_     (	self							:				Optional[int])->   int:



						'''simple docstring'''


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

				@property
				def 							lowercase_     (	self							:				int)->   str:



						'''simple docstring'''


						__lowerCAmelCase:       str          =			1
						__lowerCAmelCase:       Union[str, Any]          =			3
						__lowerCAmelCase:       Union[str, Any]          =			(3_2, 3_2)

						__lowerCAmelCase:       Union[str, Any]          =			floats_tensor((batch_size, num_channels) + sizes     ,  rng=random.Random(0)).to(UpperCamelCase__)
						return image

				@property
				def 							lowercase_     (	self							:				Tuple)->   str:



						'''simple docstring'''


						torch.manual_seed(0)
						__lowerCAmelCase:       Optional[Any]          =			UNetaDConditionModel(
						    block_out_channels=(3_2, 3_2, 6_4)     ,  layers_per_block=2     ,  sample_size=3_2     ,  in_channels=7     ,  out_channels=4     ,  down_block_types=("DownBlock2D", "CrossAttnDownBlock2D", "CrossAttnDownBlock2D")     ,  up_block_types=("CrossAttnUpBlock2D", "CrossAttnUpBlock2D", "UpBlock2D")     ,  cross_attention_dim=3_2     ,  attention_head_dim=8     ,  use_linear_projection=UpperCamelCase__     ,  only_cross_attention=(True, True, False)     ,  num_class_embeds=1_0_0     ,  )
						return model

				@property
				def 							lowercase_     (	self							:				Any)->   Optional[Any]:



						'''simple docstring'''


						torch.manual_seed(0)
						__lowerCAmelCase:       Tuple          =			AutoencoderKL(
						    block_out_channels=[3_2, 3_2, 6_4]     ,  in_channels=3     ,  out_channels=3     ,  down_block_types=["DownEncoderBlock2D", "DownEncoderBlock2D", "DownEncoderBlock2D"]     ,  up_block_types=["UpDecoderBlock2D", "UpDecoderBlock2D", "UpDecoderBlock2D"]     ,  latent_channels=4     ,  )
						return model

				@property
				def 							lowercase_     (	self							:				Any)->   Optional[Any]:



						'''simple docstring'''


						torch.manual_seed(0)
						__lowerCAmelCase:       List[str]          =			CLIPTextConfig(
						    bos_token_id=0     ,  eos_token_id=2     ,  hidden_size=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     ,  hidden_act="gelu"     ,  projection_dim=5_1_2     ,  )
						return CLIPTextModel(UpperCamelCase__)

				def 							lowercase_     (	self							:				List[str])->   Dict:



						'''simple docstring'''


						__lowerCAmelCase:       Tuple          =			"cpu"  # ensure determinism for the device-dependent torch.Generator
						__lowerCAmelCase:       int          =			self.dummy_cond_unet_upscale
						__lowerCAmelCase:       int          =			DDPMScheduler()
						__lowerCAmelCase:       List[str]          =			DDIMScheduler(prediction_type="v_prediction")
						__lowerCAmelCase:       Tuple          =			self.dummy_vae
						__lowerCAmelCase:       Optional[Any]          =			self.dummy_text_encoder
						__lowerCAmelCase:       Any          =			CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip")

						__lowerCAmelCase:       Tuple          =			self.dummy_image.cpu().permute(0     ,  2     ,  3     ,  1)[0]
						__lowerCAmelCase:       List[Any]          =			Image.fromarray(np.uinta(UpperCamelCase__)).convert("RGB").resize((6_4, 6_4))

						# make sure here that pndm scheduler skips prk
						__lowerCAmelCase:       Optional[int]          =			StableDiffusionUpscalePipeline(
						    unet=UpperCamelCase__     ,  low_res_scheduler=UpperCamelCase__     ,  scheduler=UpperCamelCase__     ,  vae=UpperCamelCase__     ,  text_encoder=UpperCamelCase__     ,  tokenizer=UpperCamelCase__     ,  max_noise_level=3_5_0     ,  )
						__lowerCAmelCase:       Tuple          =			sd_pipe.to(UpperCamelCase__)
						sd_pipe.set_progress_bar_config(disable=UpperCamelCase__)

						__lowerCAmelCase:       Any          =			"A painting of a squirrel eating a burger"
						__lowerCAmelCase:       str          =			torch.Generator(device=UpperCamelCase__).manual_seed(0)
						__lowerCAmelCase:       Optional[int]          =			sd_pipe(
						    [prompt]     ,  image=UpperCamelCase__     ,  generator=UpperCamelCase__     ,  guidance_scale=6.0     ,  noise_level=2_0     ,  num_inference_steps=2     ,  output_type="np"     ,  )

						__lowerCAmelCase:       List[str]          =			output.images

						__lowerCAmelCase:       Union[str, Any]          =			torch.Generator(device=UpperCamelCase__).manual_seed(0)
						__lowerCAmelCase:       List[str]          =			sd_pipe(
						    [prompt]     ,  image=UpperCamelCase__     ,  generator=UpperCamelCase__     ,  guidance_scale=6.0     ,  noise_level=2_0     ,  num_inference_steps=2     ,  output_type="np"     ,  return_dict=UpperCamelCase__     ,  )[0]

						__lowerCAmelCase:       int          =			image[0, -3:, -3:, -1]
						__lowerCAmelCase:       Dict          =			image_from_tuple[0, -3:, -3:, -1]

						__lowerCAmelCase:       Dict          =			low_res_image.size[0] * 4
						assert image.shape == (1, expected_height_width, expected_height_width, 3)
						__lowerCAmelCase:       List[Any]          =			np.array([0.3113, 0.3910, 0.4272, 0.4859, 0.5061, 0.4652, 0.5362, 0.5715, 0.5661])

						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 							lowercase_     (	self							:				List[str])->   Optional[Any]:



						'''simple docstring'''


						__lowerCAmelCase:       Union[str, Any]          =			"cpu"  # ensure determinism for the device-dependent torch.Generator
						__lowerCAmelCase:       Dict          =			self.dummy_cond_unet_upscale
						__lowerCAmelCase:       List[str]          =			DDPMScheduler()
						__lowerCAmelCase:       Union[str, Any]          =			DDIMScheduler(prediction_type="v_prediction")
						__lowerCAmelCase:       Optional[int]          =			self.dummy_vae
						__lowerCAmelCase:       List[Any]          =			self.dummy_text_encoder
						__lowerCAmelCase:       Optional[int]          =			CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip")

						__lowerCAmelCase:       List[Any]          =			self.dummy_image.cpu().permute(0     ,  2     ,  3     ,  1)[0]
						__lowerCAmelCase:       str          =			Image.fromarray(np.uinta(UpperCamelCase__)).convert("RGB").resize((6_4, 6_4))

						# make sure here that pndm scheduler skips prk
						__lowerCAmelCase:       Optional[int]          =			StableDiffusionUpscalePipeline(
						    unet=UpperCamelCase__     ,  low_res_scheduler=UpperCamelCase__     ,  scheduler=UpperCamelCase__     ,  vae=UpperCamelCase__     ,  text_encoder=UpperCamelCase__     ,  tokenizer=UpperCamelCase__     ,  max_noise_level=3_5_0     ,  )
						__lowerCAmelCase:       Optional[int]          =			sd_pipe.to(UpperCamelCase__)
						sd_pipe.set_progress_bar_config(disable=UpperCamelCase__)

						__lowerCAmelCase:       List[str]          =			"A painting of a squirrel eating a burger"
						__lowerCAmelCase:       List[Any]          =			sd_pipe(
						    2 * [prompt]     ,  image=2 * [low_res_image]     ,  guidance_scale=6.0     ,  noise_level=2_0     ,  num_inference_steps=2     ,  output_type="np"     ,  )
						__lowerCAmelCase:       List[Any]          =			output.images
						assert image.shape[0] == 2

						__lowerCAmelCase:       Dict          =			torch.Generator(device=UpperCamelCase__).manual_seed(0)
						__lowerCAmelCase:       Optional[Any]          =			sd_pipe(
						    [prompt]     ,  image=UpperCamelCase__     ,  generator=UpperCamelCase__     ,  num_images_per_prompt=2     ,  guidance_scale=6.0     ,  noise_level=2_0     ,  num_inference_steps=2     ,  output_type="np"     ,  )
						__lowerCAmelCase:       List[Any]          =			output.images
						assert image.shape[0] == 2




				@unittest.skipIf(torch_device != "cuda"     ,  "This test requires a GPU")
				def 							lowercase_     (	self							:				Tuple)->   Any:



						'''simple docstring'''


						__lowerCAmelCase:       Union[str, Any]          =			self.dummy_cond_unet_upscale
						__lowerCAmelCase:       int          =			DDPMScheduler()
						__lowerCAmelCase:       int          =			DDIMScheduler(prediction_type="v_prediction")
						__lowerCAmelCase:       Dict          =			self.dummy_vae
						__lowerCAmelCase:       int          =			self.dummy_text_encoder
						__lowerCAmelCase:       List[Any]          =			CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip")

						__lowerCAmelCase:       List[Any]          =			self.dummy_image.cpu().permute(0     ,  2     ,  3     ,  1)[0]
						__lowerCAmelCase:       Optional[int]          =			Image.fromarray(np.uinta(UpperCamelCase__)).convert("RGB").resize((6_4, 6_4))

						# put models in fp16, except vae as it overflows in fp16
						__lowerCAmelCase:       List[Any]          =			unet.half()
						__lowerCAmelCase:       List[str]          =			text_encoder.half()

						# make sure here that pndm scheduler skips prk
						__lowerCAmelCase:       List[Any]          =			StableDiffusionUpscalePipeline(
						    unet=UpperCamelCase__     ,  low_res_scheduler=UpperCamelCase__     ,  scheduler=UpperCamelCase__     ,  vae=UpperCamelCase__     ,  text_encoder=UpperCamelCase__     ,  tokenizer=UpperCamelCase__     ,  max_noise_level=3_5_0     ,  )
						__lowerCAmelCase:       str          =			sd_pipe.to(UpperCamelCase__)
						sd_pipe.set_progress_bar_config(disable=UpperCamelCase__)

						__lowerCAmelCase:       Optional[Any]          =			"A painting of a squirrel eating a burger"
						__lowerCAmelCase:       str          =			torch.manual_seed(0)
						__lowerCAmelCase:       Dict          =			sd_pipe(
						    [prompt]     ,  image=UpperCamelCase__     ,  generator=UpperCamelCase__     ,  num_inference_steps=2     ,  output_type="np"     ,  ).images

						__lowerCAmelCase:       Optional[Any]          =			low_res_image.size[0] * 4
						assert image.shape == (1, expected_height_width, expected_height_width, 3)





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

				def 							lowercase_     (	self							:				Tuple)->   Tuple:



						'''simple docstring'''


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

				def 							lowercase_     (	self							:				List[Any])->   Tuple:



						'''simple docstring'''


						__lowerCAmelCase:       Dict          =			load_image(
						    "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main"
						    "/sd2-upscale/low_res_cat.png")
						__lowerCAmelCase:       Optional[Any]          =			load_numpy(
						    "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/sd2-upscale"
						    "/upsampled_cat.npy")

						__lowerCAmelCase:       str          =			"stabilityai/stable-diffusion-x4-upscaler"
						__lowerCAmelCase:       Optional[int]          =			StableDiffusionUpscalePipeline.from_pretrained(UpperCamelCase__)
						pipe.to(UpperCamelCase__)
						pipe.set_progress_bar_config(disable=UpperCamelCase__)
						pipe.enable_attention_slicing()

						__lowerCAmelCase:       Tuple          =			"a cat sitting on a park bench"

						__lowerCAmelCase:       int          =			torch.manual_seed(0)
						__lowerCAmelCase:       List[Any]          =			pipe(
						    prompt=UpperCamelCase__     ,  image=UpperCamelCase__     ,  generator=UpperCamelCase__     ,  output_type="np"     ,  )
						__lowerCAmelCase:       Dict          =			output.images[0]

						assert image.shape == (5_1_2, 5_1_2, 3)
						assert np.abs(expected_image - image).max() < 1e-3

				def 							lowercase_     (	self							:				Optional[int])->   Any:



						'''simple docstring'''


						__lowerCAmelCase:       Dict          =			load_image(
						    "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main"
						    "/sd2-upscale/low_res_cat.png")
						__lowerCAmelCase:       Tuple          =			load_numpy(
						    "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/sd2-upscale"
						    "/upsampled_cat_fp16.npy")

						__lowerCAmelCase:       Optional[Any]          =			"stabilityai/stable-diffusion-x4-upscaler"
						__lowerCAmelCase:       Tuple          =			StableDiffusionUpscalePipeline.from_pretrained(
						    UpperCamelCase__     ,  torch_dtype=torch.floataa     ,  )
						pipe.to(UpperCamelCase__)
						pipe.set_progress_bar_config(disable=UpperCamelCase__)
						pipe.enable_attention_slicing()

						__lowerCAmelCase:       str          =			"a cat sitting on a park bench"

						__lowerCAmelCase:       List[str]          =			torch.manual_seed(0)
						__lowerCAmelCase:       Optional[Any]          =			pipe(
						    prompt=UpperCamelCase__     ,  image=UpperCamelCase__     ,  generator=UpperCamelCase__     ,  output_type="np"     ,  )
						__lowerCAmelCase:       Union[str, Any]          =			output.images[0]

						assert image.shape == (5_1_2, 5_1_2, 3)
						assert np.abs(expected_image - image).max() < 5e-1




				def 							lowercase_     (	self							:				Optional[int])->   Dict:



						'''simple docstring'''


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

						__lowerCAmelCase:       Tuple          =			load_image(
						    "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main"
						    "/sd2-upscale/low_res_cat.png")

						__lowerCAmelCase:       Union[str, Any]          =			"stabilityai/stable-diffusion-x4-upscaler"
						__lowerCAmelCase:       Any          =			StableDiffusionUpscalePipeline.from_pretrained(
						    UpperCamelCase__     ,  torch_dtype=torch.floataa     ,  )
						pipe.to(UpperCamelCase__)
						pipe.set_progress_bar_config(disable=UpperCamelCase__)
						pipe.enable_attention_slicing(1)
						pipe.enable_sequential_cpu_offload()

						__lowerCAmelCase:       int          =			"a cat sitting on a park bench"

						__lowerCAmelCase:       Dict          =			torch.manual_seed(0)
						__lowerCAmelCase:       Dict          =			pipe(
						    prompt=UpperCamelCase__     ,  image=UpperCamelCase__     ,  generator=UpperCamelCase__     ,  num_inference_steps=5     ,  output_type="np"     ,  )

						__lowerCAmelCase:       Optional[int]          =			torch.cuda.max_memory_allocated()
						# make sure that less than 2.9 GB is allocated
						assert mem_bytes < 2.9 * 1_0**9


