
"""simple docstring"""




import gc
import os
import tempfile
import unittest

import torch
import torch.nn as nn
import torch.nn.functional as F
from transformers import CLIPTextConfig, CLIPTextModel, CLIPTokenizer

from diffusers import AutoencoderKL, DDIMScheduler, StableDiffusionPipeline, UNetaDConditionModel
from diffusers.loaders import AttnProcsLayers, LoraLoaderMixin
from diffusers.models.attention_processor import (
    Attention,
    AttnProcessor,
    AttnProcessor2_0,
    LoRAAttnProcessor,
    LoRAAttnProcessor2_0,
    LoRAXFormersAttnProcessor,
    XFormersAttnProcessor,
)
from diffusers.utils import TEXT_ENCODER_ATTN_MODULE, floats_tensor, torch_device



def 				a__			(					lowerCAmelCase						:							nn.Module		):






     '''simple docstring'''
     UpperCAmelCase__				:			Tuple							=    {}
     for name in unet.attn_processors.keys():
          UpperCAmelCase__				:			Optional[int]							=    None if name.endswith("attn1.processor"		) else unet.config.cross_attention_dim
          if name.startswith("mid_block"		):
               UpperCAmelCase__				:			Any							=    unet.config.block_out_channels[-1]
          elif name.startswith("up_blocks"		):
               UpperCAmelCase__				:			str							=    int(name[len("up_blocks."		)]		)
               UpperCAmelCase__				:			Optional[int]							=    list(reversed(unet.config.block_out_channels		)		)[block_id]
          elif name.startswith("down_blocks"		):
               UpperCAmelCase__				:			Union[str, Any]							=    int(name[len("down_blocks."		)]		)
               UpperCAmelCase__				:			List[Any]							=    unet.config.block_out_channels[block_id]
          UpperCAmelCase__				:			Optional[int]							=    (
              LoRAAttnProcessor2_0 if hasattr(lowerCAmelCase							,     "scaled_dot_product_attention"		) else LoRAAttnProcessor
          )
          UpperCAmelCase__				:			Any							=    lora_attn_processor_class(
              hidden_size=lowerCAmelCase							,     cross_attention_dim=lowerCAmelCase		)
     UpperCAmelCase__				:			List[str]							=    AttnProcsLayers(lowerCAmelCase		)
     return lora_attn_procs, unet_lora_layers



def 				a__			(					lowerCAmelCase						:							nn.Module		):






     '''simple docstring'''
     UpperCAmelCase__				:			Optional[Any]							=    {}
     UpperCAmelCase__				:			Optional[int]							=    (
         LoRAAttnProcessor2_0 if hasattr(lowerCAmelCase							,     "scaled_dot_product_attention"		) else LoRAAttnProcessor
     )
     for name, module in text_encoder.named_modules():
          if name.endswith(lowerCAmelCase		):
               UpperCAmelCase__				:			Optional[Any]							=    lora_attn_processor_class(
                   hidden_size=module.out_proj.out_features							,     cross_attention_dim=lowerCAmelCase		)
     return text_lora_attn_procs



def 				a__			(					lowerCAmelCase						:							nn.Module		):






     '''simple docstring'''
     UpperCAmelCase__				:			List[str]							=    create_text_encoder_lora_attn_procs(lowerCAmelCase		)
     UpperCAmelCase__				:			int							=    AttnProcsLayers(lowerCAmelCase		)
     return text_encoder_lora_layers



def 				a__			(					lowerCAmelCase						:							str							,     lowerCAmelCase						:							List[str]=False		):






     '''simple docstring'''
     for _, attn_proc in text_lora_attn_procs.items():
          # set up.weights
          for layer_name, layer_module in attn_proc.named_modules():
               if layer_name.endswith("_lora"		):
                    UpperCAmelCase__				:			Any							=    (
                        torch.randn_like(layer_module.up.weight		)
                        if randn_weight
                        else torch.zeros_like(layer_module.up.weight		)
                    )
                    UpperCAmelCase__				:			int							=    torch.nn.Parameter(lowerCAmelCase		)







