





import inspect
import os
from argparse import ArgumentParser, Namespace
from importlib import import_module

import huggingface_hub
import numpy as np
from packaging import version

from .. import (
    FEATURE_EXTRACTOR_MAPPING,
    IMAGE_PROCESSOR_MAPPING,
    PROCESSOR_MAPPING,
    TOKENIZER_MAPPING,
    AutoConfig,
    AutoFeatureExtractor,
    AutoImageProcessor,
    AutoProcessor,
    AutoTokenizer,
    is_datasets_available,
    is_tf_available,
    is_torch_available,
)
from ..utils import TF2_WEIGHTS_INDEX_NAME, TF2_WEIGHTS_NAME, logging
from . import BaseTransformersCLICommand


if is_tf_available():
  import tensorflow as tf

  tf.config.experimental.enable_tensor_float_aa_execution(False)

if is_torch_available():
  import torch

if is_datasets_available():
  from datasets import load_dataset


__UpperCAmelCase						=      5e-5  # larger error tolerance than in our internal tests, to avoid flaky user-facing errors
def 			__UpperCamelCase				(			lowercase__	:       Namespace		)	->      Optional[int]:
    '''simple docstring'''





    return PTtoTFCommand(
        args.model_name		,    args.local_dir		,    args.max_error		,    args.new_weights		,    args.no_pr		,    args.push		,    args.extra_commit_description		,    args.override_model_class		,    )

