





'''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 				UpperCAmelCase_				(		self		)		->						Dict:






        """simple docstring"""





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







  @property
  def 				UpperCAmelCase_				(		self		)		->						Optional[Any]:






        """simple docstring"""





        _snake_case		:	Optional[Any]					     =     1
        _snake_case		:	Union[str, Any]					     =     3
        _snake_case		:	List[Any]					     =     (32, 32)

        _snake_case		:	Dict					     =     floats_tensor((batch_size, num_channels) + sizes     ,     rng=random.Random(0		)		).to(lowercase__		)
        return image







  @property
  def 				UpperCAmelCase_				(		self		)		->						Tuple:






        """simple docstring"""





        torch.manual_seed(0		)
        _snake_case		:	List[str]					     =     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=lowercase__     ,     only_cross_attention=(True, True, False)     ,     num_class_embeds=100     ,     )
        return model







  @property
  def 				UpperCAmelCase_				(		self		)		->						List[str]:






        """simple docstring"""





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






        """simple docstring"""





        torch.manual_seed(0		)
        _snake_case		:	Union[str, Any]					     =     CLIPTextConfig(
            bos_token_id=0     ,     eos_token_id=2     ,     hidden_size=32     ,     intermediate_size=37     ,     layer_norm_eps=1E-0_5     ,     num_attention_heads=4     ,     num_hidden_layers=5     ,     pad_token_id=1     ,     vocab_size=1_000     ,     hidden_act='''gelu'''     ,     projection_dim=512     ,     )
        return CLIPTextModel(lowercase__		)







  def 				UpperCAmelCase_				(		self		)		->						Any:






        """simple docstring"""





        _snake_case		:	List[str]					     =     '''cpu'''  # ensure determinism for the device-dependent torch.Generator
        _snake_case		:	Union[str, Any]					     =     self.dummy_cond_unet_upscale
        _snake_case		:	List[Any]					     =     DDPMScheduler()
        _snake_case		:	Dict					     =     DDIMScheduler(prediction_type='''v_prediction'''		)
        _snake_case		:	Union[str, Any]					     =     self.dummy_vae
        _snake_case		:	Optional[Any]					     =     self.dummy_text_encoder
        _snake_case		:	List[Any]					     =     CLIPTokenizer.from_pretrained('''hf-internal-testing/tiny-random-clip'''		)

        _snake_case		:	int					     =     self.dummy_image.cpu().permute(0     ,     2     ,     3     ,     1		)[0]
        _snake_case		:	List[str]					     =     Image.fromarray(np.uinta(lowercase__		)		).convert('''RGB'''		).resize((64, 64)		)

        # make sure here that pndm scheduler skips prk
        _snake_case		:	int					     =     StableDiffusionUpscalePipeline(
            unet=lowercase__     ,     low_res_scheduler=lowercase__     ,     scheduler=lowercase__     ,     vae=lowercase__     ,     text_encoder=lowercase__     ,     tokenizer=lowercase__     ,     max_noise_level=350     ,     )
        _snake_case		:	Tuple					     =     sd_pipe.to(lowercase__		)
        sd_pipe.set_progress_bar_config(disable=lowercase__		)

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

        _snake_case		:	Dict					     =     output.images

        _snake_case		:	Optional[int]					     =     torch.Generator(device=lowercase__		).manual_seed(0		)
        _snake_case		:	List[str]					     =     sd_pipe(
            [prompt]     ,     image=lowercase__     ,     generator=lowercase__     ,     guidance_scale=6.0     ,     noise_level=20     ,     num_inference_steps=2     ,     output_type='''np'''     ,     return_dict=lowercase__     ,     )[0]

        _snake_case		:	Optional[Any]					     =     image[0, -3:, -3:, -1]
        _snake_case		:	Union[str, Any]					     =     image_from_tuple[0, -3:, -3:, -1]

        _snake_case		:	Dict					     =     low_res_image.size[0] * 4
        assert image.shape == (1, expected_height_width, expected_height_width, 3)
        _snake_case		:	int					     =     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		)		->						str:






        """simple docstring"""





        _snake_case		:	Any					     =     '''cpu'''  # ensure determinism for the device-dependent torch.Generator
        _snake_case		:	Dict					     =     self.dummy_cond_unet_upscale
        _snake_case		:	Any					     =     DDPMScheduler()
        _snake_case		:	List[str]					     =     DDIMScheduler(prediction_type='''v_prediction'''		)
        _snake_case		:	Optional[Any]					     =     self.dummy_vae
        _snake_case		:	List[str]					     =     self.dummy_text_encoder
        _snake_case		:	Optional[Any]					     =     CLIPTokenizer.from_pretrained('''hf-internal-testing/tiny-random-clip'''		)

        _snake_case		:	str					     =     self.dummy_image.cpu().permute(0     ,     2     ,     3     ,     1		)[0]
        _snake_case		:	Optional[Any]					     =     Image.fromarray(np.uinta(lowercase__		)		).convert('''RGB'''		).resize((64, 64)		)

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

        _snake_case		:	str					     =     '''A painting of a squirrel eating a burger'''
        _snake_case		:	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'''     ,     )
        _snake_case		:	Any					     =     output.images
        assert image.shape[0] == 2

        _snake_case		:	str					     =     torch.Generator(device=lowercase__		).manual_seed(0		)
        _snake_case		:	Optional[Any]					     =     sd_pipe(
            [prompt]     ,     image=lowercase__     ,     generator=lowercase__     ,     num_images_per_prompt=2     ,     guidance_scale=6.0     ,     noise_level=20     ,     num_inference_steps=2     ,     output_type='''np'''     ,     )
        _snake_case		:	str					     =     output.images
        assert image.shape[0] == 2







  @unittest.skipIf(torch_device != '''cuda'''     ,     '''This test requires a GPU'''		)
  def 				UpperCAmelCase_				(		self		)		->						List[str]:






        """simple docstring"""





        _snake_case		:	str					     =     self.dummy_cond_unet_upscale
        _snake_case		:	Dict					     =     DDPMScheduler()
        _snake_case		:	int					     =     DDIMScheduler(prediction_type='''v_prediction'''		)
        _snake_case		:	Dict					     =     self.dummy_vae
        _snake_case		:	Optional[int]					     =     self.dummy_text_encoder
        _snake_case		:	List[Any]					     =     CLIPTokenizer.from_pretrained('''hf-internal-testing/tiny-random-clip'''		)

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

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

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

        _snake_case		:	List[Any]					     =     '''A painting of a squirrel eating a burger'''
        _snake_case		:	int					     =     torch.manual_seed(0		)
        _snake_case		:	Tuple					     =     sd_pipe(
            [prompt]     ,     image=lowercase__     ,     generator=lowercase__     ,     num_inference_steps=2     ,     output_type='''np'''     ,     ).images

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






        """simple docstring"""





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







  def 				UpperCAmelCase_				(		self		)		->						List[str]:






        """simple docstring"""





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

        _snake_case		:	Union[str, Any]					     =     '''stabilityai/stable-diffusion-x4-upscaler'''
        _snake_case		:	Tuple					     =     StableDiffusionUpscalePipeline.from_pretrained(lowercase__		)
        pipe.to(lowercase__		)
        pipe.set_progress_bar_config(disable=lowercase__		)
        pipe.enable_attention_slicing()

        _snake_case		:	Union[str, Any]					     =     '''a cat sitting on a park bench'''

        _snake_case		:	Any					     =     torch.manual_seed(0		)
        _snake_case		:	Tuple					     =     pipe(
            prompt=lowercase__     ,     image=lowercase__     ,     generator=lowercase__     ,     output_type='''np'''     ,     )
        _snake_case		:	Optional[Any]					     =     output.images[0]

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







  def 				UpperCAmelCase_				(		self		)		->						str:






        """simple docstring"""





        _snake_case		:	str					     =     load_image(
            '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'''
            '''/sd2-upscale/low_res_cat.png'''		)
        _snake_case		:	Optional[Any]					     =     load_numpy(
            '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/sd2-upscale'''
            '''/upsampled_cat_fp16.npy'''		)

        _snake_case		:	int					     =     '''stabilityai/stable-diffusion-x4-upscaler'''
        _snake_case		:	Union[str, Any]					     =     StableDiffusionUpscalePipeline.from_pretrained(
            lowercase__     ,     torch_dtype=torch.floataa     ,     )
        pipe.to(lowercase__		)
        pipe.set_progress_bar_config(disable=lowercase__		)
        pipe.enable_attention_slicing()

        _snake_case		:	List[str]					     =     '''a cat sitting on a park bench'''

        _snake_case		:	List[str]					     =     torch.manual_seed(0		)
        _snake_case		:	str					     =     pipe(
            prompt=lowercase__     ,     image=lowercase__     ,     generator=lowercase__     ,     output_type='''np'''     ,     )
        _snake_case		:	Optional[int]					     =     output.images[0]

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







  def 				UpperCAmelCase_				(		self		)		->						Optional[Any]:






        """simple docstring"""





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

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

        _snake_case		:	Optional[int]					     =     '''stabilityai/stable-diffusion-x4-upscaler'''
        _snake_case		:	Union[str, Any]					     =     StableDiffusionUpscalePipeline.from_pretrained(
            lowercase__     ,     torch_dtype=torch.floataa     ,     )
        pipe.to(lowercase__		)
        pipe.set_progress_bar_config(disable=lowercase__		)
        pipe.enable_attention_slicing(1		)
        pipe.enable_sequential_cpu_offload()

        _snake_case		:	str					     =     '''a cat sitting on a park bench'''

        _snake_case		:	Tuple					     =     torch.manual_seed(0		)
        _snake_case		:	List[str]					     =     pipe(
            prompt=lowercase__     ,     image=lowercase__     ,     generator=lowercase__     ,     num_inference_steps=5     ,     output_type='''np'''     ,     )

        _snake_case		:	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 __future__ import annotations