class 	_lowercase    (							unittest.TestCase					):


     '''simple docstring'''





     def 			lowerCAmelCase__							(   self	)->       Optional[int]:
          torch.manual_seed(0	)
          UpperCAmelCase__				:			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")			,      cross_attention_dim=32			,      )
          UpperCAmelCase__				:			Union[str, Any]							=    DDIMScheduler(
              beta_start=0.0_0085			,      beta_end=0.012			,      beta_schedule="scaled_linear"			,      clip_sample=__UpperCamelCase			,      set_alpha_to_one=__UpperCamelCase			,      steps_offset=1			,      )
          torch.manual_seed(0	)
          UpperCAmelCase__				:			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			,      )
          UpperCAmelCase__				:			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			,      )
          UpperCAmelCase__				:			Dict							=    CLIPTextModel(__UpperCamelCase	)
          UpperCAmelCase__				:			int							=    CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip"	)

          UpperCAmelCase__				,						UpperCAmelCase__				:			int							=    create_unet_lora_layers(__UpperCamelCase	)
          UpperCAmelCase__				:			Optional[Any]							=    create_text_encoder_lora_layers(__UpperCamelCase	)

          UpperCAmelCase__				:			Any							=    {
              "unet": unet,
              "scheduler": scheduler,
              "vae": vae,
              "text_encoder": text_encoder,
              "tokenizer": tokenizer,
              "safety_checker": None,
              "feature_extractor": None,
          }
          UpperCAmelCase__				:			Union[str, Any]							=    {
              "unet_lora_layers": unet_lora_layers,
              "text_encoder_lora_layers": text_encoder_lora_layers,
              "unet_lora_attn_procs": unet_lora_attn_procs,
          }
          return pipeline_components, lora_components



     def 			lowerCAmelCase__							(   self	)->       List[str]:
          UpperCAmelCase__				:			Dict							=    1
          UpperCAmelCase__				:			int							=    10
          UpperCAmelCase__				:			Optional[Any]							=    4
          UpperCAmelCase__				:			Optional[int]							=    (32, 32)

          UpperCAmelCase__				:			Any							=    torch.manual_seed(0	)
          UpperCAmelCase__				:			str							=    floats_tensor((batch_size, num_channels) + sizes	)
          UpperCAmelCase__				:			Dict							=    torch.randint(1			,      __UpperCamelCase			,      size=(batch_size, sequence_length)			,      generator=__UpperCamelCase	)

          UpperCAmelCase__				:			Optional[int]							=    {
              "prompt": "A painting of a squirrel eating a burger",
              "generator": generator,
              "num_inference_steps": 2,
              "guidance_scale": 6.0,
              "output_type": "numpy",
          }

          return noise, input_ids, pipeline_inputs

          # copied from: https://colab.research.google.com/gist/sayakpaul/df2ef6e1ae6d8c10a49d859883b10860/scratchpad.ipynb



     def 			lowerCAmelCase__							(   self	)->       Tuple:
          UpperCAmelCase__				:			int							=    77

          UpperCAmelCase__				:			Union[str, Any]							=    torch.randint(2			,      56			,      size=(1, max_seq_length)			,      generator=torch.manual_seed(0	)	)

          UpperCAmelCase__				:			List[str]							=    {}
          UpperCAmelCase__				:			Any							=    inputs
          return prepared_inputs



     def 			lowerCAmelCase__							(   self			,      __UpperCamelCase	)->       Optional[Any]:
          UpperCAmelCase__				,						UpperCAmelCase__				:			Any							=    self.get_dummy_components()
          LoraLoaderMixin.save_lora_weights(
              save_directory=__UpperCamelCase			,      unet_lora_layers=lora_components["unet_lora_layers"]			,      text_encoder_lora_layers=lora_components["text_encoder_lora_layers"]			,      )
          self.assertTrue(os.path.isfile(os.path.join(__UpperCamelCase			,      "pytorch_lora_weights.bin"	)	)	)



     def 			lowerCAmelCase__							(   self	)->       Dict:
          UpperCAmelCase__				,						UpperCAmelCase__				:			int							=    self.get_dummy_components()
          UpperCAmelCase__				:			str							=    StableDiffusionPipeline(**__UpperCamelCase	)
          UpperCAmelCase__				:			Any							=    sd_pipe.to(__UpperCamelCase	)
          sd_pipe.set_progress_bar_config(disable=__UpperCamelCase	)

          UpperCAmelCase__				,						UpperCAmelCase__				,						UpperCAmelCase__				:			Tuple							=    self.get_dummy_inputs()

          UpperCAmelCase__				:			str							=    sd_pipe(**__UpperCamelCase	).images
          UpperCAmelCase__				:			Tuple							=    original_images[0, -3:, -3:, -1]

          with tempfile.TemporaryDirectory() as tmpdirname:
               LoraLoaderMixin.save_lora_weights(
                   save_directory=__UpperCamelCase			,      unet_lora_layers=lora_components["unet_lora_layers"]			,      text_encoder_lora_layers=lora_components["text_encoder_lora_layers"]			,      )
               self.assertTrue(os.path.isfile(os.path.join(__UpperCamelCase			,      "pytorch_lora_weights.bin"	)	)	)
               sd_pipe.load_lora_weights(__UpperCamelCase	)

          UpperCAmelCase__				:			Optional[Any]							=    sd_pipe(**__UpperCamelCase	).images
          UpperCAmelCase__				:			Union[str, Any]							=    lora_images[0, -3:, -3:, -1]

          # Outputs shouldn't match.
          self.assertFalse(torch.allclose(torch.from_numpy(__UpperCamelCase	)			,      torch.from_numpy(__UpperCamelCase	)	)	)



     def 			lowerCAmelCase__							(   self	)->       Dict:
          UpperCAmelCase__				,						UpperCAmelCase__				:			Optional[int]							=    self.get_dummy_components()
          UpperCAmelCase__				:			Optional[int]							=    StableDiffusionPipeline(**__UpperCamelCase	)
          UpperCAmelCase__				:			List[Any]							=    sd_pipe.to(__UpperCamelCase	)
          sd_pipe.set_progress_bar_config(disable=__UpperCamelCase	)

          UpperCAmelCase__				,						UpperCAmelCase__				,						UpperCAmelCase__				:			str							=    self.get_dummy_inputs()

          UpperCAmelCase__				:			Optional[Any]							=    sd_pipe(**__UpperCamelCase	).images
          UpperCAmelCase__				:			Dict							=    original_images[0, -3:, -3:, -1]

          with tempfile.TemporaryDirectory() as tmpdirname:
               LoraLoaderMixin.save_lora_weights(
                   save_directory=__UpperCamelCase			,      unet_lora_layers=lora_components["unet_lora_layers"]			,      text_encoder_lora_layers=lora_components["text_encoder_lora_layers"]			,      safe_serialization=__UpperCamelCase			,      )
               self.assertTrue(os.path.isfile(os.path.join(__UpperCamelCase			,      "pytorch_lora_weights.safetensors"	)	)	)
               sd_pipe.load_lora_weights(__UpperCamelCase	)

          UpperCAmelCase__				:			int							=    sd_pipe(**__UpperCamelCase	).images
          UpperCAmelCase__				:			List[Any]							=    lora_images[0, -3:, -3:, -1]

          # Outputs shouldn't match.
          self.assertFalse(torch.allclose(torch.from_numpy(__UpperCamelCase	)			,      torch.from_numpy(__UpperCamelCase	)	)	)



     def 			lowerCAmelCase__							(   self	)->       Optional[int]:
          UpperCAmelCase__				,						UpperCAmelCase__				:			Any							=    self.get_dummy_components()
          UpperCAmelCase__				:			Any							=    lora_components["unet_lora_attn_procs"]
          UpperCAmelCase__				:			Union[str, Any]							=    StableDiffusionPipeline(**__UpperCamelCase	)
          UpperCAmelCase__				:			Optional[Any]							=    sd_pipe.to(__UpperCamelCase	)
          sd_pipe.set_progress_bar_config(disable=__UpperCamelCase	)

          UpperCAmelCase__				,						UpperCAmelCase__				,						UpperCAmelCase__				:			Any							=    self.get_dummy_inputs()

          UpperCAmelCase__				:			Optional[Any]							=    sd_pipe(**__UpperCamelCase	).images
          UpperCAmelCase__				:			str							=    original_images[0, -3:, -3:, -1]

          with tempfile.TemporaryDirectory() as tmpdirname:
               UpperCAmelCase__				:			int							=    sd_pipe.unet
               unet.set_attn_processor(__UpperCamelCase	)
               unet.save_attn_procs(__UpperCamelCase	)
               self.assertTrue(os.path.isfile(os.path.join(__UpperCamelCase			,      "pytorch_lora_weights.bin"	)	)	)
               sd_pipe.load_lora_weights(__UpperCamelCase	)

          UpperCAmelCase__				:			Union[str, Any]							=    sd_pipe(**__UpperCamelCase	).images
          UpperCAmelCase__				:			Optional[Any]							=    lora_images[0, -3:, -3:, -1]

          # Outputs shouldn't match.
          self.assertFalse(torch.allclose(torch.from_numpy(__UpperCamelCase	)			,      torch.from_numpy(__UpperCamelCase	)	)	)



     def 			lowerCAmelCase__							(   self	)->       Dict:
          UpperCAmelCase__				,						UpperCAmelCase__				:			int							=    self.get_dummy_components()
          UpperCAmelCase__				:			Tuple							=    StableDiffusionPipeline(**__UpperCamelCase	)

          UpperCAmelCase__				:			Optional[int]							=    self.get_dummy_tokens()

          # inference without lora
          UpperCAmelCase__				:			int							=    pipe.text_encoder(**__UpperCamelCase	)[0]
          assert outputs_without_lora.shape == (1, 77, 32)

          # create lora_attn_procs with zeroed out up.weights
          UpperCAmelCase__				:			Optional[int]							=    create_text_encoder_lora_attn_procs(pipe.text_encoder	)
          set_lora_up_weights(__UpperCamelCase			,      randn_weight=__UpperCamelCase	)

          # monkey patch
          pipe._modify_text_encoder(__UpperCamelCase	)

          # verify that it's okay to release the text_attn_procs which holds the LoRAAttnProcessor.
          del text_attn_procs
          gc.collect()

          # inference with lora
          UpperCAmelCase__				:			Optional[int]							=    pipe.text_encoder(**__UpperCamelCase	)[0]
          assert outputs_with_lora.shape == (1, 77, 32)

          assert torch.allclose(
              __UpperCamelCase			,      __UpperCamelCase	), "lora_up_weight are all zero, so the lora outputs should be the same to without lora outputs"

          # create lora_attn_procs with randn up.weights
          UpperCAmelCase__				:			Optional[Any]							=    create_text_encoder_lora_attn_procs(pipe.text_encoder	)
          set_lora_up_weights(__UpperCamelCase			,      randn_weight=__UpperCamelCase	)

          # monkey patch
          pipe._modify_text_encoder(__UpperCamelCase	)

          # verify that it's okay to release the text_attn_procs which holds the LoRAAttnProcessor.
          del text_attn_procs
          gc.collect()

          # inference with lora
          UpperCAmelCase__				:			List[str]							=    pipe.text_encoder(**__UpperCamelCase	)[0]
          assert outputs_with_lora.shape == (1, 77, 32)

          assert not torch.allclose(
              __UpperCamelCase			,      __UpperCamelCase	), "lora_up_weight are not zero, so the lora outputs should be different to without lora outputs"



     def 			lowerCAmelCase__							(   self	)->       Optional[Any]:
          UpperCAmelCase__				,						UpperCAmelCase__				:			Optional[Any]							=    self.get_dummy_components()
          UpperCAmelCase__				:			Union[str, Any]							=    StableDiffusionPipeline(**__UpperCamelCase	)

          UpperCAmelCase__				:			List[Any]							=    self.get_dummy_tokens()

          # inference without lora
          UpperCAmelCase__				:			Union[str, Any]							=    pipe.text_encoder(**__UpperCamelCase	)[0]
          assert outputs_without_lora.shape == (1, 77, 32)

          # create lora_attn_procs with randn up.weights
          UpperCAmelCase__				:			Union[str, Any]							=    create_text_encoder_lora_attn_procs(pipe.text_encoder	)
          set_lora_up_weights(__UpperCamelCase			,      randn_weight=__UpperCamelCase	)

          # monkey patch
          pipe._modify_text_encoder(__UpperCamelCase	)

          # verify that it's okay to release the text_attn_procs which holds the LoRAAttnProcessor.
          del text_attn_procs
          gc.collect()

          # inference with lora
          UpperCAmelCase__				:			Any							=    pipe.text_encoder(**__UpperCamelCase	)[0]
          assert outputs_with_lora.shape == (1, 77, 32)

          assert not torch.allclose(
              __UpperCamelCase			,      __UpperCamelCase	), "lora outputs should be different to without lora outputs"

          # remove monkey patch
          pipe._remove_text_encoder_monkey_patch()

          # inference with removed lora
          UpperCAmelCase__				:			Dict							=    pipe.text_encoder(**__UpperCamelCase	)[0]
          assert outputs_without_lora_removed.shape == (1, 77, 32)

          assert torch.allclose(
              __UpperCamelCase			,      __UpperCamelCase	), "remove lora monkey patch should restore the original outputs"



     def 			lowerCAmelCase__							(   self	)->       List[str]:
          UpperCAmelCase__				,						UpperCAmelCase__				:			Union[str, Any]							=    self.get_dummy_components()
          UpperCAmelCase__				:			List[str]							=    StableDiffusionPipeline(**__UpperCamelCase	)
          UpperCAmelCase__				:			Optional[int]							=    sd_pipe.to(__UpperCamelCase	)
          sd_pipe.set_progress_bar_config(disable=__UpperCamelCase	)

          UpperCAmelCase__				,						UpperCAmelCase__				,						UpperCAmelCase__				:			Tuple							=    self.get_dummy_inputs()

          with tempfile.TemporaryDirectory() as tmpdirname:
               LoraLoaderMixin.save_lora_weights(
                   save_directory=__UpperCamelCase			,      unet_lora_layers=lora_components["unet_lora_layers"]			,      text_encoder_lora_layers=lora_components["text_encoder_lora_layers"]			,      )
               self.assertTrue(os.path.isfile(os.path.join(__UpperCamelCase			,      "pytorch_lora_weights.bin"	)	)	)
               sd_pipe.load_lora_weights(__UpperCamelCase	)

          UpperCAmelCase__				:			int							=    sd_pipe(**__UpperCamelCase	).images
          UpperCAmelCase__				:			Dict							=    lora_images[0, -3:, -3:, -1]

          UpperCAmelCase__				:			Optional[Any]							=    sd_pipe(**__UpperCamelCase			,      cross_attention_kwargs={"scale": 0.5}	).images
          UpperCAmelCase__				:			Union[str, Any]							=    lora_images_with_scale[0, -3:, -3:, -1]

          # Outputs shouldn't match.
          self.assertFalse(
              torch.allclose(torch.from_numpy(__UpperCamelCase	)			,      torch.from_numpy(__UpperCamelCase	)	)	)



     def 			lowerCAmelCase__							(   self	)->       List[str]:
          with tempfile.TemporaryDirectory() as tmpdirname:
               self.create_lora_weight_file(__UpperCamelCase	)

               UpperCAmelCase__				,						UpperCAmelCase__				:			Tuple							=    self.get_dummy_components()
               UpperCAmelCase__				:			Optional[int]							=    StableDiffusionPipeline(**__UpperCamelCase	)
               UpperCAmelCase__				:			Any							=    sd_pipe.to(__UpperCamelCase	)
               sd_pipe.set_progress_bar_config(disable=__UpperCamelCase	)

               # check if vanilla attention processors are used
               for _, module in sd_pipe.unet.named_modules():
                    if isinstance(__UpperCamelCase			,      __UpperCamelCase	):
                         self.assertIsInstance(module.processor			,      (AttnProcessor, AttnProcessor2_0)	)

            # load LoRA weight file
               sd_pipe.load_lora_weights(__UpperCamelCase	)

               # check if lora attention processors are used
               for _, module in sd_pipe.unet.named_modules():
                    if isinstance(__UpperCamelCase			,      __UpperCamelCase	):
                         UpperCAmelCase__				:			str							=    (
                             LoRAAttnProcessor2_0 if hasattr(__UpperCamelCase			,      "scaled_dot_product_attention"	) else LoRAAttnProcessor
                         )
                         self.assertIsInstance(module.processor			,      __UpperCamelCase	)



     @unittest.skipIf(torch_device != "cuda"			,      "This test is supposed to run on GPU"	)
     def 			lowerCAmelCase__							(   self	)->       Optional[int]:
          with tempfile.TemporaryDirectory() as tmpdirname:
               self.create_lora_weight_file(__UpperCamelCase	)

               UpperCAmelCase__				,						UpperCAmelCase__				:			Optional[Any]							=    self.get_dummy_components()
               UpperCAmelCase__				:			Tuple							=    StableDiffusionPipeline(**__UpperCamelCase	)
               UpperCAmelCase__				:			Union[str, Any]							=    sd_pipe.to(__UpperCamelCase	)
               sd_pipe.set_progress_bar_config(disable=__UpperCamelCase	)

               # enable XFormers
               sd_pipe.enable_xformers_memory_efficient_attention()

               # check if xFormers attention processors are used
               for _, module in sd_pipe.unet.named_modules():
                    if isinstance(__UpperCamelCase			,      __UpperCamelCase	):
                         self.assertIsInstance(module.processor			,      __UpperCamelCase	)

            # load LoRA weight file
               sd_pipe.load_lora_weights(__UpperCamelCase	)

               # check if lora attention processors are used
               for _, module in sd_pipe.unet.named_modules():
                    if isinstance(__UpperCamelCase			,      __UpperCamelCase	):
                         self.assertIsInstance(module.processor			,      __UpperCamelCase	)



     @unittest.skipIf(torch_device != "cuda"			,      "This test is supposed to run on GPU"	)
     def 			lowerCAmelCase__							(   self	)->       str:
          UpperCAmelCase__				,						UpperCAmelCase__				:			Union[str, Any]							=    self.get_dummy_components()
          UpperCAmelCase__				:			Tuple							=    StableDiffusionPipeline(**__UpperCamelCase	)
          UpperCAmelCase__				:			Union[str, Any]							=    sd_pipe.to(__UpperCamelCase	)
          sd_pipe.set_progress_bar_config(disable=__UpperCamelCase	)

          UpperCAmelCase__				,						UpperCAmelCase__				,						UpperCAmelCase__				:			Optional[int]							=    self.get_dummy_inputs()

          # enable XFormers
          sd_pipe.enable_xformers_memory_efficient_attention()

          UpperCAmelCase__				:			str							=    sd_pipe(**__UpperCamelCase	).images
          UpperCAmelCase__				:			str							=    original_images[0, -3:, -3:, -1]

          with tempfile.TemporaryDirectory() as tmpdirname:
               LoraLoaderMixin.save_lora_weights(
                   save_directory=__UpperCamelCase			,      unet_lora_layers=lora_components["unet_lora_layers"]			,      text_encoder_lora_layers=lora_components["text_encoder_lora_layers"]			,      )
               self.assertTrue(os.path.isfile(os.path.join(__UpperCamelCase			,      "pytorch_lora_weights.bin"	)	)	)
               sd_pipe.load_lora_weights(__UpperCamelCase	)

          UpperCAmelCase__				:			int							=    sd_pipe(**__UpperCamelCase	).images
          UpperCAmelCase__				:			List[str]							=    lora_images[0, -3:, -3:, -1]

          # Outputs shouldn't match.
          self.assertFalse(torch.allclose(torch.from_numpy(__UpperCamelCase	)			,      torch.from_numpy(__UpperCamelCase	)	)	)