class 		__a							(	__UpperCamelCase      ):



       @staticmethod
       def  A							(     UpperCAmelCase	:   ArgumentParser			):
           lowerCAmelCase_					:	List[Any]						=   parser.add_parser(
               """pt-to-tf"""					,			help=(
                   """CLI tool to run convert a transformers model from a PyTorch checkpoint to a TensorFlow checkpoint."""
                   """ Can also be used to validate existing weights without opening PRs, with --no-pr."""
               )					,			)
           train_parser.add_argument(
               """--model-name"""					,			type=UpperCAmelCase					,			required=UpperCAmelCase					,			help="""The model name, including owner/organization, as seen on the hub."""					,			)
           train_parser.add_argument(
               """--local-dir"""					,			type=UpperCAmelCase					,			default=""""""					,			help="""Optional local directory of the model repository. Defaults to /tmp/{model_name}"""					,			)
           train_parser.add_argument(
               """--max-error"""					,			type=UpperCAmelCase					,			default=UpperCAmelCase					,			help=(
                   F'Maximum error tolerance. Defaults to {MAX_ERROR}. This flag should be avoided, use at your own risk.'
               )					,			)
           train_parser.add_argument(
               """--new-weights"""					,			action="""store_true"""					,			help="""Optional flag to create new TensorFlow weights, even if they already exist."""					,			)
           train_parser.add_argument(
               """--no-pr"""					,			action="""store_true"""					,			help="""Optional flag to NOT open a PR with converted weights."""			)
           train_parser.add_argument(
               """--push"""					,			action="""store_true"""					,			help="""Optional flag to push the weights directly to `main` (requires permissions)"""					,			)
           train_parser.add_argument(
               """--extra-commit-description"""					,			type=UpperCAmelCase					,			default=""""""					,			help="""Optional additional commit description to use when opening a PR (e.g. to tag the owner)."""					,			)
           train_parser.add_argument(
               """--override-model-class"""					,			type=UpperCAmelCase					,			default=UpperCAmelCase					,			help="""If you think you know better than the auto-detector, you can specify the model class here. """
               """Can be either an AutoModel class or a specific model class like BertForSequenceClassification."""					,			)
           train_parser.set_defaults(func=UpperCAmelCase			)



       @staticmethod
       def  A							(     UpperCAmelCase	:   Dict					,			UpperCAmelCase	:   Dict			):
           lowerCAmelCase_					:	Any						=   set(pt_outputs.keys()			)
           lowerCAmelCase_					:	Tuple						=   set(tf_outputs.keys()			)
           if pt_out_attrs != tf_out_attrs:
               raise ValueError(
                   F'The model outputs have different attributes, aborting. (Pytorch: {pt_out_attrs}, TensorFlow:'
                   F' {tf_out_attrs})'			)

           # 2. For each output attribute, computes the difference
           def _find_pt_tf_differences(UpperCAmelCase	:   int					,			UpperCAmelCase	:   str					,			UpperCAmelCase	:   List[str]					,			UpperCAmelCase	:   int=""			):
               # If the current attribute is a tensor, it is a leaf and we make the comparison. Otherwise, we will dig in
               # recursivelly, keeping the name of the attribute.
               if isinstance(UpperCAmelCase					,			torch.Tensor			):
                   lowerCAmelCase_					:	int						=   np.max(np.abs(pt_out.numpy() - tf_out.numpy()			)			)
                   lowerCAmelCase_					:	Tuple						=   tensor_difference
               else:
                   lowerCAmelCase_					:	Any						=   attr_name
                   for i, pt_item in enumerate(UpperCAmelCase			):
                       # If it is a named attribute, we keep the name. Otherwise, just its index.
                       if isinstance(UpperCAmelCase					,			UpperCAmelCase			):
                           lowerCAmelCase_					:	str						=   root_name + pt_item
                           lowerCAmelCase_					:	Dict						=   tf_out[pt_item]
                           lowerCAmelCase_					:	List[str]						=   pt_out[pt_item]
                       else:
                           lowerCAmelCase_					:	Union[str, Any]						=   root_name + F'[{i}]'
                           lowerCAmelCase_					:	List[Any]						=   tf_out[i]
                       lowerCAmelCase_					:	Optional[int]						=   _find_pt_tf_differences(UpperCAmelCase					,			UpperCAmelCase					,			UpperCAmelCase					,			UpperCAmelCase			)

               return differences

           return _find_pt_tf_differences(UpperCAmelCase					,			UpperCAmelCase					,			{}			)



       def __init__(     self	:   str					,			UpperCAmelCase	:   str					,			UpperCAmelCase	:   str					,			UpperCAmelCase	:   float					,			UpperCAmelCase	:   bool					,			UpperCAmelCase	:   bool					,			UpperCAmelCase	:   bool					,			UpperCAmelCase	:   str					,			UpperCAmelCase	:   str					,			*UpperCAmelCase	:   Any					,			):
           lowerCAmelCase_					:	Dict						=   logging.get_logger("""transformers-cli/pt_to_tf"""			)
           lowerCAmelCase_					:	Optional[int]						=   model_name
           lowerCAmelCase_					:	str						=   local_dir if local_dir else os.path.join("""/tmp"""					,			UpperCAmelCase			)
           lowerCAmelCase_					:	Optional[int]						=   max_error
           lowerCAmelCase_					:	int						=   new_weights
           lowerCAmelCase_					:	int						=   no_pr
           lowerCAmelCase_					:	int						=   push
           lowerCAmelCase_					:	Optional[Any]						=   extra_commit_description
           lowerCAmelCase_					:	str						=   override_model_class



       def  A							(     self	:   Optional[int]					,			UpperCAmelCase	:   List[Any]					,			UpperCAmelCase	:   str					,			UpperCAmelCase	:   int			):

           def _get_audio_input():
               lowerCAmelCase_					:	Optional[Any]						=   load_dataset("""hf-internal-testing/librispeech_asr_dummy"""					,			"""clean"""					,			split="""validation"""			)
               lowerCAmelCase_					:	Tuple						=   ds.sort("""id"""			).select(range(2			)			)[:2]["""audio"""]
               lowerCAmelCase_					:	Tuple						=   [x["""array"""] for x in speech_samples]
               return raw_samples

           lowerCAmelCase_					:	Dict						=   type(pt_model.config			)
           if model_config_class in PROCESSOR_MAPPING:
               lowerCAmelCase_					:	Tuple						=   AutoProcessor.from_pretrained(self._local_dir			)
               if model_config_class in TOKENIZER_MAPPING and processor.tokenizer.pad_token is None:
                   lowerCAmelCase_					:	Tuple						=   processor.tokenizer.eos_token
           elif model_config_class in IMAGE_PROCESSOR_MAPPING:
               lowerCAmelCase_					:	Optional[Any]						=   AutoImageProcessor.from_pretrained(self._local_dir			)
           elif model_config_class in FEATURE_EXTRACTOR_MAPPING:
               lowerCAmelCase_					:	List[Any]						=   AutoFeatureExtractor.from_pretrained(self._local_dir			)
           elif model_config_class in TOKENIZER_MAPPING:
               lowerCAmelCase_					:	int						=   AutoTokenizer.from_pretrained(self._local_dir			)
               if processor.pad_token is None:
                   lowerCAmelCase_					:	Optional[int]						=   processor.eos_token
           else:
               raise ValueError(F'Unknown data processing type (model config type: {model_config_class})'			)

           lowerCAmelCase_					:	Union[str, Any]						=   set(inspect.signature(pt_model.forward			).parameters.keys()			)
           lowerCAmelCase_					:	Tuple						=   {}
           if "input_ids" in model_forward_signature:
               processor_inputs.update(
                   {
                       """text""": ["""Hi there!""", """I am a batch with more than one row and different input lengths."""],
                       """padding""": True,
                       """truncation""": True,
                   }			)
           if "pixel_values" in model_forward_signature:
               lowerCAmelCase_					:	str						=   load_dataset("""cifar10"""					,			"""plain_text"""					,			split="""test"""			)[:2]["""img"""]
               processor_inputs.update({"""images""": sample_images}			)
           if "input_features" in model_forward_signature:
               lowerCAmelCase_					:	Any						=   inspect.signature(processor.feature_extractor			).parameters
               # Pad to the largest input length by default but take feature extractor default
               # padding value if it exists e.g. "max_length" and is not False or None
               if "padding" in feature_extractor_signature:
                   lowerCAmelCase_					:	Dict						=   feature_extractor_signature["""padding"""].default
                   if default_strategy is not False and default_strategy is not None:
                       lowerCAmelCase_					:	Optional[Any]						=   default_strategy
                   else:
                       lowerCAmelCase_					:	int						=   True
               else:
                   lowerCAmelCase_					:	Optional[int]						=   True
               processor_inputs.update({"""audio""": _get_audio_input(), """padding""": padding_strategy}			)
           if "input_values" in model_forward_signature:  # Wav2Vec2 audio input
               processor_inputs.update({"""audio""": _get_audio_input(), """padding""": True}			)
           lowerCAmelCase_					:	Dict						=   processor(**UpperCAmelCase					,			return_tensors="""pt"""			)
           lowerCAmelCase_					:	Any						=   processor(**UpperCAmelCase					,			return_tensors="""tf"""			)

           # Extra input requirements, in addition to the input modality
           if (
               config.is_encoder_decoder
               or (hasattr(UpperCAmelCase					,			"""encoder"""			) and hasattr(UpperCAmelCase					,			"""decoder"""			))
               or "decoder_input_ids" in tf_dummy_inputs
           ):
               lowerCAmelCase_					:	List[Any]						=   np.asarray([[1], [1]]					,			dtype=UpperCAmelCase			) * (pt_model.config.decoder_start_token_id or 0)
               pt_input.update({"""decoder_input_ids""": torch.tensor(UpperCAmelCase			)}			)
               tf_input.update({"""decoder_input_ids""": tf.convert_to_tensor(UpperCAmelCase			)}			)

           return pt_input, tf_input



       def  A							(     self	:   str			):
           # hub version 0.9.0 introduced the possibility of programmatically opening PRs with normal write tokens.
           if version.parse(huggingface_hub.__version__			) < version.parse("""0.9.0"""			):
               raise ImportError(
                   """The huggingface_hub version must be >= 0.9.0 to use this command. Please update your huggingface_hub"""
                   """ installation."""			)
           else:
               from huggingface_hub import Repository, create_commit
               from huggingface_hub._commit_api import CommitOperationAdd

           # Fetch remote data
           lowerCAmelCase_					:	Dict						=   Repository(local_dir=self._local_dir					,			clone_from=self._model_name			)

           # Load config and get the appropriate architecture -- the latter is needed to convert the head's weights
           lowerCAmelCase_					:	Optional[int]						=   AutoConfig.from_pretrained(self._local_dir			)
           lowerCAmelCase_					:	Union[str, Any]						=   config.architectures
           if self._override_model_class is not None:
               if self._override_model_class.startswith("""TF"""			):
                   lowerCAmelCase_					:	str						=   [self._override_model_class[2:]]
               else:
                   lowerCAmelCase_					:	Any						=   [self._override_model_class]
               try:
                   lowerCAmelCase_					:	Optional[int]						=   getattr(import_module("""transformers"""			)					,			architectures[0]			)
               except AttributeError:
                   raise ValueError(F'Model class {self._override_model_class} not found in transformers.'			)
               try:
                   lowerCAmelCase_					:	Optional[int]						=   getattr(import_module("""transformers"""			)					,			"""TF""" + architectures[0]			)
               except AttributeError:
                   raise ValueError(F'TF model class TF{self._override_model_class} not found in transformers.'			)
           elif architectures is None:  # No architecture defined -- use auto classes
               lowerCAmelCase_					:	Tuple						=   getattr(import_module("""transformers"""			)					,			"""AutoModel"""			)
               lowerCAmelCase_					:	int						=   getattr(import_module("""transformers"""			)					,			"""TFAutoModel"""			)
               self._logger.warning("""No detected architecture, using AutoModel/TFAutoModel"""			)
           else:  # Architecture defined -- use it
               if len(UpperCAmelCase			) > 1:
                   raise ValueError(F'More than one architecture was found, aborting. (architectures = {architectures})'			)
               self._logger.warning(F'Detected architecture: {architectures[0]}'			)
               lowerCAmelCase_					:	Union[str, Any]						=   getattr(import_module("""transformers"""			)					,			architectures[0]			)
               try:
                   lowerCAmelCase_					:	Union[str, Any]						=   getattr(import_module("""transformers"""			)					,			"""TF""" + architectures[0]			)
               except AttributeError:
                   raise AttributeError(F'The TensorFlow equivalent of {architectures[0]} doesn\'t exist in transformers.'			)

        # Check the TF dummy inputs to see what keys we need in the forward pass
           lowerCAmelCase_					:	str						=   tf_class.from_config(UpperCAmelCase			)
           lowerCAmelCase_					:	Dict						=   tf_from_pt_model.dummy_inputs

           del tf_from_pt_model  # Try to keep only one model in memory at a time

           # Load the model and get some basic inputs
           lowerCAmelCase_					:	Union[str, Any]						=   pt_class.from_pretrained(self._local_dir			)
           pt_model.eval()

           lowerCAmelCase_		,		lowerCAmelCase_					:	Tuple						=   self.get_inputs(UpperCAmelCase					,			UpperCAmelCase					,			UpperCAmelCase			)

           with torch.no_grad():
               lowerCAmelCase_					:	Union[str, Any]						=   pt_model(**UpperCAmelCase					,			output_hidden_states=UpperCAmelCase			)
           del pt_model  # will no longer be used, and may have a large memory footprint

           lowerCAmelCase_					:	Tuple						=   tf_class.from_pretrained(self._local_dir					,			from_pt=UpperCAmelCase			)
           lowerCAmelCase_					:	str						=   tf_from_pt_model(**UpperCAmelCase					,			output_hidden_states=UpperCAmelCase					,			training=UpperCAmelCase			)

           # Confirms that cross loading PT weights into TF worked.
           lowerCAmelCase_					:	Dict						=   self.find_pt_tf_differences(UpperCAmelCase					,			UpperCAmelCase			)
           lowerCAmelCase_					:	str						=   {k: v for k, v in crossload_differences.items() if """hidden""" not in k}
           lowerCAmelCase_					:	Tuple						=   {k: v for k, v in crossload_differences.items() if """hidden""" in k}
           if len(UpperCAmelCase			) == 0 and architectures is not None:
               raise ValueError(
                   F'Something went wrong -- the config file has architectures ({architectures}), but no model head'
                   """ output was found. All outputs start with 'hidden'"""			)
           lowerCAmelCase_					:	List[Any]						=   max(output_differences.values()			) if output_differences else 0.0
           lowerCAmelCase_					:	Any						=   max(hidden_differences.values()			)
           if max_crossload_output_diff > self._max_error or max_crossload_hidden_diff > self._max_error:
               raise ValueError(
                   """The cross-loaded TensorFlow model has different outputs, something went wrong!\n"""
                   + F'\nList of maximum output differences above the threshold ({self._max_error}):\n'
                   + """\n""".join([F'{k}: {v:.3e}' for k, v in output_differences.items() if v > self._max_error]			)
                   + F'\n\nList of maximum hidden layer differences above the threshold ({self._max_error}):\n'
                   + """\n""".join([F'{k}: {v:.3e}' for k, v in hidden_differences.items() if v > self._max_error]			)			)

           # Save the weights in a TF format (if needed) and confirms that the results are still good
           lowerCAmelCase_					:	str						=   os.path.join(self._local_dir					,			UpperCAmelCase			)
           lowerCAmelCase_					:	Optional[Any]						=   os.path.join(self._local_dir					,			UpperCAmelCase			)
           if (not os.path.exists(UpperCAmelCase			) and not os.path.exists(UpperCAmelCase			)) or self._new_weights:
               tf_from_pt_model.save_pretrained(self._local_dir			)
           del tf_from_pt_model  # will no longer be used, and may have a large memory footprint

           lowerCAmelCase_					:	Optional[int]						=   tf_class.from_pretrained(self._local_dir			)
           lowerCAmelCase_					:	Dict						=   tf_model(**UpperCAmelCase					,			output_hidden_states=UpperCAmelCase			)

           lowerCAmelCase_					:	str						=   self.find_pt_tf_differences(UpperCAmelCase					,			UpperCAmelCase			)
           lowerCAmelCase_					:	Optional[int]						=   {k: v for k, v in conversion_differences.items() if """hidden""" not in k}
           lowerCAmelCase_					:	Optional[int]						=   {k: v for k, v in conversion_differences.items() if """hidden""" in k}
           if len(UpperCAmelCase			) == 0 and architectures is not None:
               raise ValueError(
                   F'Something went wrong -- the config file has architectures ({architectures}), but no model head'
                   """ output was found. All outputs start with 'hidden'"""			)
           lowerCAmelCase_					:	Optional[int]						=   max(output_differences.values()			) if output_differences else 0.0
           lowerCAmelCase_					:	List[Any]						=   max(hidden_differences.values()			)
           if max_conversion_output_diff > self._max_error or max_conversion_hidden_diff > self._max_error:
               raise ValueError(
                   """The converted TensorFlow model has different outputs, something went wrong!\n"""
                   + F'\nList of maximum output differences above the threshold ({self._max_error}):\n'
                   + """\n""".join([F'{k}: {v:.3e}' for k, v in output_differences.items() if v > self._max_error]			)
                   + F'\n\nList of maximum hidden layer differences above the threshold ({self._max_error}):\n'
                   + """\n""".join([F'{k}: {v:.3e}' for k, v in hidden_differences.items() if v > self._max_error]			)			)

           lowerCAmelCase_					:	Dict						=   """Update TF weights""" if self._new_weights else """Add TF weights"""
           if self._push:
               repo.git_add(auto_lfs_track=UpperCAmelCase			)
               repo.git_commit(UpperCAmelCase			)
               repo.git_push(blocking=UpperCAmelCase			)  # this prints a progress bar with the upload
               self._logger.warning(F'TF weights pushed into {self._model_name}'			)
           elif not self._no_pr:
               self._logger.warning("""Uploading the weights into a new PR..."""			)
               lowerCAmelCase_					:	Dict						=   (
                   """Model converted by the [`transformers`' `pt_to_tf`"""
                   """ CLI](https://github.com/huggingface/transformers/blob/main/src/transformers/commands/pt_to_tf.py). """
                   """All converted model outputs and hidden layers were validated against its PyTorch counterpart.\n\n"""
                   F'Maximum crossload output difference={max_crossload_output_diff:.3e}; '
                   F'Maximum crossload hidden layer difference={max_crossload_hidden_diff:.3e};\n'
                   F'Maximum conversion output difference={max_conversion_output_diff:.3e}; '
                   F'Maximum conversion hidden layer difference={max_conversion_hidden_diff:.3e};\n'
               )
               if self._max_error > MAX_ERROR:
                   commit_descrition += (
                       F'\n\nCAUTION: The maximum admissible error was manually increased to {self._max_error}!'
                   )
               if self._extra_commit_description:
                   commit_descrition += "\n\n" + self._extra_commit_description

               # sharded model -> adds all related files (index and .h5 shards)
               if os.path.exists(UpperCAmelCase			):
                   lowerCAmelCase_					:	Optional[Any]						=   [
                       CommitOperationAdd(path_in_repo=UpperCAmelCase					,			path_or_fileobj=UpperCAmelCase			)
                   ]
                   for shard_path in tf.io.gfile.glob(self._local_dir + """/tf_model-*.h5"""			):
                       operations += [
                           CommitOperationAdd(path_in_repo=os.path.basename(UpperCAmelCase			)					,			path_or_fileobj=UpperCAmelCase			)
                       ]
               else:
                   lowerCAmelCase_					:	Optional[Any]						=   [CommitOperationAdd(path_in_repo=UpperCAmelCase					,			path_or_fileobj=UpperCAmelCase			)]

               lowerCAmelCase_					:	int						=   create_commit(
                   repo_id=self._model_name					,			operations=UpperCAmelCase					,			commit_message=UpperCAmelCase					,			commit_description=UpperCAmelCase					,			repo_type="""model"""					,			create_pr=UpperCAmelCase					,			).pr_url
               self._logger.warning(F'PR open in {hub_pr_url}'			)









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 			__UpperCamelCase				(			lowercase__	:       nn.Module		)	->      str:
    '''simple docstring'''





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





    lowerCAmelCase_					:	Optional[int]						=   {}
    lowerCAmelCase_					:	str						=   (
        LoRAAttnProcessor2_0 if hasattr(lowercase__		,    """scaled_dot_product_attention"""		) else LoRAAttnProcessor
    )
    for name, module in text_encoder.named_modules():
        if name.endswith(lowercase__		):
            lowerCAmelCase_					:	Optional[int]						=   lora_attn_processor_class(
                hidden_size=module.out_proj.out_features		,    cross_attention_dim=lowercase__		)
    return text_lora_attn_procs
