






from dataclasses import dataclass
from typing import Optional

import torch
from torch import nn

from ..configuration_utils import ConfigMixin, register_to_config
from ..utils import BaseOutput
from .attention import BasicTransformerBlock
from .modeling_utils import ModelMixin


@dataclass
class       __SCREAMING_SNAKE_CASE   (				__lowercase):
						_SCREAMING_SNAKE_CASE		:    torch.FloatTensor




class       __SCREAMING_SNAKE_CASE   (				__lowercase						,   __lowercase):






						@register_to_config
						def __init__(		self       ,  _UpperCamelCase = 16       ,  _UpperCamelCase = 88       ,  _UpperCamelCase = None       ,  _UpperCamelCase = None       ,  _UpperCamelCase = 1       ,  _UpperCamelCase = 0.0       ,  _UpperCamelCase = 32       ,  _UpperCamelCase = None       ,  _UpperCamelCase = False       ,  _UpperCamelCase = None       ,  _UpperCamelCase = "geglu"       ,  _UpperCamelCase = True       ,  _UpperCamelCase = True       ,  ):






												"""simple docstring"""
												super().__init__()
												lowerCAmelCase__        =  num_attention_heads
												lowerCAmelCase__        =  attention_head_dim
												lowerCAmelCase__        =  num_attention_heads * attention_head_dim

												lowerCAmelCase__        =  in_channels

												lowerCAmelCase__        =  torch.nn.GroupNorm(num_groups=_UpperCamelCase       ,  num_channels=_UpperCamelCase       ,  eps=1E-6       ,  affine=_UpperCamelCase				)
												lowerCAmelCase__        =  nn.Linear(_UpperCamelCase       ,  _UpperCamelCase				)

												# 3. Define transformers blocks
												lowerCAmelCase__        =  nn.ModuleList(
												    [
												        BasicTransformerBlock(
												            _UpperCamelCase       ,  _UpperCamelCase       ,  _UpperCamelCase       ,  dropout=_UpperCamelCase       ,  cross_attention_dim=_UpperCamelCase       ,  activation_fn=_UpperCamelCase       ,  attention_bias=_UpperCamelCase       ,  double_self_attention=_UpperCamelCase       ,  norm_elementwise_affine=_UpperCamelCase       ,  )
												        for d in range(_UpperCamelCase				)
												    ]				)

												lowerCAmelCase__        =  nn.Linear(_UpperCamelCase       ,  _UpperCamelCase				)







						def  UpperCamelCase__ (		self       ,  _UpperCamelCase       ,  _UpperCamelCase=None       ,  _UpperCamelCase=None       ,  _UpperCamelCase=None       ,  _UpperCamelCase=1       ,  _UpperCamelCase=None       ,  _UpperCamelCase = True       ,  ):






												"""simple docstring"""
												lowerCAmelCase__  ,    lowerCAmelCase__  ,    lowerCAmelCase__  ,    lowerCAmelCase__        =  hidden_states.shape
												lowerCAmelCase__        =  batch_frames // num_frames

												lowerCAmelCase__        =  hidden_states

												lowerCAmelCase__        =  hidden_states[None, :].reshape(_UpperCamelCase       ,  _UpperCamelCase       ,  _UpperCamelCase       ,  _UpperCamelCase       ,  _UpperCamelCase				)
												lowerCAmelCase__        =  hidden_states.permute(0       ,  2       ,  1       ,  3       ,  4				)

												lowerCAmelCase__        =  self.norm(_UpperCamelCase				)
												lowerCAmelCase__        =  hidden_states.permute(0       ,  3       ,  4       ,  2       ,  1				).reshape(batch_size * height * width       ,  _UpperCamelCase       ,  _UpperCamelCase				)

												lowerCAmelCase__        =  self.proj_in(_UpperCamelCase				)

												# 2. Blocks
												for block in self.transformer_blocks:
																		lowerCAmelCase__        =  block(
																		    _UpperCamelCase       ,  encoder_hidden_states=_UpperCamelCase       ,  timestep=_UpperCamelCase       ,  cross_attention_kwargs=_UpperCamelCase       ,  class_labels=_UpperCamelCase       ,  )

												# 3. Output
												lowerCAmelCase__        =  self.proj_out(_UpperCamelCase				)
												lowerCAmelCase__        =  (
												    hidden_states[None, None, :]
												    .reshape(_UpperCamelCase       ,  _UpperCamelCase       ,  _UpperCamelCase       ,  _UpperCamelCase       ,  _UpperCamelCase				)
												    .permute(0       ,  3       ,  4       ,  1       ,  2				)
												    .contiguous()
												)
												lowerCAmelCase__        =  hidden_states.reshape(_UpperCamelCase       ,  _UpperCamelCase       ,  _UpperCamelCase       ,  _UpperCamelCase				)

												lowerCAmelCase__        =  hidden_states + residual

												if not return_dict:
																		return (output,)

												return TransformerTemporalModelOutput(sample=_UpperCamelCase				)













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






						def  UpperCamelCase__ (		self				):






												"""simple docstring"""
												# clean up the VRAM after each test
												super().tearDown()
												gc.collect()
												torch.cuda.empty_cache()






						@property
						def  UpperCamelCase__ (		self				):






												"""simple docstring"""
												lowerCAmelCase__        =  1
												lowerCAmelCase__        =  3
												lowerCAmelCase__        =  (32, 32)

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






						@property
						def  UpperCamelCase__ (		self				):






												"""simple docstring"""
												torch.manual_seed(0				)
												lowerCAmelCase__        =  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=_UpperCamelCase       ,  only_cross_attention=(True, True, False)       ,  num_class_embeds=1_00       ,  )
												return model






						@property
						def  UpperCamelCase__ (		self				):






												"""simple docstring"""
												torch.manual_seed(0				)
												lowerCAmelCase__        =  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  UpperCamelCase__ (		self				):






												"""simple docstring"""
												torch.manual_seed(0				)
												lowerCAmelCase__        =  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=10_00       ,  hidden_act='gelu'       ,  projection_dim=5_12       ,  )
												return CLIPTextModel(_UpperCamelCase				)






						def  UpperCamelCase__ (		self				):






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

												lowerCAmelCase__        =  self.dummy_image.cpu().permute(0       ,  2       ,  3       ,  1				)[0]
												lowerCAmelCase__        =  Image.fromarray(np.uinta(_UpperCamelCase				)				).convert('RGB'				).resize((64, 64)				)

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

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

												lowerCAmelCase__        =  output.images

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

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

												lowerCAmelCase__        =  low_res_image.size[0] * 4
												assert image.shape == (1, expected_height_width, expected_height_width, 3)
												lowerCAmelCase__        =  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  UpperCamelCase__ (		self				):






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

												lowerCAmelCase__        =  self.dummy_image.cpu().permute(0       ,  2       ,  3       ,  1				)[0]
												lowerCAmelCase__        =  Image.fromarray(np.uinta(_UpperCamelCase				)				).convert('RGB'				).resize((64, 64)				)

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

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

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







						@unittest.skipIf(torch_device != 'cuda'       ,  'This test requires a GPU'				)
						def  UpperCamelCase__ (		self				):






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

												lowerCAmelCase__        =  self.dummy_image.cpu().permute(0       ,  2       ,  3       ,  1				)[0]
												lowerCAmelCase__        =  Image.fromarray(np.uinta(_UpperCamelCase				)				).convert('RGB'				).resize((64, 64)				)

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

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

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

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




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






						def  UpperCamelCase__ (		self				):






												"""simple docstring"""
												# clean up the VRAM after each test
												super().tearDown()
												gc.collect()
												torch.cuda.empty_cache()






						def  UpperCamelCase__ (		self				):






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

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

												lowerCAmelCase__        =  'a cat sitting on a park bench'

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

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






						def  UpperCamelCase__ (		self				):






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

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

												lowerCAmelCase__        =  'a cat sitting on a park bench'

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

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







						def  UpperCamelCase__ (		self				):






												"""simple docstring"""
												torch.cuda.empty_cache()
												torch.cuda.reset_max_memory_allocated()
												torch.cuda.reset_peak_memory_stats()

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

												lowerCAmelCase__        =  'stabilityai/stable-diffusion-x4-upscaler'
												lowerCAmelCase__        =  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__        =  'a cat sitting on a park bench'

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

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