"""simple docstring"""




from dataclasses import dataclass
from typing import Optional, Tuple, Union

import torch
import torch.nn as nn

from ..configuration_utils import ConfigMixin, register_to_config
from ..utils import BaseOutput
from .embeddings import GaussianFourierProjection, TimestepEmbedding, Timesteps
from .modeling_utils import ModelMixin
from .unet_ad_blocks import UNetMidBlockaD, get_down_block, get_up_block







@dataclass
class 	_lowercase    (							lowerCAmelCase_					):


     '''simple docstring'''





     _A		    =				42







class 	_lowercase    (							lowerCAmelCase_						, lowerCAmelCase_					):


     '''simple docstring'''





     @register_to_config
     def __init__(   self			,      __UpperCamelCase = None			,      __UpperCamelCase = 3			,      __UpperCamelCase = 3			,      __UpperCamelCase = False			,      __UpperCamelCase = "positional"			,      __UpperCamelCase = 0			,      __UpperCamelCase = True			,      __UpperCamelCase = ("DownBlock2D", "AttnDownBlock2D", "AttnDownBlock2D", "AttnDownBlock2D")			,      __UpperCamelCase = ("AttnUpBlock2D", "AttnUpBlock2D", "AttnUpBlock2D", "UpBlock2D")			,      __UpperCamelCase = (2_24, 4_48, 6_72, 8_96)			,      __UpperCamelCase = 2			,      __UpperCamelCase = 1			,      __UpperCamelCase = 1			,      __UpperCamelCase = "conv"			,      __UpperCamelCase = "conv"			,      __UpperCamelCase = "silu"			,      __UpperCamelCase = 8			,      __UpperCamelCase = 32			,      __UpperCamelCase = 1E-5			,      __UpperCamelCase = "default"			,      __UpperCamelCase = True			,      __UpperCamelCase = None			,      __UpperCamelCase = None			,      )->       Any:
          super().__init__()

          UpperCAmelCase__				:			Optional[int]							=    sample_size
          UpperCAmelCase__				:			Any							=    block_out_channels[0] * 4

          # Check inputs
          if len(__UpperCamelCase	) != len(__UpperCamelCase	):
               raise ValueError(
                   F"Must provide the same number of `down_block_types` as `up_block_types`. `down_block_types`: {down_block_types}. `up_block_types`: {up_block_types}."	)

          if len(__UpperCamelCase	) != len(__UpperCamelCase	):
               raise ValueError(
                   F"Must provide the same number of `block_out_channels` as `down_block_types`. `block_out_channels`: {block_out_channels}. `down_block_types`: {down_block_types}."	)

          # input
          UpperCAmelCase__				:			List[Any]							=    nn.Convad(__UpperCamelCase			,      block_out_channels[0]			,      kernel_size=3			,      padding=(1, 1)	)

          # time
          if time_embedding_type == "fourier":
               UpperCAmelCase__				:			int							=    GaussianFourierProjection(embedding_size=block_out_channels[0]			,      scale=16	)
               UpperCAmelCase__				:			Dict							=    2 * block_out_channels[0]
          elif time_embedding_type == "positional":
               UpperCAmelCase__				:			List[str]							=    Timesteps(block_out_channels[0]			,      __UpperCamelCase			,      __UpperCamelCase	)
               UpperCAmelCase__				:			Any							=    block_out_channels[0]

          UpperCAmelCase__				:			List[str]							=    TimestepEmbedding(__UpperCamelCase			,      __UpperCamelCase	)

          # class embedding
          if class_embed_type is None and num_class_embeds is not None:
               UpperCAmelCase__				:			Optional[Any]							=    nn.Embedding(__UpperCamelCase			,      __UpperCamelCase	)
          elif class_embed_type == "timestep":
               UpperCAmelCase__				:			str							=    TimestepEmbedding(__UpperCamelCase			,      __UpperCamelCase	)
          elif class_embed_type == "identity":
               UpperCAmelCase__				:			Union[str, Any]							=    nn.Identity(__UpperCamelCase			,      __UpperCamelCase	)
          else:
               UpperCAmelCase__				:			Optional[int]							=    None

          UpperCAmelCase__				:			str							=    nn.ModuleList([]	)
          UpperCAmelCase__				:			Any							=    None
          UpperCAmelCase__				:			Tuple							=    nn.ModuleList([]	)

          # down
          UpperCAmelCase__				:			Any							=    block_out_channels[0]
          for i, down_block_type in enumerate(__UpperCamelCase	):
               UpperCAmelCase__				:			List[Any]							=    output_channel
               UpperCAmelCase__				:			Optional[int]							=    block_out_channels[i]
               UpperCAmelCase__				:			Any							=    i == len(__UpperCamelCase	) - 1

               UpperCAmelCase__				:			List[Any]							=    get_down_block(
                   __UpperCamelCase			,      num_layers=__UpperCamelCase			,      in_channels=__UpperCamelCase			,      out_channels=__UpperCamelCase			,      temb_channels=__UpperCamelCase			,      add_downsample=not is_final_block			,      resnet_eps=__UpperCamelCase			,      resnet_act_fn=__UpperCamelCase			,      resnet_groups=__UpperCamelCase			,      attention_head_dim=attention_head_dim if attention_head_dim is not None else output_channel			,      downsample_padding=__UpperCamelCase			,      resnet_time_scale_shift=__UpperCamelCase			,      downsample_type=__UpperCamelCase			,      )
               self.down_blocks.append(__UpperCamelCase	)

          # mid
          UpperCAmelCase__				:			List[str]							=    UNetMidBlockaD(
              in_channels=block_out_channels[-1]			,      temb_channels=__UpperCamelCase			,      resnet_eps=__UpperCamelCase			,      resnet_act_fn=__UpperCamelCase			,      output_scale_factor=__UpperCamelCase			,      resnet_time_scale_shift=__UpperCamelCase			,      attention_head_dim=attention_head_dim if attention_head_dim is not None else block_out_channels[-1]			,      resnet_groups=__UpperCamelCase			,      add_attention=__UpperCamelCase			,      )

          # up
          UpperCAmelCase__				:			Tuple							=    list(reversed(__UpperCamelCase	)	)
          UpperCAmelCase__				:			Optional[Any]							=    reversed_block_out_channels[0]
          for i, up_block_type in enumerate(__UpperCamelCase	):
               UpperCAmelCase__				:			int							=    output_channel
               UpperCAmelCase__				:			List[Any]							=    reversed_block_out_channels[i]
               UpperCAmelCase__				:			Dict							=    reversed_block_out_channels[min(i + 1			,      len(__UpperCamelCase	) - 1	)]

               UpperCAmelCase__				:			int							=    i == len(__UpperCamelCase	) - 1

               UpperCAmelCase__				:			Dict							=    get_up_block(
                   __UpperCamelCase			,      num_layers=layers_per_block + 1			,      in_channels=__UpperCamelCase			,      out_channels=__UpperCamelCase			,      prev_output_channel=__UpperCamelCase			,      temb_channels=__UpperCamelCase			,      add_upsample=not is_final_block			,      resnet_eps=__UpperCamelCase			,      resnet_act_fn=__UpperCamelCase			,      resnet_groups=__UpperCamelCase			,      attention_head_dim=attention_head_dim if attention_head_dim is not None else output_channel			,      resnet_time_scale_shift=__UpperCamelCase			,      upsample_type=__UpperCamelCase			,      )
               self.up_blocks.append(__UpperCamelCase	)
               UpperCAmelCase__				:			Dict							=    output_channel

          # out
          UpperCAmelCase__				:			str							=    norm_num_groups if norm_num_groups is not None else min(block_out_channels[0] // 4			,      32	)
          UpperCAmelCase__				:			Tuple							=    nn.GroupNorm(num_channels=block_out_channels[0]			,      num_groups=__UpperCamelCase			,      eps=__UpperCamelCase	)
          UpperCAmelCase__				:			Any							=    nn.SiLU()
          UpperCAmelCase__				:			Optional[int]							=    nn.Convad(block_out_channels[0]			,      __UpperCamelCase			,      kernel_size=3			,      padding=1	)



     def 			lowerCAmelCase__							(   self			,      __UpperCamelCase			,      __UpperCamelCase			,      __UpperCamelCase = None			,      __UpperCamelCase = True			,      )->       Union[UNetaDOutput, Tuple]:
          # 0. center input if necessary
          if self.config.center_input_sample:
               UpperCAmelCase__				:			Any							=    2 * sample - 1.0

          # 1. time
          UpperCAmelCase__				:			Tuple							=    timestep
          if not torch.is_tensor(__UpperCamelCase	):
               UpperCAmelCase__				:			Optional[int]							=    torch.tensor([timesteps]			,      dtype=torch.long			,      device=sample.device	)
          elif torch.is_tensor(__UpperCamelCase	) and len(timesteps.shape	) == 0:
               UpperCAmelCase__				:			List[Any]							=    timesteps[None].to(sample.device	)

          # broadcast to batch dimension in a way that's compatible with ONNX/Core ML
          UpperCAmelCase__				:			Dict							=    timesteps * torch.ones(sample.shape[0]			,      dtype=timesteps.dtype			,      device=timesteps.device	)

          UpperCAmelCase__				:			List[str]							=    self.time_proj(__UpperCamelCase	)

          # timesteps does not contain any weights and will always return f32 tensors
          # but time_embedding might actually be running in fp16. so we need to cast here.
          # there might be better ways to encapsulate this.
          UpperCAmelCase__				:			Dict							=    t_emb.to(dtype=self.dtype	)
          UpperCAmelCase__				:			Optional[int]							=    self.time_embedding(__UpperCamelCase	)

          if self.class_embedding is not None:
               if class_labels is None:
                    raise ValueError("class_labels should be provided when doing class conditioning"	)

               if self.config.class_embed_type == "timestep":
                    UpperCAmelCase__				:			List[Any]							=    self.time_proj(__UpperCamelCase	)

               UpperCAmelCase__				:			List[str]							=    self.class_embedding(__UpperCamelCase	).to(dtype=self.dtype	)
               UpperCAmelCase__				:			Optional[int]							=    emb + class_emb

          # 2. pre-process
          UpperCAmelCase__				:			List[str]							=    sample
          UpperCAmelCase__				:			str							=    self.conv_in(__UpperCamelCase	)

          # 3. down
          UpperCAmelCase__				:			List[str]							=    (sample,)
          for downsample_block in self.down_blocks:
               if hasattr(__UpperCamelCase			,      "skip_conv"	):
                    UpperCAmelCase__				,						UpperCAmelCase__				,						UpperCAmelCase__				:			Optional[Any]							=    downsample_block(
                        hidden_states=__UpperCamelCase			,      temb=__UpperCamelCase			,      skip_sample=__UpperCamelCase	)
               else:
                    UpperCAmelCase__				,						UpperCAmelCase__				:			Optional[Any]							=    downsample_block(hidden_states=__UpperCamelCase			,      temb=__UpperCamelCase	)

               down_block_res_samples += res_samples

          # 4. mid
          UpperCAmelCase__				:			Union[str, Any]							=    self.mid_block(__UpperCamelCase			,      __UpperCamelCase	)

          # 5. up
          UpperCAmelCase__				:			Dict							=    None
          for upsample_block in self.up_blocks:
               UpperCAmelCase__				:			Union[str, Any]							=    down_block_res_samples[-len(upsample_block.resnets	) :]
               UpperCAmelCase__				:			Optional[Any]							=    down_block_res_samples[: -len(upsample_block.resnets	)]

               if hasattr(__UpperCamelCase			,      "skip_conv"	):
                    UpperCAmelCase__				,						UpperCAmelCase__				:			str							=    upsample_block(__UpperCamelCase			,      __UpperCamelCase			,      __UpperCamelCase			,      __UpperCamelCase	)
               else:
                    UpperCAmelCase__				:			Tuple							=    upsample_block(__UpperCamelCase			,      __UpperCamelCase			,      __UpperCamelCase	)

        # 6. post-process
          UpperCAmelCase__				:			List[Any]							=    self.conv_norm_out(__UpperCamelCase	)
          UpperCAmelCase__				:			Any							=    self.conv_act(__UpperCamelCase	)
          UpperCAmelCase__				:			Optional[Any]							=    self.conv_out(__UpperCamelCase	)

          if skip_sample is not None:
               sample += skip_sample

          if self.config.time_embedding_type == "fourier":
               UpperCAmelCase__				:			str							=    timesteps.reshape((sample.shape[0], *([1] * len(sample.shape[1:]	)))	)
               UpperCAmelCase__				:			int							=    sample / timesteps

          if not return_dict:
               return (sample,)

          return UNetaDOutput(sample=__UpperCamelCase	)