from collections.abc import Callable
from typing import Generic, TypeVar

UpperCAmelCase					:		Any		=							TypeVar('T')
UpperCAmelCase					:		str		=							TypeVar('U')





class lowerCamelCase						(Generic[T, U]					):







  def __init__(		self     ,     lowercase__     ,     lowercase__		)		->						List[Any]:






        """simple docstring"""





        _snake_case		:	str					     =     key
        _snake_case		:	Optional[int]					     =     val
        _snake_case		:	DoubleLinkedListNode[T, U] | None					     =     None
        _snake_case		:	DoubleLinkedListNode[T, U] | None					     =     None







  def __repr__(		self		)		->						str:






        """simple docstring"""





        return (
            F'''Node: key: {self.key}, val: {self.val}, '''
            F'''has next: {bool(self.next		)}, has prev: {bool(self.prev		)}'''
        )





class lowerCamelCase						(Generic[T, U]					):







  def __init__(		self		)		->						None:






        """simple docstring"""





        _snake_case		:	DoubleLinkedListNode[T, U]					     =     DoubleLinkedListNode(lowercase__     ,     lowercase__		)
        _snake_case		:	DoubleLinkedListNode[T, U]					     =     DoubleLinkedListNode(lowercase__     ,     lowercase__		)
        _snake_case     ,							_snake_case		:	Union[str, Any]					     =     self.rear, self.head







  def __repr__(		self		)		->						str:






        """simple docstring"""





        _snake_case		:	List[Any]					     =     ['''DoubleLinkedList''']
        _snake_case		:	str					     =     self.head
        while node.next is not None:
              rep.append(str(lowercase__		)		)
              _snake_case		:	List[str]					     =     node.next
        rep.append(str(self.rear		)		)
        return ",\n    ".join(lowercase__		)







  def 				UpperCAmelCase_				(		self     ,     lowercase__		)		->						None:






        """simple docstring"""





        _snake_case		:	Tuple					     =     self.rear.prev

        # All nodes other than self.head are guaranteed to have non-None previous
        assert previous is not None

        _snake_case		:	Union[str, Any]					     =     node
        _snake_case		:	Optional[Any]					     =     previous
        _snake_case		:	int					     =     node
        _snake_case		:	Union[str, Any]					     =     self.rear







  def 				UpperCAmelCase_				(		self     ,     lowercase__		)		->						DoubleLinkedListNode[T, U] | None:






        """simple docstring"""





        if node.prev is None or node.next is None:
              return None

        _snake_case		:	Optional[int]					     =     node.next
        _snake_case		:	Any					     =     node.prev
        _snake_case		:	List[str]					     =     None
        _snake_case		:	Optional[int]					     =     None
        return node





