
'''simple docstring'''
import gc
import random
import unittest

import numpy as np
import torch
from transformers import CLIPTextConfig, CLIPTextModel, CLIPTextModelWithProjection, CLIPTokenizer

from diffusers import (
    AutoencoderKL,
    DiffusionPipeline,
    EulerDiscreteScheduler,
    StableDiffusionXLImgaImgPipeline,
    UNetaDConditionModel,
)
from diffusers.utils import floats_tensor, slow, torch_device
from diffusers.utils.testing_utils import enable_full_determinism, require_torch_gpu

from ..pipeline_params import (
    IMAGE_TO_IMAGE_IMAGE_PARAMS,
    TEXT_GUIDED_IMAGE_VARIATION_BATCH_PARAMS,
    TEXT_GUIDED_IMAGE_VARIATION_PARAMS,
)
from ..test_pipelines_common import PipelineLatentTesterMixin, PipelineTesterMixin


enable_full_determinism()






class    UpperCAmelCase_				(	snake_case__    ,						snake_case__    ,						unittest.TestCase					):
 UpperCAmelCase_		    =						StableDiffusionXLImgaImgPipeline
 UpperCAmelCase_		    =						TEXT_GUIDED_IMAGE_VARIATION_PARAMS - {"""height""", """width"""}
 UpperCAmelCase_		    =						PipelineTesterMixin.required_optional_params - {"""latents"""}
 UpperCAmelCase_		    =						TEXT_GUIDED_IMAGE_VARIATION_BATCH_PARAMS
 UpperCAmelCase_		    =						IMAGE_TO_IMAGE_IMAGE_PARAMS
 UpperCAmelCase_		    =						IMAGE_TO_IMAGE_IMAGE_PARAMS





 def        snake_case__		(    self):
     torch.manual_seed(0)
     snake_case_ :    str     			=						UNetaDConditionModel(
         block_out_channels=(32, 64) ,       layers_per_block=2 ,       sample_size=32 ,       in_channels=4 ,       out_channels=4 ,       down_block_types=("DownBlock2D", "CrossAttnDownBlock2D") ,       up_block_types=("CrossAttnUpBlock2D", "UpBlock2D") ,       attention_head_dim=(2, 4) ,       use_linear_projection=lowercase_ ,       addition_embed_type="text_time" ,       addition_time_embed_dim=8 ,       transformer_layers_per_block=(1, 2) ,       projection_class_embeddings_input_dim=80 ,       cross_attention_dim=64 ,       )
     snake_case_ :    Dict     			=						EulerDiscreteScheduler(
         beta_start=0.00_085 ,       beta_end=0.012 ,       steps_offset=1 ,       beta_schedule="scaled_linear" ,       timestep_spacing="leading" ,       )
     torch.manual_seed(0)
     snake_case_ :    int     			=						AutoencoderKL(
         block_out_channels=[32, 64] ,       in_channels=3 ,       out_channels=3 ,       down_block_types=["DownEncoderBlock2D", "DownEncoderBlock2D"] ,       up_block_types=["UpDecoderBlock2D", "UpDecoderBlock2D"] ,       latent_channels=4 ,       sample_size=1_28 ,       )
     torch.manual_seed(0)
     snake_case_ :    int     			=						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=32 ,       )
     snake_case_ :    Union[str, Any]     			=						CLIPTextModel(lowercase_)
     snake_case_ :    Optional[int]     			=						CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip" ,       local_files_only=lowercase_)

     snake_case_ :    List[str]     			=						CLIPTextModelWithProjection(lowercase_)
     snake_case_ :    Optional[int]     			=						CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip" ,       local_files_only=lowercase_)

     snake_case_ :    Dict     			=						{
         "unet": unet,
         "scheduler": scheduler,
         "vae": vae,
         "text_encoder": text_encoder,
         "tokenizer": tokenizer,
         "text_encoder_2": text_encoder_a,
         "tokenizer_2": tokenizer_a,
         # "safety_checker": None,
         # "feature_extractor": None,
     }
     return components





 def        snake_case__		(    self ,       lowercase_ ,       lowercase_=0):
     snake_case_ :    List[Any]     			=						floats_tensor((1, 3, 32, 32) ,       rng=random.Random(lowercase_)).to(lowercase_)
     snake_case_ :    str     			=						image / 2 + 0.5
     if str(lowercase_).startswith("mps"):
         snake_case_ :    List[str]     			=						torch.manual_seed(lowercase_)
     else:
         snake_case_ :    List[str]     			=						torch.Generator(device=lowercase_).manual_seed(lowercase_)
     snake_case_ :    Union[str, Any]     			=						{
         "prompt": "A painting of a squirrel eating a burger",
         "image": image,
         "generator": generator,
         "num_inference_steps": 2,
         "guidance_scale": 5.0,
         "output_type": "numpy",
         "strength": 0.75,
     }
     return inputs





 def        snake_case__		(    self):
     snake_case_ :    Optional[Any]     			=						"cpu"  # ensure determinism for the device-dependent torch.Generator
     snake_case_ :    Any     			=						self.get_dummy_components()
     snake_case_ :    List[Any]     			=						StableDiffusionXLImgaImgPipeline(**lowercase_)
     snake_case_ :    Any     			=						sd_pipe.to(lowercase_)
     sd_pipe.set_progress_bar_config(disable=lowercase_)

     snake_case_ :    Dict     			=						self.get_dummy_inputs(lowercase_)
     snake_case_ :    Tuple     			=						sd_pipe(**lowercase_).images
     snake_case_ :    Optional[int]     			=						image[0, -3:, -3:, -1]

     assert image.shape == (1, 32, 32, 3)

     snake_case_ :    Union[str, Any]     			=						np.array([0.4_656, 0.4_840, 0.4_439, 0.6_698, 0.5_574, 0.4_524, 0.5_799, 0.5_943, 0.5_165])

     assert np.abs(image_slice.flatten() - expected_slice).max() < 1E-2





 def        snake_case__		(    self):
     super().test_attention_slicing_forward_pass(expected_max_diff=3E-3)





 def        snake_case__		(    self):
     super().test_inference_batch_single_identical(expected_max_diff=3E-3)





 def        snake_case__		(    self):
     pass





 def        snake_case__		(    self):
     snake_case_ :    int     			=						self.get_dummy_components()
     snake_case_ :    Tuple     			=						StableDiffusionXLImgaImgPipeline(**lowercase_)
     snake_case_ :    List[Any]     			=						sd_pipe.to(lowercase_)
     snake_case_ :    int     			=						sd_pipe.to(lowercase_)
     sd_pipe.set_progress_bar_config(disable=lowercase_)

     # forward without prompt embeds
     snake_case_ :    int     			=						self.get_dummy_inputs(lowercase_)
     snake_case_ :    Optional[Any]     			=						3 * ["this is a negative prompt"]
     snake_case_ :    Optional[int]     			=						negative_prompt
     snake_case_ :    Optional[Any]     			=						3 * [inputs["prompt"]]

     snake_case_ :    List[str]     			=						sd_pipe(**lowercase_)
     snake_case_ :    int     			=						output.images[0, -3:, -3:, -1]

     # forward with prompt embeds
     snake_case_ :    int     			=						self.get_dummy_inputs(lowercase_)
     snake_case_ :    Optional[Any]     			=						3 * ["this is a negative prompt"]
     snake_case_ :    List[Any]     			=						3 * [inputs.pop("prompt")]

     (
         (
         snake_case_
     )			, (
         snake_case_
     )			, (
         snake_case_
     )			, (
         snake_case_
     )			, 
     ) :    Dict     			=						sd_pipe.encode_prompt(lowercase_ ,       negative_prompt=lowercase_)

     snake_case_ :    Optional[Any]     			=						sd_pipe(
         **lowercase_ ,       prompt_embeds=lowercase_ ,       negative_prompt_embeds=lowercase_ ,       pooled_prompt_embeds=lowercase_ ,       negative_pooled_prompt_embeds=lowercase_ ,       )
     snake_case_ :    str     			=						output.images[0, -3:, -3:, -1]

     # make sure that it's equal
     assert np.abs(image_slice_a.flatten() - image_slice_a.flatten()).max() < 1E-4






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





 def        snake_case__		(    self):
     super().tearDown()
     gc.collect()
     torch.cuda.empty_cache()





 def        snake_case__		(    self ,       lowercase_ ,       lowercase_="cpu" ,       lowercase_=torch.floataa ,       lowercase_=0):
     snake_case_ :    Dict     			=						torch.Generator(device=lowercase_).manual_seed(lowercase_)
     snake_case_ :    List[Any]     			=						np.random.RandomState(lowercase_).standard_normal((1, 4, 64, 64))
     snake_case_ :    str     			=						torch.from_numpy(lowercase_).to(device=lowercase_ ,       dtype=lowercase_)
     snake_case_ :    Tuple     			=						{
         "prompt": "a photograph of an astronaut riding a horse",
         "latents": latents,
         "generator": generator,
         "num_inference_steps": 3,
         "guidance_scale": 7.5,
         "output_type": "numpy",
     }
     return inputs





 def        snake_case__		(    self):
     snake_case_ :    Optional[Any]     			=						DiffusionPipeline.from_pretrained("stabilityai/stable-diffusion-2-base")
     pipe.to(lowercase_)
     pipe.set_progress_bar_config(disable=lowercase_)

     snake_case_ :    Tuple     			=						self.get_inputs(lowercase_)
     snake_case_ :    int     			=						pipe(**lowercase_).images
     snake_case_ :    Optional[int]     			=						image[0, -3:, -3:, -1].flatten()

     assert image.shape == (1, 5_12, 5_12, 3)
     snake_case_ :    Tuple     			=						np.array([0.49_493, 0.47_896, 0.40_798, 0.54_214, 0.53_212, 0.48_202, 0.47_656, 0.46_329, 0.48_506])
     assert np.abs(image_slice - expected_slice).max() < 7E-3






