


from ...configuration_utils import PretrainedConfig
from ...utils import logging
from ...utils.backbone_utils import BackboneConfigMixin, get_aligned_output_features_output_indices


__magic_name__  	=     logging.get_logger(__name__)





class        lowerCAmelCase__ (					__lowerCamelCase,							__lowerCamelCase  ):






						"""simple docstring"""
						__UpperCAmelCase       :						int    	=			'''maskformer-swin'''

						__UpperCAmelCase       :						Union[str, Any]    	=			{
						    '''num_attention_heads''': '''num_heads''',
						    '''num_hidden_layers''': '''num_layers''',
						}




						def __init__(       self ,				a_=224 ,				a_=4 ,				a_=3 ,				a_=96 ,				a_=[2, 2, 6, 2] ,				a_=[3, 6, 12, 24] ,				a_=7 ,				a_=4.0 ,				a_=True ,				a_=0.0 ,				a_=0.0 ,				a_=0.1 ,				a_="gelu" ,				a_=False ,				a_=0.02 ,				a_=1E-5 ,				a_=None ,				a_=None ,				**a_ ,				):
							super().__init__(**snake_case__      )

							lowerCamelCase_					:				Union[str, Any]              =  image_size
							lowerCamelCase_					:				Union[str, Any]              =  patch_size
							lowerCamelCase_					:				int              =  num_channels
							lowerCamelCase_					:				Any              =  embed_dim
							lowerCamelCase_					:				int              =  depths
							lowerCamelCase_					:				Any              =  len(snake_case__      )
							lowerCamelCase_					:				List[Any]              =  num_heads
							lowerCamelCase_					:				Any              =  window_size
							lowerCamelCase_					:				Optional[Any]              =  mlp_ratio
							lowerCamelCase_					:				Any              =  qkv_bias
							lowerCamelCase_					:				Tuple              =  hidden_dropout_prob
							lowerCamelCase_					:				List[str]              =  attention_probs_dropout_prob
							lowerCamelCase_					:				str              =  drop_path_rate
							lowerCamelCase_					:				Any              =  hidden_act
							lowerCamelCase_					:				List[Any]              =  use_absolute_embeddings
							lowerCamelCase_					:				List[str]              =  layer_norm_eps
							lowerCamelCase_					:				Union[str, Any]              =  initializer_range
							# we set the hidden_size attribute in order to make Swin work with VisionEncoderDecoderModel
							# this indicates the channel dimension after the last stage of the model
							lowerCamelCase_					:				Union[str, Any]              =  int(embed_dim * 2 ** (len(snake_case__      ) - 1)      )
							lowerCamelCase_					:				Dict              =  ["stem"] + [F"""stage{idx}""" for idx in range(1 ,				len(snake_case__      ) + 1      )]
							lowerCamelCase_       ,lowerCamelCase_					:				Any              =  get_aligned_output_features_output_indices(
							    out_features=snake_case__ ,				out_indices=snake_case__ ,				stage_names=self.stage_names      )




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   UpperCamelCase			(  unittest.TestCase ):

    def UpperCamelCase      (      self			:   str     ):

     """simple docstring"""






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

    @property
    def UpperCamelCase      (      self			:   Optional[int]     ):

     """simple docstring"""






     SCREAMING_SNAKE_CASE		    =      1
     SCREAMING_SNAKE_CASE		    =      3
     SCREAMING_SNAKE_CASE		    =      (3_2, 3_2)

     SCREAMING_SNAKE_CASE		    =      floats_tensor((batch_size, num_channels) + sizes				,     rng=random.Random(0     )     ).to(snake_case__     )
     return image

    @property
    def UpperCamelCase      (      self			:   int     ):

     """simple docstring"""






     torch.manual_seed(0     )
     SCREAMING_SNAKE_CASE		    =      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=snake_case__				,     only_cross_attention=(True, True, False)				,     num_class_embeds=1_0_0				,     )
     return model

    @property
    def UpperCamelCase      (      self			:   Tuple     ):

     """simple docstring"""






     torch.manual_seed(0     )
     SCREAMING_SNAKE_CASE		    =      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 UpperCamelCase      (      self			:   List[Any]     ):

     """simple docstring"""






     torch.manual_seed(0     )
     SCREAMING_SNAKE_CASE		    =      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(snake_case__     )

    def UpperCamelCase      (      self			:   Dict     ):

     """simple docstring"""






     SCREAMING_SNAKE_CASE		    =      'cpu'  # ensure determinism for the device-dependent torch.Generator
     SCREAMING_SNAKE_CASE		    =      self.dummy_cond_unet_upscale
     SCREAMING_SNAKE_CASE		    =      DDPMScheduler()
     SCREAMING_SNAKE_CASE		    =      DDIMScheduler(prediction_type='v_prediction'     )
     SCREAMING_SNAKE_CASE		    =      self.dummy_vae
     SCREAMING_SNAKE_CASE		    =      self.dummy_text_encoder
     SCREAMING_SNAKE_CASE		    =      CLIPTokenizer.from_pretrained('hf-internal-testing/tiny-random-clip'     )

     SCREAMING_SNAKE_CASE		    =      self.dummy_image.cpu().permute(0				,     2				,     3				,     1     )[0]
     SCREAMING_SNAKE_CASE		    =      Image.fromarray(np.uinta(snake_case__     )     ).convert('RGB'     ).resize((6_4, 6_4)     )

     # make sure here that pndm scheduler skips prk
     SCREAMING_SNAKE_CASE		    =      StableDiffusionUpscalePipeline(
         unet=snake_case__				,     low_res_scheduler=snake_case__				,     scheduler=snake_case__				,     vae=snake_case__				,     text_encoder=snake_case__				,     tokenizer=snake_case__				,     max_noise_level=3_5_0				,     )
     SCREAMING_SNAKE_CASE		    =      sd_pipe.to(snake_case__     )
     sd_pipe.set_progress_bar_config(disable=snake_case__     )

     SCREAMING_SNAKE_CASE		    =      'A painting of a squirrel eating a burger'
     SCREAMING_SNAKE_CASE		    =      torch.Generator(device=snake_case__     ).manual_seed(0     )
     SCREAMING_SNAKE_CASE		    =      sd_pipe(
         [prompt]				,     image=snake_case__				,     generator=snake_case__				,     guidance_scale=6.0				,     noise_level=2_0				,     num_inference_steps=2				,     output_type='np'				,     )

     SCREAMING_SNAKE_CASE		    =      output.images

     SCREAMING_SNAKE_CASE		    =      torch.Generator(device=snake_case__     ).manual_seed(0     )
     SCREAMING_SNAKE_CASE		    =      sd_pipe(
         [prompt]				,     image=snake_case__				,     generator=snake_case__				,     guidance_scale=6.0				,     noise_level=2_0				,     num_inference_steps=2				,     output_type='np'				,     return_dict=snake_case__				,     )[0]

     SCREAMING_SNAKE_CASE		    =      image[0, -3:, -3:, -1]
     SCREAMING_SNAKE_CASE		    =      image_from_tuple[0, -3:, -3:, -1]

     SCREAMING_SNAKE_CASE		    =      low_res_image.size[0] * 4
     assert image.shape == (1, expected_height_width, expected_height_width, 3)
     SCREAMING_SNAKE_CASE		    =      np.array([0.3_113, 0.3_910, 0.4_272, 0.4_859, 0.5_061, 0.4_652, 0.5_362, 0.5_715, 0.5_661]     )

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

     """simple docstring"""






     SCREAMING_SNAKE_CASE		    =      'cpu'  # ensure determinism for the device-dependent torch.Generator
     SCREAMING_SNAKE_CASE		    =      self.dummy_cond_unet_upscale
     SCREAMING_SNAKE_CASE		    =      DDPMScheduler()
     SCREAMING_SNAKE_CASE		    =      DDIMScheduler(prediction_type='v_prediction'     )
     SCREAMING_SNAKE_CASE		    =      self.dummy_vae
     SCREAMING_SNAKE_CASE		    =      self.dummy_text_encoder
     SCREAMING_SNAKE_CASE		    =      CLIPTokenizer.from_pretrained('hf-internal-testing/tiny-random-clip'     )

     SCREAMING_SNAKE_CASE		    =      self.dummy_image.cpu().permute(0				,     2				,     3				,     1     )[0]
     SCREAMING_SNAKE_CASE		    =      Image.fromarray(np.uinta(snake_case__     )     ).convert('RGB'     ).resize((6_4, 6_4)     )

     # make sure here that pndm scheduler skips prk
     SCREAMING_SNAKE_CASE		    =      StableDiffusionUpscalePipeline(
         unet=snake_case__				,     low_res_scheduler=snake_case__				,     scheduler=snake_case__				,     vae=snake_case__				,     text_encoder=snake_case__				,     tokenizer=snake_case__				,     max_noise_level=3_5_0				,     )
     SCREAMING_SNAKE_CASE		    =      sd_pipe.to(snake_case__     )
     sd_pipe.set_progress_bar_config(disable=snake_case__     )

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

     SCREAMING_SNAKE_CASE		    =      torch.Generator(device=snake_case__     ).manual_seed(0     )
     SCREAMING_SNAKE_CASE		    =      sd_pipe(
         [prompt]				,     image=snake_case__				,     generator=snake_case__				,     num_images_per_prompt=2				,     guidance_scale=6.0				,     noise_level=2_0				,     num_inference_steps=2				,     output_type='np'				,     )
     SCREAMING_SNAKE_CASE		    =      output.images
     assert image.shape[0] == 2




    @unittest.skipIf(torch_device != 'cuda'				,     'This test requires a GPU'     )
    def UpperCamelCase      (      self			:   Union[str, Any]     ):

     """simple docstring"""






     SCREAMING_SNAKE_CASE		    =      self.dummy_cond_unet_upscale
     SCREAMING_SNAKE_CASE		    =      DDPMScheduler()
     SCREAMING_SNAKE_CASE		    =      DDIMScheduler(prediction_type='v_prediction'     )
     SCREAMING_SNAKE_CASE		    =      self.dummy_vae
     SCREAMING_SNAKE_CASE		    =      self.dummy_text_encoder
     SCREAMING_SNAKE_CASE		    =      CLIPTokenizer.from_pretrained('hf-internal-testing/tiny-random-clip'     )

     SCREAMING_SNAKE_CASE		    =      self.dummy_image.cpu().permute(0				,     2				,     3				,     1     )[0]
     SCREAMING_SNAKE_CASE		    =      Image.fromarray(np.uinta(snake_case__     )     ).convert('RGB'     ).resize((6_4, 6_4)     )

     # put models in fp16, except vae as it overflows in fp16
     SCREAMING_SNAKE_CASE		    =      unet.half()
     SCREAMING_SNAKE_CASE		    =      text_encoder.half()

     # make sure here that pndm scheduler skips prk
     SCREAMING_SNAKE_CASE		    =      StableDiffusionUpscalePipeline(
         unet=snake_case__				,     low_res_scheduler=snake_case__				,     scheduler=snake_case__				,     vae=snake_case__				,     text_encoder=snake_case__				,     tokenizer=snake_case__				,     max_noise_level=3_5_0				,     )
     SCREAMING_SNAKE_CASE		    =      sd_pipe.to(snake_case__     )
     sd_pipe.set_progress_bar_config(disable=snake_case__     )

     SCREAMING_SNAKE_CASE		    =      'A painting of a squirrel eating a burger'
     SCREAMING_SNAKE_CASE		    =      torch.manual_seed(0     )
     SCREAMING_SNAKE_CASE		    =      sd_pipe(
         [prompt]				,     image=snake_case__				,     generator=snake_case__				,     num_inference_steps=2				,     output_type='np'				,     ).images

     SCREAMING_SNAKE_CASE		    =      low_res_image.size[0] * 4
     assert image.shape == (1, expected_height_width, expected_height_width, 3)




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

    def UpperCamelCase      (      self			:   List[str]     ):

     """simple docstring"""






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

    def UpperCamelCase      (      self			:   Optional[int]     ):

     """simple docstring"""






     SCREAMING_SNAKE_CASE		    =      load_image(
         'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'
         '/sd2-upscale/low_res_cat.png'     )
     SCREAMING_SNAKE_CASE		    =      load_numpy(
         'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/sd2-upscale'
         '/upsampled_cat.npy'     )

     SCREAMING_SNAKE_CASE		    =      'stabilityai/stable-diffusion-x4-upscaler'
     SCREAMING_SNAKE_CASE		    =      StableDiffusionUpscalePipeline.from_pretrained(snake_case__     )
     pipe.to(snake_case__     )
     pipe.set_progress_bar_config(disable=snake_case__     )
     pipe.enable_attention_slicing()

     SCREAMING_SNAKE_CASE		    =      'a cat sitting on a park bench'

     SCREAMING_SNAKE_CASE		    =      torch.manual_seed(0     )
     SCREAMING_SNAKE_CASE		    =      pipe(
         prompt=snake_case__				,     image=snake_case__				,     generator=snake_case__				,     output_type='np'				,     )
     SCREAMING_SNAKE_CASE		    =      output.images[0]

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

    def UpperCamelCase      (      self			:   str     ):

     """simple docstring"""






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

     SCREAMING_SNAKE_CASE		    =      'stabilityai/stable-diffusion-x4-upscaler'
     SCREAMING_SNAKE_CASE		    =      StableDiffusionUpscalePipeline.from_pretrained(
         snake_case__				,     torch_dtype=torch.floataa				,     )
     pipe.to(snake_case__     )
     pipe.set_progress_bar_config(disable=snake_case__     )
     pipe.enable_attention_slicing()

     SCREAMING_SNAKE_CASE		    =      'a cat sitting on a park bench'

     SCREAMING_SNAKE_CASE		    =      torch.manual_seed(0     )
     SCREAMING_SNAKE_CASE		    =      pipe(
         prompt=snake_case__				,     image=snake_case__				,     generator=snake_case__				,     output_type='np'				,     )
     SCREAMING_SNAKE_CASE		    =      output.images[0]

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




    def UpperCamelCase      (      self			:   Tuple     ):

     """simple docstring"""






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

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

     SCREAMING_SNAKE_CASE		    =      'stabilityai/stable-diffusion-x4-upscaler'
     SCREAMING_SNAKE_CASE		    =      StableDiffusionUpscalePipeline.from_pretrained(
         snake_case__				,     torch_dtype=torch.floataa				,     )
     pipe.to(snake_case__     )
     pipe.set_progress_bar_config(disable=snake_case__     )
     pipe.enable_attention_slicing(1     )
     pipe.enable_sequential_cpu_offload()

     SCREAMING_SNAKE_CASE		    =      'a cat sitting on a park bench'

     SCREAMING_SNAKE_CASE		    =      torch.manual_seed(0     )
     SCREAMING_SNAKE_CASE		    =      pipe(
         prompt=snake_case__				,     image=snake_case__				,     generator=snake_case__				,     num_inference_steps=5				,     output_type='np'				,     )

     SCREAMING_SNAKE_CASE		    =      torch.cuda.max_memory_allocated()
     # make sure that less than 2.9 GB is allocated
     assert mem_bytes < 2.9 * 1_0**9

