


"""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    lowerCamelCase__						(		unittest.TestCase    ):
      def SCREAMING_SNAKE_CASE_	(		self	:	Tuple							):






       '''simple docstring'''


       super().tearDown()
       gc.collect()
       torch.cuda.empty_cache()
      @property
      def SCREAMING_SNAKE_CASE_	(		self	:	Dict							):






       '''simple docstring'''


       __lowercase									=	1
       __lowercase									=	3
       __lowercase									=	(3_2, 3_2)

       __lowercase									=	floats_tensor((batch_size, num_channels) + sizes      ,       rng=random.Random(0							)							).to(A_							)
       return image
      @property
      def SCREAMING_SNAKE_CASE_	(		self	:	str							):






       '''simple docstring'''


       torch.manual_seed(0							)
       __lowercase									=	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=A_      ,       only_cross_attention=(True, True, False)      ,       num_class_embeds=1_0_0      ,       )
       return model
      @property
      def SCREAMING_SNAKE_CASE_	(		self	:	Optional[Any]							):






       '''simple docstring'''


       torch.manual_seed(0							)
       __lowercase									=	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 SCREAMING_SNAKE_CASE_	(		self	:	Dict							):






       '''simple docstring'''


       torch.manual_seed(0							)
       __lowercase									=	CLIPTextConfig(
           bos_token_id=0      ,       eos_token_id=2      ,       hidden_size=3_2      ,       intermediate_size=3_7      ,       layer_norm_eps=1e-0_5      ,       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(A_							)
      def SCREAMING_SNAKE_CASE_	(		self	:	str							):






       '''simple docstring'''


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

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

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

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

       __lowercase									=	output.images

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

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

       __lowercase									=	low_res_image.size[0] * 4
       assert image.shape == (1, expected_height_width, expected_height_width, 3)
       __lowercase									=	np.array([0.31_13, 0.39_10, 0.42_72, 0.48_59, 0.50_61, 0.46_52, 0.53_62, 0.57_15, 0.56_61]							)

       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 SCREAMING_SNAKE_CASE_	(		self	:	str							):






       '''simple docstring'''


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

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

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

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

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






      @unittest.skipIf(torch_device != """cuda"""      ,       """This test requires a GPU"""							)
      def SCREAMING_SNAKE_CASE_	(		self	:	List[str]							):






       '''simple docstring'''


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

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

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

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

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

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






@slow
@require_torch_gpu
class    lowerCamelCase__						(		unittest.TestCase    ):
      def SCREAMING_SNAKE_CASE_	(		self	:	Any							):






       '''simple docstring'''


       super().tearDown()
       gc.collect()
       torch.cuda.empty_cache()
      def SCREAMING_SNAKE_CASE_	(		self	:	List[Any]							):






       '''simple docstring'''


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

       __lowercase									=	"""stabilityai/stable-diffusion-x4-upscaler"""
       __lowercase									=	StableDiffusionUpscalePipeline.from_pretrained(A_							)
       pipe.to(A_							)
       pipe.set_progress_bar_config(disable=A_							)
       pipe.enable_attention_slicing()

       __lowercase									=	"""a cat sitting on a park bench"""

       __lowercase									=	torch.manual_seed(0							)
       __lowercase									=	pipe(
           prompt=A_      ,       image=A_      ,       generator=A_      ,       output_type="""np"""      ,       )
       __lowercase									=	output.images[0]

       assert image.shape == (5_1_2, 5_1_2, 3)
       assert np.abs(expected_image - image							).max() < 1e-3
      def SCREAMING_SNAKE_CASE_	(		self	:	List[Any]							):






       '''simple docstring'''


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

       __lowercase									=	"""stabilityai/stable-diffusion-x4-upscaler"""
       __lowercase									=	StableDiffusionUpscalePipeline.from_pretrained(
           A_      ,       torch_dtype=torch.floataa      ,       )
       pipe.to(A_							)
       pipe.set_progress_bar_config(disable=A_							)
       pipe.enable_attention_slicing()

       __lowercase									=	"""a cat sitting on a park bench"""

       __lowercase									=	torch.manual_seed(0							)
       __lowercase									=	pipe(
           prompt=A_      ,       image=A_      ,       generator=A_      ,       output_type="""np"""      ,       )
       __lowercase									=	output.images[0]

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






      def SCREAMING_SNAKE_CASE_	(		self	:	Optional[int]							):






       '''simple docstring'''


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

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

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

       __lowercase									=	"""a cat sitting on a park bench"""

       __lowercase									=	torch.manual_seed(0							)
       __lowercase									=	pipe(
           prompt=A_      ,       image=A_      ,       generator=A_      ,       num_inference_steps=5      ,       output_type="""np"""      ,       )

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




"""simple docstring"""

import warnings
from collections import OrderedDict
from typing import Mapping

from packaging import version

from ...configuration_utils import PretrainedConfig
from ...onnx import OnnxConfig
from ...utils import logging


UpperCAmelCase__  							=logging.get_logger(__name__)

UpperCAmelCase__  							={
    "nvidia/segformer-b0-finetuned-ade-512-512": (
        "https://huggingface.co/nvidia/segformer-b0-finetuned-ade-512-512/resolve/main/config.json"
    ),
    # See all SegFormer models at https://huggingface.co/models?filter=segformer
}





class    lowerCamelCase__						(		_a    ):
      a :     Optional[int]											=		"""segformer"""
      def __init__(		self	:	str      ,       A_	:	Optional[int]=3      ,       A_	:	Any=4      ,       A_	:	Optional[int]=[2, 2, 2, 2]      ,       A_	:	str=[8, 4, 2, 1]      ,       A_	:	Any=[3_2, 6_4, 1_6_0, 2_5_6]      ,       A_	:	Union[str, Any]=[7, 3, 3, 3]      ,       A_	:	Any=[4, 2, 2, 2]      ,       A_	:	int=[1, 2, 5, 8]      ,       A_	:	Union[str, Any]=[4, 4, 4, 4]      ,       A_	:	Optional[Any]="gelu"      ,       A_	:	List[str]=0.0      ,       A_	:	Any=0.0      ,       A_	:	Optional[Any]=0.1      ,       A_	:	List[str]=0.02      ,       A_	:	Tuple=0.1      ,       A_	:	Optional[Any]=1e-6      ,       A_	:	Optional[int]=2_5_6      ,       A_	:	Optional[Any]=2_5_5      ,       **A_	:	Union[str, Any]      ,       ):






       '''simple docstring'''


       super().__init__(**A_							)

       if "reshape_last_stage" in kwargs and kwargs["reshape_last_stage"] is False:
        warnings.warn(
            """Reshape_last_stage is set to False in this config. This argument is deprecated and will soon be"""
            """ removed, as the behaviour will default to that of reshape_last_stage = True."""      ,       A_      ,       )

       __lowercase									=	num_channels
       __lowercase									=	num_encoder_blocks
       __lowercase									=	depths
       __lowercase									=	sr_ratios
       __lowercase									=	hidden_sizes
       __lowercase									=	patch_sizes
       __lowercase									=	strides
       __lowercase									=	mlp_ratios
       __lowercase									=	num_attention_heads
       __lowercase									=	hidden_act
       __lowercase									=	hidden_dropout_prob
       __lowercase									=	attention_probs_dropout_prob
       __lowercase									=	classifier_dropout_prob
       __lowercase									=	initializer_range
       __lowercase									=	drop_path_rate
       __lowercase									=	layer_norm_eps
       __lowercase									=	decoder_hidden_size
       __lowercase									=	kwargs.get("""reshape_last_stage"""      ,       A_							)
       __lowercase									=	semantic_loss_ignore_index






class    lowerCamelCase__						(		_a    ):
      a :     List[str]											=		version.parse("""1.11"""    )
      @property
      def SCREAMING_SNAKE_CASE_	(		self	:	Any							):






       '''simple docstring'''


       return OrderedDict(
           [
               ("""pixel_values""", {0: """batch""", 1: """num_channels""", 2: """height""", 3: """width"""}),
           ]							)
      @property
      def SCREAMING_SNAKE_CASE_	(		self	:	int							):






       '''simple docstring'''


       return 1e-4






      @property
      def SCREAMING_SNAKE_CASE_	(		self	:	int							):






       '''simple docstring'''


       return 1_2