'''simple docstring'''
from __future__ import annotations

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

a_											=						TypeVar("T")
a_											=						TypeVar("U")






class    UpperCAmelCase_				(	Generic[T, U]					):





 def __init__(    self ,       lowercase_ ,       lowercase_):
     snake_case_ :    Any     			=						key
     snake_case_ :    List[Any]     			=						val
     snake_case_ :    DoubleLinkedListNode[T, U] | None     			=						None
     snake_case_ :    DoubleLinkedListNode[T, U] | None     			=						None





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






class    UpperCAmelCase_				(	Generic[T, U]					):





 def __init__(    self):
     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):
     snake_case_ :    Dict     			=						["DoubleLinkedList"]
     snake_case_ :    Dict     			=						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        snake_case__		(    self ,       lowercase_):
     snake_case_ :    List[str]     			=						self.rear.prev

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

     snake_case_ :    Tuple     			=						node
     snake_case_ :    str     			=						previous
     snake_case_ :    Optional[Any]     			=						node
     snake_case_ :    Any     			=						self.rear





 def        snake_case__		(    self ,       lowercase_):

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

     snake_case_ :    Union[str, Any]     			=						node.next
     snake_case_ :    Optional[int]     			=						node.prev
     snake_case_ :    str     			=						None
     snake_case_ :    int     			=						None
     return node






