





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






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






							@property
							def     __a   (   self      )	-> Tuple:
								a      :		Dict	=				1
								a      :		Any	=				3
								a      :		Optional[Any]	=				(32, 32)

								a      :		Tuple	=				floats_tensor((batch_size, num_channels) + sizes      ,					rng=random.Random(0      )      ).to(lowerCAmelCase__      )
								return image






							@property
							def     __a   (   self      )	-> Union[str, Any]:
								torch.manual_seed(0      )
								a      :		Any	=				UNetaDConditionModel(
								    block_out_channels=(32, 32, 64)      ,					layers_per_block=2      ,					sample_size=32      ,					in_channels=7      ,					out_channels=4      ,					down_block_types=("DownBlock2D", "CrossAttnDownBlock2D", "CrossAttnDownBlock2D")      ,					up_block_types=("CrossAttnUpBlock2D", "CrossAttnUpBlock2D", "UpBlock2D")      ,					cross_attention_dim=32      ,					attention_head_dim=8      ,					use_linear_projection=lowerCAmelCase__      ,					only_cross_attention=(True, True, False)      ,					num_class_embeds=100      ,					)
								return model






							@property
							def     __a   (   self      )	-> Dict:
								torch.manual_seed(0      )
								a      :		Dict	=				AutoencoderKL(
								    block_out_channels=[32, 32, 64]      ,					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     __a   (   self      )	-> int:
								torch.manual_seed(0      )
								a      :		str	=				CLIPTextConfig(
								    bos_token_id=0      ,					eos_token_id=2      ,					hidden_size=32      ,					intermediate_size=37      ,					layer_norm_eps=1E-05      ,					num_attention_heads=4      ,					num_hidden_layers=5      ,					pad_token_id=1      ,					vocab_size=1000      ,					hidden_act="gelu"      ,					projection_dim=512      ,					)
								return CLIPTextModel(lowerCAmelCase__      )






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

								a      :		int	=				self.dummy_image.cpu().permute(0      ,					2      ,					3      ,					1      )[0]
								a      :		List[Any]	=				Image.fromarray(np.uinta(lowerCAmelCase__      )      ).convert("RGB"      ).resize((64, 64)      )

								# make sure here that pndm scheduler skips prk
								a      :		List[str]	=				StableDiffusionUpscalePipeline(
								    unet=lowerCAmelCase__      ,					low_res_scheduler=lowerCAmelCase__      ,					scheduler=lowerCAmelCase__      ,					vae=lowerCAmelCase__      ,					text_encoder=lowerCAmelCase__      ,					tokenizer=lowerCAmelCase__      ,					max_noise_level=350      ,					)
								a      :		List[str]	=				sd_pipe.to(lowerCAmelCase__      )
								sd_pipe.set_progress_bar_config(disable=lowerCAmelCase__      )

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

								a      :		Dict	=				output.images

								a      :		Optional[Any]	=				torch.Generator(device=lowerCAmelCase__      ).manual_seed(0      )
								a      :		Dict	=				sd_pipe(
								    [prompt]      ,					image=lowerCAmelCase__      ,					generator=lowerCAmelCase__      ,					guidance_scale=6.0      ,					noise_level=20      ,					num_inference_steps=2      ,					output_type="np"      ,					return_dict=lowerCAmelCase__      ,					)[0]

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

								a      :		str	=				low_res_image.size[0] * 4
								assert image.shape == (1, expected_height_width, expected_height_width, 3)
								a      :		Any	=				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     __a   (   self      )	-> Tuple:
								a      :		Tuple	=				"cpu"  # ensure determinism for the device-dependent torch.Generator
								a      :		Optional[Any]	=				self.dummy_cond_unet_upscale
								a      :		str	=				DDPMScheduler()
								a      :		str	=				DDIMScheduler(prediction_type="v_prediction"      )
								a      :		Tuple	=				self.dummy_vae
								a      :		Any	=				self.dummy_text_encoder
								a      :		Union[str, Any]	=				CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip"      )

								a      :		Optional[Any]	=				self.dummy_image.cpu().permute(0      ,					2      ,					3      ,					1      )[0]
								a      :		Optional[int]	=				Image.fromarray(np.uinta(lowerCAmelCase__      )      ).convert("RGB"      ).resize((64, 64)      )

								# make sure here that pndm scheduler skips prk
								a      :		Optional[int]	=				StableDiffusionUpscalePipeline(
								    unet=lowerCAmelCase__      ,					low_res_scheduler=lowerCAmelCase__      ,					scheduler=lowerCAmelCase__      ,					vae=lowerCAmelCase__      ,					text_encoder=lowerCAmelCase__      ,					tokenizer=lowerCAmelCase__      ,					max_noise_level=350      ,					)
								a      :		Union[str, Any]	=				sd_pipe.to(lowerCAmelCase__      )
								sd_pipe.set_progress_bar_config(disable=lowerCAmelCase__      )

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

								a      :		int	=				torch.Generator(device=lowerCAmelCase__      ).manual_seed(0      )
								a      :		int	=				sd_pipe(
								    [prompt]      ,					image=lowerCAmelCase__      ,					generator=lowerCAmelCase__      ,					num_images_per_prompt=2      ,					guidance_scale=6.0      ,					noise_level=20      ,					num_inference_steps=2      ,					output_type="np"      ,					)
								a      :		List[Any]	=				output.images
								assert image.shape[0] == 2






							@unittest.skipIf(torch_device != "cuda"      ,					"This test requires a GPU"      )
							def     __a   (   self      )	-> Tuple:
								a      :		Any	=				self.dummy_cond_unet_upscale
								a      :		Optional[int]	=				DDPMScheduler()
								a      :		int	=				DDIMScheduler(prediction_type="v_prediction"      )
								a      :		List[Any]	=				self.dummy_vae
								a      :		Any	=				self.dummy_text_encoder
								a      :		Optional[int]	=				CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip"      )

								a      :		int	=				self.dummy_image.cpu().permute(0      ,					2      ,					3      ,					1      )[0]
								a      :		Dict	=				Image.fromarray(np.uinta(lowerCAmelCase__      )      ).convert("RGB"      ).resize((64, 64)      )

								# put models in fp16, except vae as it overflows in fp16
								a      :		int	=				unet.half()
								a      :		str	=				text_encoder.half()

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

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

								a      :		Any	=				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     __a   (   self      )	-> Optional[int]:
								# clean up the VRAM after each test
								super().tearDown()
								gc.collect()
								torch.cuda.empty_cache()






							def     __a   (   self      )	-> str:
								a      :		Optional[Any]	=				load_image(
								    "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main"
								    "/sd2-upscale/low_res_cat.png"      )
								a      :		Tuple	=				load_numpy(
								    "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/sd2-upscale"
								    "/upsampled_cat.npy"      )

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

								a      :		Dict	=				"a cat sitting on a park bench"

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

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






							def     __a   (   self      )	-> Optional[Any]:
								a      :		Optional[int]	=				load_image(
								    "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main"
								    "/sd2-upscale/low_res_cat.png"      )
								a      :		List[str]	=				load_numpy(
								    "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/sd2-upscale"
								    "/upsampled_cat_fp16.npy"      )

								a      :		List[Any]	=				"stabilityai/stable-diffusion-x4-upscaler"
								a      :		List[str]	=				StableDiffusionUpscalePipeline.from_pretrained(
								    lowerCAmelCase__      ,					torch_dtype=torch.floataa      ,					)
								pipe.to(lowerCAmelCase__      )
								pipe.set_progress_bar_config(disable=lowerCAmelCase__      )
								pipe.enable_attention_slicing()

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

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

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






							def     __a   (   self      )	-> Optional[Any]:
								torch.cuda.empty_cache()
								torch.cuda.reset_max_memory_allocated()
								torch.cuda.reset_peak_memory_stats()

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

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

								a      :		Dict	=				"a cat sitting on a park bench"

								a      :		List[Any]	=				torch.manual_seed(0      )
								a      :		Optional[Any]	=				pipe(
								    prompt=lowerCAmelCase__      ,					image=lowerCAmelCase__      ,					generator=lowerCAmelCase__      ,					num_inference_steps=5      ,					output_type="np"      ,					)

								a      :		List[Any]	=				torch.cuda.max_memory_allocated()
								# make sure that less than 2.9 GB is allocated
								assert mem_bytes < 2.9 * 10**9









"""simple docstring"""




from typing import Dict, List, Optional, Tuple, Union

import torch

from ...models import AutoencoderKL, TransformeraDModel
from ...schedulers import KarrasDiffusionSchedulers
from ...utils import randn_tensor
from ..pipeline_utils import DiffusionPipeline, ImagePipelineOutput


class        __UpperCamelCase				(						a__						):






							def __init__(   self      ,					lowerCAmelCase__      ,					lowerCAmelCase__      ,					lowerCAmelCase__      ,					lowerCAmelCase__ = None      ,					)	-> int:
								super().__init__()
								self.register_modules(transformer=lowerCAmelCase__      ,					vae=lowerCAmelCase__      ,					scheduler=lowerCAmelCase__      )

								# create a imagenet -> id dictionary for easier use
								a      :		str	=				{}
								if idalabel is not None:
									for key, value in idalabel.items():
										for label in value.split(","      ):
											a      :		int	=				int(lowerCAmelCase__      )
									a      :		Any	=				dict(sorted(self.labels.items()      )      )






							def     __a   (   self      ,					lowerCAmelCase__      )	-> List[int]:

								if not isinstance(lowerCAmelCase__      ,					lowerCAmelCase__      ):
									a      :		Optional[Any]	=				list(lowerCAmelCase__      )

								for l in label:
									if l not in self.labels:
										raise ValueError(
										    f"""{l} does not exist. Please make sure to select one of the following labels: \n {self.labels}."""      )

								return [self.labels[l] for l in label]






							@torch.no_grad()
							def __call__(   self      ,					lowerCAmelCase__      ,					lowerCAmelCase__ = 4.0      ,					lowerCAmelCase__ = None      ,					lowerCAmelCase__ = 50      ,					lowerCAmelCase__ = "pil"      ,					lowerCAmelCase__ = True      ,					)	-> Union[ImagePipelineOutput, Tuple]:
								a      :		Dict	=				len(lowerCAmelCase__      )
								a      :		Tuple	=				self.transformer.config.sample_size
								a      :		Tuple	=				self.transformer.config.in_channels

								a      :		Optional[int]	=				randn_tensor(
								    shape=(batch_size, latent_channels, latent_size, latent_size)      ,					generator=lowerCAmelCase__      ,					device=self.device      ,					dtype=self.transformer.dtype      ,					)
								a      :		Tuple	=				torch.cat([latents] * 2      ) if guidance_scale > 1 else latents

								a      :		List[str]	=				torch.tensor(lowerCAmelCase__      ,					device=self.device      ).reshape(-1      )
								a      :		str	=				torch.tensor([1000] * batch_size      ,					device=self.device      )
								a      :		Dict	=				torch.cat([class_labels, class_null]      ,					0      ) if guidance_scale > 1 else class_labels

								# set step values
								self.scheduler.set_timesteps(lowerCAmelCase__      )

								for t in self.progress_bar(self.scheduler.timesteps      ):
									if guidance_scale > 1:
										a      :		Any	=				latent_model_input[: len(lowerCAmelCase__      ) // 2]
										a      :		Tuple	=				torch.cat([half, half]      ,					dim=0      )
									a      :		List[str]	=				self.scheduler.scale_model_input(lowerCAmelCase__      ,					lowerCAmelCase__      )

									a      :		Dict	=				t
									if not torch.is_tensor(lowerCAmelCase__      ):
										# TODO: this requires sync between CPU and GPU. So try to pass timesteps as tensors if you can
										# This would be a good case for the `match` statement (Python 3.10+)
										a      :		Optional[int]	=				latent_model_input.device.type == "mps"
										if isinstance(lowerCAmelCase__      ,					lowerCAmelCase__      ):
											a      :		Union[str, Any]	=				torch.floataa if is_mps else torch.floataa
										else:
											a      :		str	=				torch.intaa if is_mps else torch.intaa
										a      :		List[str]	=				torch.tensor([timesteps]      ,					dtype=lowerCAmelCase__      ,					device=latent_model_input.device      )
									elif len(timesteps.shape      ) == 0:
										a      :		List[Any]	=				timesteps[None].to(latent_model_input.device      )
									# broadcast to batch dimension in a way that's compatible with ONNX/Core ML
									a      :		Dict	=				timesteps.expand(latent_model_input.shape[0]      )
									# predict noise model_output
									a      :		Union[str, Any]	=				self.transformer(
									    lowerCAmelCase__      ,					timestep=lowerCAmelCase__      ,					class_labels=lowerCAmelCase__      ).sample

									# perform guidance
									if guidance_scale > 1:
										a,       a      :		Any	=				noise_pred[:, :latent_channels], noise_pred[:, latent_channels:]
										a,       a      :		Union[str, Any]	=				torch.split(lowerCAmelCase__      ,					len(lowerCAmelCase__      ) // 2      ,					dim=0      )

										a      :		Dict	=				uncond_eps + guidance_scale * (cond_eps - uncond_eps)
										a      :		Optional[int]	=				torch.cat([half_eps, half_eps]      ,					dim=0      )

										a      :		Optional[int]	=				torch.cat([eps, rest]      ,					dim=1      )

									# learned sigma
									if self.transformer.config.out_channels // 2 == latent_channels:
										a,       a      :		str	=				torch.split(lowerCAmelCase__      ,					lowerCAmelCase__      ,					dim=1      )
									else:
										a      :		Any	=				noise_pred

									# compute previous image: x_t -> x_t-1
									a      :		Optional[int]	=				self.scheduler.step(lowerCAmelCase__      ,					lowerCAmelCase__      ,					lowerCAmelCase__      ).prev_sample

								if guidance_scale > 1:
									a,       a      :		Tuple	=				latent_model_input.chunk(2      ,					dim=0      )
								else:
									a      :		Tuple	=				latent_model_input

								a      :		Optional[Any]	=				1 / self.vae.config.scaling_factor * latents
								a      :		Any	=				self.vae.decode(lowerCAmelCase__      ).sample

								a      :		Tuple	=				(samples / 2 + 0.5).clamp(0      ,					1      )

								# we always cast to float32 as this does not cause significant overhead and is compatible with bfloat16
								a      :		Optional[Any]	=				samples.cpu().permute(0      ,					2      ,					3      ,					1      ).float().numpy()

								if output_type == "pil":
									a      :		int	=				self.numpy_to_pil(lowerCAmelCase__      )

								if not return_dict:
									return (samples,)

								return ImagePipelineOutput(images=lowerCAmelCase__      )


