




import gc
import unittest

import numpy as np
import torch

from diffusers import StableDiffusionKDiffusionPipeline
from diffusers.utils import slow, torch_device
from diffusers.utils.testing_utils import enable_full_determinism, require_torch_gpu


enable_full_determinism()






@slow
@require_torch_gpu
class  __magic_name__						(  unittest.TestCase):
				'''simple docstring'''







				def     _A							(							self:   Dict		):
											# clean up the VRAM after each test
											super().tearDown()
											gc.collect()
											torch.cuda.empty_cache()


				def     _A							(							self:   Dict		):
											SCREAMING_SNAKE_CASE_   	=    StableDiffusionKDiffusionPipeline.from_pretrained('''CompVis/stable-diffusion-v1-4'''		)
											SCREAMING_SNAKE_CASE_   	=    sd_pipe.to(_lowerCamelCase		)
											sd_pipe.set_progress_bar_config(disable=_lowerCamelCase		)

											sd_pipe.set_scheduler('''sample_euler'''		)

											SCREAMING_SNAKE_CASE_   	=    '''A painting of a squirrel eating a burger'''
											SCREAMING_SNAKE_CASE_   	=    torch.manual_seed(0		)
											SCREAMING_SNAKE_CASE_   	=    sd_pipe([prompt]		,					generator=_lowerCamelCase		,					guidance_scale=9.0		,					num_inference_steps=20		,					output_type='''np'''		)

											SCREAMING_SNAKE_CASE_   	=    output.images

											SCREAMING_SNAKE_CASE_   	=    image[0, -3:, -3:, -1]

											assert image.shape == (1, 5_12, 5_12, 3)
											SCREAMING_SNAKE_CASE_   	=    np.array([0.04_47, 0.04_92, 0.04_68, 0.04_08, 0.03_83, 0.04_08, 0.03_54, 0.03_80, 0.03_39]		)

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


				def     _A							(							self:   List[str]		):
											SCREAMING_SNAKE_CASE_   	=    StableDiffusionKDiffusionPipeline.from_pretrained('''stabilityai/stable-diffusion-2-1-base'''		)
											SCREAMING_SNAKE_CASE_   	=    sd_pipe.to(_lowerCamelCase		)
											sd_pipe.set_progress_bar_config(disable=_lowerCamelCase		)

											sd_pipe.set_scheduler('''sample_euler'''		)

											SCREAMING_SNAKE_CASE_   	=    '''A painting of a squirrel eating a burger'''
											SCREAMING_SNAKE_CASE_   	=    torch.manual_seed(0		)
											SCREAMING_SNAKE_CASE_   	=    sd_pipe([prompt]		,					generator=_lowerCamelCase		,					guidance_scale=9.0		,					num_inference_steps=20		,					output_type='''np'''		)

											SCREAMING_SNAKE_CASE_   	=    output.images

											SCREAMING_SNAKE_CASE_   	=    image[0, -3:, -3:, -1]

											assert image.shape == (1, 5_12, 5_12, 3)
											SCREAMING_SNAKE_CASE_   	=    np.array([0.12_37, 0.13_20, 0.14_38, 0.13_59, 0.13_90, 0.11_32, 0.12_77, 0.11_75, 0.11_12]		)

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





				def     _A							(							self:   Tuple		):
											SCREAMING_SNAKE_CASE_   	=    StableDiffusionKDiffusionPipeline.from_pretrained('''stabilityai/stable-diffusion-2-1-base'''		)
											SCREAMING_SNAKE_CASE_   	=    sd_pipe.to(_lowerCamelCase		)
											sd_pipe.set_progress_bar_config(disable=_lowerCamelCase		)

											sd_pipe.set_scheduler('''sample_dpmpp_2m'''		)

											SCREAMING_SNAKE_CASE_   	=    '''A painting of a squirrel eating a burger'''
											SCREAMING_SNAKE_CASE_   	=    torch.manual_seed(0		)
											SCREAMING_SNAKE_CASE_   	=    sd_pipe(
											    [prompt]		,					generator=_lowerCamelCase		,					guidance_scale=7.5		,					num_inference_steps=15		,					output_type='''np'''		,					use_karras_sigmas=_lowerCamelCase		,					)

											SCREAMING_SNAKE_CASE_   	=    output.images

											SCREAMING_SNAKE_CASE_   	=    image[0, -3:, -3:, -1]

											assert image.shape == (1, 5_12, 5_12, 3)
											SCREAMING_SNAKE_CASE_   	=    np.array(
											    [0.11_38_16_89, 0.12_11_29_21, 0.1_38_94_57, 0.12_54_96_06, 0.1_24_49_64, 0.10_83_15_17, 0.11_56_28_66, 0.10_86_78_16, 0.10_49_90_48]		)

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