class    UpperCAmelCase_				(	Generic[T, U]					):
 UpperCAmelCase_		    =						{}





 def __init__(    self ,       lowercase_):
     snake_case_ :    DoubleLinkedList[T, U]     			=						DoubleLinkedList()
     snake_case_ :    List[str]     			=						capacity
     snake_case_ :    Any     			=						0
     snake_case_ :    Dict     			=						0
     snake_case_ :    Union[str, Any]     			=						0
     snake_case_ :    dict[T, DoubleLinkedListNode[T, U]]     			=						{}





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





 def __contains__(    self ,       lowercase_):
     return key in self.cache





 def        snake_case__		(    self ,       lowercase_):
     # Note: pythonic interface would throw KeyError rather than return None

     if key in self.cache:
         self.hits += 1
         snake_case_ :    DoubleLinkedListNode[T, U]     			=						self.cache[key]
         snake_case_ :    Optional[Any]     			=						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        snake_case__		(    self ,       lowercase_ ,       lowercase_):

     if key not in self.cache:
         if self.num_keys >= self.capacity:
             # delete first node (oldest) when over capacity
             snake_case_ :    Optional[Any]     			=						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_ :    Dict     			=						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_ :    List[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        snake_case__		(    cls ,       lowercase_ = 1_28):

     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_ :    List[str]     			=						LRUCache(lowercase_)

             snake_case_ :    Optional[Any]     			=						cls.decorator_function_to_instance_map[func].get(args[0])
             if result is None:
                 snake_case_ :    Any     			=						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()




