

"""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    (    UpperCamelCase_	:			nn.Module)      ->		Tuple:
   '''simple docstring'''


   __lowercase								=		{}
   for name in unet.attn_processors.keys():
      __lowercase								=		None if name.endswith("attn1.processor") else unet.config.cross_attention_dim
      if name.startswith("mid_block"):
         __lowercase								=		unet.config.block_out_channels[-1]
      elif name.startswith("up_blocks"):
         __lowercase								=		int(name[len("up_blocks.")])
         __lowercase								=		list(reversed(unet.config.block_out_channels))[block_id]
      elif name.startswith("down_blocks"):
         __lowercase								=		int(name[len("down_blocks.")])
         __lowercase								=		unet.config.block_out_channels[block_id]
      __lowercase								=		(
          LoRAAttnProcessor2_0 if hasattr(UpperCamelCase_,     "scaled_dot_product_attention") else LoRAAttnProcessor
      )
      __lowercase								=		lora_attn_processor_class(
          hidden_size=UpperCamelCase_,     cross_attention_dim=UpperCamelCase_)
   __lowercase								=		AttnProcsLayers(UpperCamelCase_)
   return lora_attn_procs, unet_lora_layers




def     _A    (    UpperCamelCase_	:			nn.Module)      ->		Tuple:
   '''simple docstring'''


   __lowercase								=		{}
   __lowercase								=		(
       LoRAAttnProcessor2_0 if hasattr(UpperCamelCase_,     "scaled_dot_product_attention") else LoRAAttnProcessor
   )
   for name, module in text_encoder.named_modules():
      if name.endswith(UpperCamelCase_):
         __lowercase								=		lora_attn_processor_class(
             hidden_size=module.out_proj.out_features,     cross_attention_dim=UpperCamelCase_)
   return text_lora_attn_procs




def     _A    (    UpperCamelCase_	:			nn.Module)      ->		int:
   '''simple docstring'''


   __lowercase								=		create_text_encoder_lora_attn_procs(UpperCamelCase_)
   __lowercase								=		AttnProcsLayers(UpperCamelCase_)
   return text_encoder_lora_layers




def     _A    (    UpperCamelCase_	:			Union[str, Any],     UpperCamelCase_	:			Optional[Any]=False)      ->		List[Any]:
   '''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"):
            __lowercase								=		(
                torch.randn_like(layer_module.up.weight)
                if randn_weight
                else torch.zeros_like(layer_module.up.weight)
            )
            __lowercase								=		torch.nn.Parameter(UpperCamelCase_)




class    _lowerCAmelCase    (      unittest.TestCase  ):







 """simple docstring"""







 def    _lowercase					(     self				:     Optional[Any]      ):
    torch.manual_seed(0      )
    __lowercase								=		UNetaDConditionModel(
        block_out_channels=(3_2, 6_4),   layers_per_block=2,   sample_size=3_2,   in_channels=4,   out_channels=4,   down_block_types=("DownBlock2D", "CrossAttnDownBlock2D"),   up_block_types=("CrossAttnUpBlock2D", "UpBlock2D"),   cross_attention_dim=3_2,   )
    __lowercase								=		DDIMScheduler(
        beta_start=0.00_085,   beta_end=0.012,   beta_schedule="scaled_linear",   clip_sample=UpperCAmelCase__,   set_alpha_to_one=UpperCAmelCase__,   steps_offset=1,   )
    torch.manual_seed(0      )
    __lowercase								=		AutoencoderKL(
        block_out_channels=[3_2, 6_4],   in_channels=3,   out_channels=3,   down_block_types=["DownEncoderBlock2D", "DownEncoderBlock2D"],   up_block_types=["UpDecoderBlock2D", "UpDecoderBlock2D"],   latent_channels=4,   )
    __lowercase								=		CLIPTextConfig(
        bos_token_id=0,   eos_token_id=2,   hidden_size=3_2,   intermediate_size=3_7,   layer_norm_eps=1E-05,   num_attention_heads=4,   num_hidden_layers=5,   pad_token_id=1,   vocab_size=1_0_0_0,   )
    __lowercase								=		CLIPTextModel(UpperCAmelCase__      )
    __lowercase								=		CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip"      )

    __lowercase						,__lowercase								=		create_unet_lora_layers(UpperCAmelCase__      )
    __lowercase								=		create_text_encoder_lora_layers(UpperCAmelCase__      )

    __lowercase								=		{
        "unet": unet,
        "scheduler": scheduler,
        "vae": vae,
        "text_encoder": text_encoder,
        "tokenizer": tokenizer,
        "safety_checker": None,
        "feature_extractor": None,
    }
    __lowercase								=		{
        "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    _lowercase					(     self				:     List[Any]      ):
    __lowercase								=		1
    __lowercase								=		1_0
    __lowercase								=		4
    __lowercase								=		(3_2, 3_2)

    __lowercase								=		torch.manual_seed(0      )
    __lowercase								=		floats_tensor((batch_size, num_channels) + sizes      )
    __lowercase								=		torch.randint(1,   UpperCAmelCase__,   size=(batch_size, sequence_length),   generator=UpperCAmelCase__      )

    __lowercase								=		{
        "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    _lowercase					(     self				:     Union[str, Any]      ):
    __lowercase								=		7_7

    __lowercase								=		torch.randint(2,   5_6,   size=(1, max_seq_length),   generator=torch.manual_seed(0      )      )

    __lowercase								=		{}
    __lowercase								=		inputs
    return prepared_inputs


 def    _lowercase					(     self				:     List[str],   UpperCAmelCase__				:     str      ):
    __lowercase						,__lowercase								=		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    _lowercase					(     self				:     Any      ):
    __lowercase						,__lowercase								=		self.get_dummy_components()
    __lowercase								=		StableDiffusionPipeline(**UpperCAmelCase__      )
    __lowercase								=		sd_pipe.to(UpperCAmelCase__      )
    sd_pipe.set_progress_bar_config(disable=UpperCAmelCase__      )

    __lowercase						,__lowercase						,__lowercase								=		self.get_dummy_inputs()

    __lowercase								=		sd_pipe(**UpperCAmelCase__      ).images
    __lowercase								=		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__      )

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

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


 def    _lowercase					(     self				:     Tuple      ):
    __lowercase						,__lowercase								=		self.get_dummy_components()
    __lowercase								=		StableDiffusionPipeline(**UpperCAmelCase__      )
    __lowercase								=		sd_pipe.to(UpperCAmelCase__      )
    sd_pipe.set_progress_bar_config(disable=UpperCAmelCase__      )

    __lowercase						,__lowercase						,__lowercase								=		self.get_dummy_inputs()

    __lowercase								=		sd_pipe(**UpperCAmelCase__      ).images
    __lowercase								=		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__      )

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

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


 def    _lowercase					(     self				:     Optional[int]      ):
    __lowercase						,__lowercase								=		self.get_dummy_components()
    __lowercase								=		lora_components["unet_lora_attn_procs"]
    __lowercase								=		StableDiffusionPipeline(**UpperCAmelCase__      )
    __lowercase								=		sd_pipe.to(UpperCAmelCase__      )
    sd_pipe.set_progress_bar_config(disable=UpperCAmelCase__      )

    __lowercase						,__lowercase						,__lowercase								=		self.get_dummy_inputs()

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

    with tempfile.TemporaryDirectory() as tmpdirname:
       __lowercase								=		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__      )

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

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


 def    _lowercase					(     self				:     str      ):
    __lowercase						,__lowercase								=		self.get_dummy_components()
    __lowercase								=		StableDiffusionPipeline(**UpperCAmelCase__      )

    __lowercase								=		self.get_dummy_tokens()

    # inference without lora
    __lowercase								=		pipe.text_encoder(**UpperCAmelCase__      )[0]
    assert outputs_without_lora.shape == (1, 7_7, 3_2)

    # create lora_attn_procs with zeroed out up.weights
    __lowercase								=		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
    __lowercase								=		pipe.text_encoder(**UpperCAmelCase__      )[0]
    assert outputs_with_lora.shape == (1, 7_7, 3_2)

    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
    __lowercase								=		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
    __lowercase								=		pipe.text_encoder(**UpperCAmelCase__      )[0]
    assert outputs_with_lora.shape == (1, 7_7, 3_2)

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


 def    _lowercase					(     self				:     List[str]      ):
    __lowercase						,__lowercase								=		self.get_dummy_components()
    __lowercase								=		StableDiffusionPipeline(**UpperCAmelCase__      )

    __lowercase								=		self.get_dummy_tokens()

    # inference without lora
    __lowercase								=		pipe.text_encoder(**UpperCAmelCase__      )[0]
    assert outputs_without_lora.shape == (1, 7_7, 3_2)

    # create lora_attn_procs with randn up.weights
    __lowercase								=		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
    __lowercase								=		pipe.text_encoder(**UpperCAmelCase__      )[0]
    assert outputs_with_lora.shape == (1, 7_7, 3_2)

    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
    __lowercase								=		pipe.text_encoder(**UpperCAmelCase__      )[0]
    assert outputs_without_lora_removed.shape == (1, 7_7, 3_2)

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


 def    _lowercase					(     self				:     Dict      ):
    __lowercase						,__lowercase								=		self.get_dummy_components()
    __lowercase								=		StableDiffusionPipeline(**UpperCAmelCase__      )
    __lowercase								=		sd_pipe.to(UpperCAmelCase__      )
    sd_pipe.set_progress_bar_config(disable=UpperCAmelCase__      )

    __lowercase						,__lowercase						,__lowercase								=		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__      )

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

    __lowercase								=		sd_pipe(**UpperCAmelCase__,   cross_attention_kwargs={"scale": 0.5}      ).images
    __lowercase								=		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    _lowercase					(     self				:     Union[str, Any]      ):
    with tempfile.TemporaryDirectory() as tmpdirname:
       self.create_lora_weight_file(UpperCAmelCase__      )

       __lowercase						,__lowercase								=		self.get_dummy_components()
       __lowercase								=		StableDiffusionPipeline(**UpperCAmelCase__      )
       __lowercase								=		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__      ):
             __lowercase								=		(
                 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    _lowercase					(     self				:     Dict      ):
    with tempfile.TemporaryDirectory() as tmpdirname:
       self.create_lora_weight_file(UpperCAmelCase__      )

       __lowercase						,__lowercase								=		self.get_dummy_components()
       __lowercase								=		StableDiffusionPipeline(**UpperCAmelCase__      )
       __lowercase								=		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    _lowercase					(     self				:     List[str]      ):
    __lowercase						,__lowercase								=		self.get_dummy_components()
    __lowercase								=		StableDiffusionPipeline(**UpperCAmelCase__      )
    __lowercase								=		sd_pipe.to(UpperCAmelCase__      )
    sd_pipe.set_progress_bar_config(disable=UpperCAmelCase__      )

    __lowercase						,__lowercase						,__lowercase								=		self.get_dummy_inputs()

    # enable XFormers
    sd_pipe.enable_xformers_memory_efficient_attention()

    __lowercase								=		sd_pipe(**UpperCAmelCase__      ).images
    __lowercase								=		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__      )

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

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







"""simple docstring"""





import math
from decimal import Decimal, getcontext




def     _A    (    UpperCamelCase_	:			int = 4000000)      ->		int:
   '''simple docstring'''


   try:
      __lowercase								=		int(UpperCamelCase_)
   except (TypeError, ValueError):
      raise TypeError("Parameter n must be int or castable to int.")
   if n <= 0:
      raise ValueError("Parameter n must be greater than or equal to one.")
   __lowercase								=		100
   __lowercase								=		(Decimal(5) ** Decimal(0.5) + 1) / Decimal(2)

   __lowercase								=		(math.floor(math.log(n * (phi + 2),     UpperCamelCase_) - 1) // 3) * 3 + 2
   __lowercase								=		Decimal(round(phi ** Decimal(index + 1))) / (phi + 2)
   __lowercase								=		num // 2
   return int(UpperCamelCase_)


if __name__ == "__main__":
   print(F"{solution() = }")