from ...configuration_utils import PretrainedConfig
from ...utils import logging


__SCREAMING_SNAKE_CASE 						=logging.get_logger(__name__)

__SCREAMING_SNAKE_CASE 						={
    """caidas/swin2sr-classicalsr-x2-64""": (
        """https://huggingface.co/caidas/swin2sr-classicalsr-x2-64/resolve/main/config.json"""
    ),
}






class  __magic_name__						(  __UpperCAmelCase):
				'''simple docstring'''







				SCREAMING_SNAKE_CASE__					:							Optional[Any]	      =					"swin2sr"

				SCREAMING_SNAKE_CASE__					:							str	      =					{
				    "hidden_size": "embed_dim",
				    "num_attention_heads": "num_heads",
				    "num_hidden_layers": "num_layers",
				}


				def __init__(							self:   Optional[int]		,					_lowerCamelCase:   Dict=64		,					_lowerCamelCase:   str=1		,					_lowerCamelCase:   Optional[int]=3		,					_lowerCamelCase:   List[Any]=1_80		,					_lowerCamelCase:   Optional[int]=[6, 6, 6, 6, 6, 6]		,					_lowerCamelCase:   Union[str, Any]=[6, 6, 6, 6, 6, 6]		,					_lowerCamelCase:   Tuple=8		,					_lowerCamelCase:   List[Any]=2.0		,					_lowerCamelCase:   Dict=True		,					_lowerCamelCase:   Dict=0.0		,					_lowerCamelCase:   int=0.0		,					_lowerCamelCase:   Optional[int]=0.1		,					_lowerCamelCase:   Any="gelu"		,					_lowerCamelCase:   int=False		,					_lowerCamelCase:   Tuple=0.02		,					_lowerCamelCase:   List[str]=1E-5		,					_lowerCamelCase:   List[Any]=2		,					_lowerCamelCase:   str=1.0		,					_lowerCamelCase:   List[Any]="1conv"		,					_lowerCamelCase:   str="pixelshuffle"		,					**_lowerCamelCase:   List[str]		,					):
											super().__init__(**_lowerCamelCase		)

											SCREAMING_SNAKE_CASE_   	=    image_size
											SCREAMING_SNAKE_CASE_   	=    patch_size
											SCREAMING_SNAKE_CASE_   	=    num_channels
											SCREAMING_SNAKE_CASE_   	=    embed_dim
											SCREAMING_SNAKE_CASE_   	=    depths
											SCREAMING_SNAKE_CASE_   	=    len(_lowerCamelCase		)
											SCREAMING_SNAKE_CASE_   	=    num_heads
											SCREAMING_SNAKE_CASE_   	=    window_size
											SCREAMING_SNAKE_CASE_   	=    mlp_ratio
											SCREAMING_SNAKE_CASE_   	=    qkv_bias
											SCREAMING_SNAKE_CASE_   	=    hidden_dropout_prob
											SCREAMING_SNAKE_CASE_   	=    attention_probs_dropout_prob
											SCREAMING_SNAKE_CASE_   	=    drop_path_rate
											SCREAMING_SNAKE_CASE_   	=    hidden_act
											SCREAMING_SNAKE_CASE_   	=    use_absolute_embeddings
											SCREAMING_SNAKE_CASE_   	=    layer_norm_eps
											SCREAMING_SNAKE_CASE_   	=    initializer_range
											SCREAMING_SNAKE_CASE_   	=    upscale
											SCREAMING_SNAKE_CASE_   	=    img_range
											SCREAMING_SNAKE_CASE_   	=    resi_connection
											SCREAMING_SNAKE_CASE_   	=    upsampler