def 			__UpperCamelCase				(			lowercase__	:       nn.Module		)	->      List[str]:
    '''simple docstring'''





    lowerCAmelCase_					:	int						=   create_text_encoder_lora_attn_procs(lowercase__		)
    lowerCAmelCase_					:	str						=   AttnProcsLayers(lowercase__		)
    return text_encoder_lora_layers
def 			__UpperCamelCase				(			lowercase__	:       int		,    lowercase__	:       Union[str, Any]=False		)	->      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"""		):
                lowerCAmelCase_					:	Union[str, Any]						=   (
                    torch.randn_like(layer_module.up.weight		)
                    if randn_weight
                    else torch.zeros_like(layer_module.up.weight		)
                )
                lowerCAmelCase_					:	Optional[Any]						=   torch.nn.Parameter(lowercase__		)

class 		__a							(	unittest.TestCase      ):



       def  A							(     self	:   List[Any]			):
           torch.manual_seed(0			)
           lowerCAmelCase_					:	Any						=   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					,			)
           lowerCAmelCase_					:	Dict						=   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			)
           lowerCAmelCase_					:	List[str]						=   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					,			)
           lowerCAmelCase_					:	int						=   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=10_00					,			)
           lowerCAmelCase_					:	str						=   CLIPTextModel(UpperCAmelCase			)
           lowerCAmelCase_					:	Union[str, Any]						=   CLIPTokenizer.from_pretrained("""hf-internal-testing/tiny-random-clip"""			)

           lowerCAmelCase_		,		lowerCAmelCase_					:	Union[str, Any]						=   create_unet_lora_layers(UpperCAmelCase			)
           lowerCAmelCase_					:	int						=   create_text_encoder_lora_layers(UpperCAmelCase			)

           lowerCAmelCase_					:	str						=   {
               """unet""": unet,
               """scheduler""": scheduler,
               """vae""": vae,
               """text_encoder""": text_encoder,
               """tokenizer""": tokenizer,
               """safety_checker""": None,
               """feature_extractor""": None,
           }
           lowerCAmelCase_					:	List[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  A							(     self	:   Tuple			):
           lowerCAmelCase_					:	Dict						=   1
           lowerCAmelCase_					:	Optional[Any]						=   10
           lowerCAmelCase_					:	List[str]						=   4
           lowerCAmelCase_					:	Tuple						=   (32, 32)

           lowerCAmelCase_					:	int						=   torch.manual_seed(0			)
           lowerCAmelCase_					:	str						=   floats_tensor((batch_size, num_channels) + sizes			)
           lowerCAmelCase_					:	int						=   torch.randint(1					,			UpperCAmelCase					,			size=(batch_size, sequence_length)					,			generator=UpperCAmelCase			)

           lowerCAmelCase_					:	Union[str, Any]						=   {
               """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  A							(     self	:   Dict			):
           lowerCAmelCase_					:	Union[str, Any]						=   77

           lowerCAmelCase_					:	List[Any]						=   torch.randint(2					,			56					,			size=(1, max_seq_length)					,			generator=torch.manual_seed(0			)			)

           lowerCAmelCase_					:	List[Any]						=   {}
           lowerCAmelCase_					:	Optional[Any]						=   inputs
           return prepared_inputs



       def  A							(     self	:   List[str]					,			UpperCAmelCase	:   List[str]			):
           lowerCAmelCase_		,		lowerCAmelCase_					:	List[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  A							(     self	:   Optional[int]			):
           lowerCAmelCase_		,		lowerCAmelCase_					:	Optional[Any]						=   self.get_dummy_components()
           lowerCAmelCase_					:	Optional[Any]						=   StableDiffusionPipeline(**UpperCAmelCase			)
           lowerCAmelCase_					:	str						=   sd_pipe.to(UpperCAmelCase			)
           sd_pipe.set_progress_bar_config(disable=UpperCAmelCase			)

           lowerCAmelCase_		,		lowerCAmelCase_		,		lowerCAmelCase_					:	List[Any]						=   self.get_dummy_inputs()

           lowerCAmelCase_					:	Tuple						=   sd_pipe(**UpperCAmelCase			).images
           lowerCAmelCase_					:	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"""]					,			)
               self.assertTrue(os.path.isfile(os.path.join(UpperCAmelCase					,			"""pytorch_lora_weights.bin"""			)			)			)
               sd_pipe.load_lora_weights(UpperCAmelCase			)

           lowerCAmelCase_					:	Any						=   sd_pipe(**UpperCAmelCase			).images
           lowerCAmelCase_					:	Any						=   lora_images[0, -3:, -3:, -1]

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



       def  A							(     self	:   Tuple			):
           lowerCAmelCase_		,		lowerCAmelCase_					:	Optional[Any]						=   self.get_dummy_components()
           lowerCAmelCase_					:	Union[str, Any]						=   StableDiffusionPipeline(**UpperCAmelCase			)
           lowerCAmelCase_					:	Optional[Any]						=   sd_pipe.to(UpperCAmelCase			)
           sd_pipe.set_progress_bar_config(disable=UpperCAmelCase			)

           lowerCAmelCase_		,		lowerCAmelCase_		,		lowerCAmelCase_					:	Optional[Any]						=   self.get_dummy_inputs()

           lowerCAmelCase_					:	Optional[int]						=   sd_pipe(**UpperCAmelCase			).images
           lowerCAmelCase_					:	int						=   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			)

           lowerCAmelCase_					:	Any						=   sd_pipe(**UpperCAmelCase			).images
           lowerCAmelCase_					:	List[str]						=   lora_images[0, -3:, -3:, -1]

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



       def  A							(     self	:   Dict			):
           lowerCAmelCase_		,		lowerCAmelCase_					:	Optional[Any]						=   self.get_dummy_components()
           lowerCAmelCase_					:	Any						=   lora_components["""unet_lora_attn_procs"""]
           lowerCAmelCase_					:	str						=   StableDiffusionPipeline(**UpperCAmelCase			)
           lowerCAmelCase_					:	Union[str, Any]						=   sd_pipe.to(UpperCAmelCase			)
           sd_pipe.set_progress_bar_config(disable=UpperCAmelCase			)

           lowerCAmelCase_		,		lowerCAmelCase_		,		lowerCAmelCase_					:	Dict						=   self.get_dummy_inputs()

           lowerCAmelCase_					:	int						=   sd_pipe(**UpperCAmelCase			).images
           lowerCAmelCase_					:	int						=   original_images[0, -3:, -3:, -1]

           with tempfile.TemporaryDirectory() as tmpdirname:
               lowerCAmelCase_					:	List[Any]						=   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			)

           lowerCAmelCase_					:	List[str]						=   sd_pipe(**UpperCAmelCase			).images
           lowerCAmelCase_					:	List[str]						=   lora_images[0, -3:, -3:, -1]

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



       def  A							(     self	:   int			):
           lowerCAmelCase_		,		lowerCAmelCase_					:	Dict						=   self.get_dummy_components()
           lowerCAmelCase_					:	Dict						=   StableDiffusionPipeline(**UpperCAmelCase			)

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

           # inference without lora
           lowerCAmelCase_					:	List[Any]						=   pipe.text_encoder(**UpperCAmelCase			)[0]
           assert outputs_without_lora.shape == (1, 77, 32)

           # create lora_attn_procs with zeroed out up.weights
           lowerCAmelCase_					:	Tuple						=   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
           lowerCAmelCase_					:	str						=   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
           lowerCAmelCase_					:	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
           lowerCAmelCase_					:	Optional[int]						=   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  A							(     self	:   Union[str, Any]			):
           lowerCAmelCase_		,		lowerCAmelCase_					:	Optional[int]						=   self.get_dummy_components()
           lowerCAmelCase_					:	List[Any]						=   StableDiffusionPipeline(**UpperCAmelCase			)

           lowerCAmelCase_					:	Optional[Any]						=   self.get_dummy_tokens()

           # inference without lora
           lowerCAmelCase_					:	Optional[Any]						=   pipe.text_encoder(**UpperCAmelCase			)[0]
           assert outputs_without_lora.shape == (1, 77, 32)

           # create lora_attn_procs with randn up.weights
           lowerCAmelCase_					:	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
           lowerCAmelCase_					:	List[str]						=   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
           lowerCAmelCase_					:	Tuple						=   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  A							(     self	:   Any			):
           lowerCAmelCase_		,		lowerCAmelCase_					:	Any						=   self.get_dummy_components()
           lowerCAmelCase_					:	int						=   StableDiffusionPipeline(**UpperCAmelCase			)
           lowerCAmelCase_					:	int						=   sd_pipe.to(UpperCAmelCase			)
           sd_pipe.set_progress_bar_config(disable=UpperCAmelCase			)

           lowerCAmelCase_		,		lowerCAmelCase_		,		lowerCAmelCase_					:	Optional[int]						=   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			)

           lowerCAmelCase_					:	Any						=   sd_pipe(**UpperCAmelCase			).images
           lowerCAmelCase_					:	Optional[int]						=   lora_images[0, -3:, -3:, -1]

           lowerCAmelCase_					:	Dict						=   sd_pipe(**UpperCAmelCase					,			cross_attention_kwargs={"""scale""": 0.5}			).images
           lowerCAmelCase_					:	str						=   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  A							(     self	:   List[str]			):
           with tempfile.TemporaryDirectory() as tmpdirname:
               self.create_lora_weight_file(UpperCAmelCase			)

               lowerCAmelCase_		,		lowerCAmelCase_					:	List[str]						=   self.get_dummy_components()
               lowerCAmelCase_					:	int						=   StableDiffusionPipeline(**UpperCAmelCase			)
               lowerCAmelCase_					:	Union[str, 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			):
                       lowerCAmelCase_					:	Optional[int]						=   (
                           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  A							(     self	:   Any			):
           with tempfile.TemporaryDirectory() as tmpdirname:
               self.create_lora_weight_file(UpperCAmelCase			)

               lowerCAmelCase_		,		lowerCAmelCase_					:	List[str]						=   self.get_dummy_components()
               lowerCAmelCase_					:	Any						=   StableDiffusionPipeline(**UpperCAmelCase			)
               lowerCAmelCase_					:	List[str]						=   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  A							(     self	:   List[str]			):
           lowerCAmelCase_		,		lowerCAmelCase_					:	int						=   self.get_dummy_components()
           lowerCAmelCase_					:	Optional[int]						=   StableDiffusionPipeline(**UpperCAmelCase			)
           lowerCAmelCase_					:	Dict						=   sd_pipe.to(UpperCAmelCase			)
           sd_pipe.set_progress_bar_config(disable=UpperCAmelCase			)

           lowerCAmelCase_		,		lowerCAmelCase_		,		lowerCAmelCase_					:	List[str]						=   self.get_dummy_inputs()

           # enable XFormers
           sd_pipe.enable_xformers_memory_efficient_attention()

           lowerCAmelCase_					:	Any						=   sd_pipe(**UpperCAmelCase			).images
           lowerCAmelCase_					:	Union[str, Any]						=   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			)

           lowerCAmelCase_					:	Optional[Any]						=   sd_pipe(**UpperCAmelCase			).images
           lowerCAmelCase_					:	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			)			)			)