class lowerCamelCase						(Generic[T, U]					):
  _lowercase  :	dict[Callable[[T], U], LRUCache[T, U]]       					=				{}







  def __init__(		self     ,     lowercase__		)		->						Union[str, Any]:






        """simple docstring"""





        _snake_case		:	DoubleLinkedList[T, U]					     =     DoubleLinkedList()
        _snake_case		:	Union[str, Any]					     =     capacity
        _snake_case		:	int					     =     0
        _snake_case		:	Dict					     =     0
        _snake_case		:	Union[str, Any]					     =     0
        _snake_case		:	dict[T, DoubleLinkedListNode[T, U]]					     =     {}







  def __repr__(		self		)		->						str:






        """simple docstring"""





        return (
            F'''CacheInfo(hits={self.hits}, misses={self.miss}, '''
            F'''capacity={self.capacity}, current size={self.num_keys})'''
        )







  def __contains__(		self     ,     lowercase__		)		->						bool:






        """simple docstring"""





        return key in self.cache







  def 				UpperCAmelCase_				(		self     ,     lowercase__		)		->						U | None:






        """simple docstring"""





        if key in self.cache:
              self.hits += 1
              _snake_case		:	DoubleLinkedListNode[T, U]					     =     self.cache[key]
              _snake_case		:	Tuple					     =     self.list.remove(self.cache[key]		)
              assert node == value_node

              # node is guaranteed not None because it is in self.cache
              assert node is not None
              self.list.add(lowercase__		)
              return node.val
        self.miss += 1
        return None







  def 				UpperCAmelCase_				(		self     ,     lowercase__     ,     lowercase__		)		->						None:






        """simple docstring"""





        if key not in self.cache:
              if self.num_keys >= self.capacity:
                    # delete first node (oldest) when over capacity
                    _snake_case		:	Dict					     =     self.list.head.next

                    # guaranteed to have a non-None first node when num_keys > 0
                    # explain to type checker via assertions
                    assert first_node is not None
                    assert first_node.key is not None
                    assert (
                        self.list.remove(lowercase__		) is not None
                    )  # node guaranteed to be in list assert node.key is not None

                    del self.cache[first_node.key]
                    self.num_keys -= 1
              _snake_case		:	Optional[int]					     =     DoubleLinkedListNode(lowercase__     ,     lowercase__		)
              self.list.add(self.cache[key]		)
              self.num_keys += 1

        else:
              # bump node to the end of the list, update value
              _snake_case		:	Optional[Any]					     =     self.list.remove(self.cache[key]		)
              assert node is not None  # node guaranteed to be in list
              _snake_case		:	Optional[Any]					     =     value
              self.list.add(lowercase__		)







  @classmethod
  def 				UpperCAmelCase_				(		cls     ,     lowercase__ = 128		)		->						Callable[[Callable[[T], U]], Callable[..., U]]:






        """simple docstring"""





        def cache_decorator_inner(lowercase__		) -> Callable[..., U]:
              def cache_decorator_wrapper(*lowercase__		) -> U:
                    if func not in cls.decorator_function_to_instance_map:
                          _snake_case		:	Optional[Any]					     =     LRUCache(lowercase__		)

                    _snake_case		:	Union[str, Any]					     =     cls.decorator_function_to_instance_map[func].get(args[0]		)
                    if result is None:
                          _snake_case		:	Tuple					     =     func(*lowercase__		)
                          cls.decorator_function_to_instance_map[func].put(args[0]     ,     lowercase__		)
                    return result

              def cache_info() -> LRUCache[T, U]:
                    return cls.decorator_function_to_instance_map[func]

              setattr(lowercase__     ,     '''cache_info'''     ,     lowercase__		)  # noqa: B010

              return cache_decorator_wrapper

        return cache_decorator_inner


if __name__ == "__main__":
   import doctest

   doctest.testmod()


