





'''simple docstring'''






import gc
import unittest

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

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

from ..pipeline_params import TEXT_TO_IMAGE_BATCH_PARAMS, TEXT_TO_IMAGE_IMAGE_PARAMS, TEXT_TO_IMAGE_PARAMS
from ..test_pipelines_common import PipelineLatentTesterMixin, PipelineTesterMixin


enable_full_determinism()




class 				UpperCAmelCase  (    snake_case_							,	snake_case_							,	unittest.TestCase     ):
      _lowercase:    Dict		  =     StableDiffusionXLPipeline
      _lowercase:    Optional[int]		  =     TEXT_TO_IMAGE_PARAMS
      _lowercase:    str		  =     TEXT_TO_IMAGE_BATCH_PARAMS
      _lowercase:    List[Any]		  =     TEXT_TO_IMAGE_IMAGE_PARAMS
      _lowercase:    Tuple		  =     TEXT_TO_IMAGE_IMAGE_PARAMS




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

            _lowerCAmelCase							  =							CLIPTextModelWithProjection(__snake_case						)
            _lowerCAmelCase							  =							CLIPTokenizer.from_pretrained("""hf-internal-testing/tiny-random-clip"""    ,    local_files_only=__snake_case						)

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




      def        lowercase__					(    self    :				Optional[Any]    ,    __snake_case    :				Tuple    ,    __snake_case    :				Union[str, Any]=0						)       ->   List[Any]:
            if str(__snake_case						).startswith("""mps"""						):
                  _lowerCAmelCase							  =							torch.manual_seed(__snake_case						)
            else:
                  _lowerCAmelCase							  =							torch.Generator(device=__snake_case						).manual_seed(__snake_case						)
            _lowerCAmelCase							  =							{
                """prompt""": """A painting of a squirrel eating a burger""",
                """generator""": generator,
                """num_inference_steps""": 2,
                """guidance_scale""": 5.0,
                """output_type""": """numpy""",
            }
            return inputs




      def        lowercase__					(    self    :				str						)       ->   List[str]:
            _lowerCAmelCase							  =							"""cpu"""  # ensure determinism for the device-dependent torch.Generator
            _lowerCAmelCase							  =							self.get_dummy_components()
            _lowerCAmelCase							  =							StableDiffusionXLPipeline(**__snake_case						)
            _lowerCAmelCase							  =							sd_pipe.to(__snake_case						)
            sd_pipe.set_progress_bar_config(disable=__snake_case						)

            _lowerCAmelCase							  =							self.get_dummy_inputs(__snake_case						)
            _lowerCAmelCase							  =							sd_pipe(**__snake_case						).images
            _lowerCAmelCase							  =							image[0, -3:, -3:, -1]

            assert image.shape == (1, 64, 64, 3)
            _lowerCAmelCase							  =							np.array([0.58_73, 0.61_28, 0.47_97, 0.51_22, 0.56_74, 0.46_39, 0.52_27, 0.51_49, 0.47_47]						)

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




      def        lowercase__					(    self    :				Union[str, Any]						)       ->   Tuple:
            _lowerCAmelCase							  =							self.get_dummy_components()
            _lowerCAmelCase							  =							StableDiffusionXLPipeline(**__snake_case						)
            _lowerCAmelCase							  =							sd_pipe.to(__snake_case						)
            _lowerCAmelCase							  =							sd_pipe.to(__snake_case						)
            sd_pipe.set_progress_bar_config(disable=__snake_case						)

            # forward without prompt embeds
            _lowerCAmelCase							  =							self.get_dummy_inputs(__snake_case						)
            _lowerCAmelCase							  =							3 * ["""this is a negative prompt"""]
            _lowerCAmelCase							  =							negative_prompt
            _lowerCAmelCase							  =							3 * [inputs["""prompt"""]]

            _lowerCAmelCase							  =							sd_pipe(**__snake_case						)
            _lowerCAmelCase							  =							output.images[0, -3:, -3:, -1]

            # forward with prompt embeds
            _lowerCAmelCase							  =							self.get_dummy_inputs(__snake_case						)
            _lowerCAmelCase							  =							3 * ["""this is a negative prompt"""]
            _lowerCAmelCase							  =							3 * [inputs.pop("""prompt"""						)]

            (
                (
                _lowerCAmelCase
            )  ,							(
                _lowerCAmelCase
            )  ,							(
                _lowerCAmelCase
            )  ,							(
                _lowerCAmelCase
            )  ,							
            )							  =							sd_pipe.encode_prompt(__snake_case    ,    negative_prompt=__snake_case						)

            _lowerCAmelCase							  =							sd_pipe(
                **__snake_case    ,    prompt_embeds=__snake_case    ,    negative_prompt_embeds=__snake_case    ,    pooled_prompt_embeds=__snake_case    ,    negative_pooled_prompt_embeds=__snake_case    ,    )
            _lowerCAmelCase							  =							output.images[0, -3:, -3:, -1]

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




      def        lowercase__					(    self    :				Optional[Any]						)       ->   List[str]:
            super().test_attention_slicing_forward_pass(expected_max_diff=3E-3						)




      def        lowercase__					(    self    :				List[Any]						)       ->   List[Any]:
            super().test_inference_batch_single_identical(expected_max_diff=3E-3						)




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




      def        lowercase__					(    self    :				str						)       ->   List[Any]:
            super().tearDown()
            gc.collect()
            torch.cuda.empty_cache()




      def        lowercase__					(    self    :				str    ,    __snake_case    :				Any    ,    __snake_case    :				int="cpu"    ,    __snake_case    :				Union[str, Any]=torch.floataa    ,    __snake_case    :				Optional[Any]=0						)       ->   Tuple:
            _lowerCAmelCase							  =							torch.Generator(device=__snake_case						).manual_seed(__snake_case						)
            _lowerCAmelCase							  =							np.random.RandomState(__snake_case						).standard_normal((1, 4, 64, 64)						)
            _lowerCAmelCase							  =							torch.from_numpy(__snake_case						).to(device=__snake_case    ,    dtype=__snake_case						)
            _lowerCAmelCase							  =							{
                """prompt""": """a photograph of an astronaut riding a horse""",
                """latents""": latents,
                """generator""": generator,
                """num_inference_steps""": 3,
                """guidance_scale""": 7.5,
                """output_type""": """numpy""",
            }
            return inputs




      def        lowercase__					(    self    :				Dict						)       ->   str:
            _lowerCAmelCase							  =							DiffusionPipeline.from_pretrained("""stabilityai/stable-diffusion-2-base"""						)
            pipe.to(__snake_case						)
            pipe.set_progress_bar_config(disable=__snake_case						)

            _lowerCAmelCase							  =							self.get_inputs(__snake_case						)
            _lowerCAmelCase							  =							pipe(**__snake_case						).images
            _lowerCAmelCase							  =							image[0, -3:, -3:, -1].flatten()

            assert image.shape == (1, 5_12, 5_12, 3)
            _lowerCAmelCase							  =							np.array([0.4_94_93, 0.4_78_96, 0.4_07_98, 0.5_42_14, 0.5_32_12, 0.4_82_02, 0.4_76_56, 0.4_63_29, 0.4_85_06]						)
            assert np.abs(image_slice - expected_slice						).max() < 7E-3













'''simple docstring'''






import os
import tempfile
import unittest

import numpy as np
import torch

from diffusers.models import ModelMixin, UNetaDConditionModel
from diffusers.models.attention_processor import AttnProcessor, LoRAAttnProcessor
from diffusers.utils import (
    floats_tensor,
    logging,
    skip_mps,
    torch_device,
)
from diffusers.utils.import_utils import is_xformers_available
from diffusers.utils.testing_utils import enable_full_determinism

from .test_modeling_common import ModelTesterMixin, UNetTesterMixin


enable_full_determinism()

A__       :     str          =logging.get_logger(__name__)
def      UpperCamelCase__   (	lowerCAmelCase							,			lowerCAmelCase = True						):





      """simple docstring"""


      _lowerCAmelCase							  =							{}
      for name in model.attn_processors.keys():
            _lowerCAmelCase							  =							name.endswith("""attn2.processor"""						) and not (
                name.startswith("""transformer_in"""						) or """temp_attentions""" in name.split("""."""						)
            )
            _lowerCAmelCase							  =							model.config.cross_attention_dim if has_cross_attention else None
            if name.startswith("""mid_block"""						):
                  _lowerCAmelCase							  =							model.config.block_out_channels[-1]
            elif name.startswith("""up_blocks"""						):
                  _lowerCAmelCase							  =							int(name[len("""up_blocks."""						)]						)
                  _lowerCAmelCase							  =							list(reversed(model.config.block_out_channels						)						)[block_id]
            elif name.startswith("""down_blocks"""						):
                  _lowerCAmelCase							  =							int(name[len("""down_blocks."""						)]						)
                  _lowerCAmelCase							  =							model.config.block_out_channels[block_id]
            elif name.startswith("""transformer_in"""						):
                  # Note that the `8 * ...` comes from: https://github.com/huggingface/diffusers/blob/7139f0e874f10b2463caa8cbd585762a309d12d6/src/diffusers/models/unet_3d_condition.py#L148
                  _lowerCAmelCase							  =							8 * model.config.attention_head_dim

            _lowerCAmelCase							  =							LoRAAttnProcessor(hidden_size=lowerCAmelCase							,			cross_attention_dim=lowerCAmelCase						)
            _lowerCAmelCase							  =							lora_attn_procs[name].to(model.device						)

            if mock_weights:
                  # add 1 to weights to mock trained weights
                  with torch.no_grad():
                        lora_attn_procs[name].to_q_lora.up.weight += 1
                        lora_attn_procs[name].to_k_lora.up.weight += 1
                        lora_attn_procs[name].to_v_lora.up.weight += 1
                        lora_attn_procs[name].to_out_lora.up.weight += 1

      return lora_attn_procs





@skip_mps
class 				UpperCAmelCase  (    snake_case_							,	snake_case_							,	unittest.TestCase     ):
      _lowercase:    int		  =     UNetaDConditionModel
      _lowercase:    Optional[int]		  =     '''sample'''




      @property
      def        lowercase__					(    self    :				Any						)       ->   int:
            _lowerCAmelCase							  =							4
            _lowerCAmelCase							  =							4
            _lowerCAmelCase							  =							4
            _lowerCAmelCase							  =							(32, 32)

            _lowerCAmelCase							  =							floats_tensor((batch_size, num_channels, num_frames) + sizes						).to(__snake_case						)
            _lowerCAmelCase							  =							torch.tensor([10]						).to(__snake_case						)
            _lowerCAmelCase							  =							floats_tensor((batch_size, 4, 32)						).to(__snake_case						)

            return {"sample": noise, "timestep": time_step, "encoder_hidden_states": encoder_hidden_states}




      @property
      def        lowercase__					(    self    :				Dict						)       ->   Any:
            return (4, 4, 32, 32)




      @property
      def        lowercase__					(    self    :				List[str]						)       ->   Optional[int]:
            return (4, 4, 32, 32)




      def        lowercase__					(    self    :				Optional[Any]						)       ->   Optional[Any]:
            _lowerCAmelCase							  =							{
                """block_out_channels""": (32, 64),
                """down_block_types""": (
                    """CrossAttnDownBlock3D""",
                    """DownBlock3D""",
                ),
                """up_block_types""": ("""UpBlock3D""", """CrossAttnUpBlock3D"""),
                """cross_attention_dim""": 32,
                """attention_head_dim""": 8,
                """out_channels""": 4,
                """in_channels""": 4,
                """layers_per_block""": 1,
                """sample_size""": 32,
            }
            _lowerCAmelCase							  =							self.dummy_input
            return init_dict, inputs_dict




      @unittest.skipIf(
          torch_device != """cuda""" or not is_xformers_available()    ,    reason="""XFormers attention is only available with CUDA and `xformers` installed"""    ,    )
      def        lowercase__					(    self    :				List[str]						)       ->   int:
            _lowerCAmelCase  ,							_lowerCAmelCase							  =							self.prepare_init_args_and_inputs_for_common()
            _lowerCAmelCase							  =							self.model_class(**__snake_case						)

            model.enable_xformers_memory_efficient_attention()

            assert (
                model.mid_block.attentions[0].transformer_blocks[0].attna.processor.__class__.__name__
                == "XFormersAttnProcessor"
            ), "xformers is not enabled"




      def        lowercase__					(    self    :				str						)       ->   int:
            _lowerCAmelCase  ,							_lowerCAmelCase							  =							self.prepare_init_args_and_inputs_for_common()

            _lowerCAmelCase							  =							32

            _lowerCAmelCase							  =							self.model_class(**__snake_case						)
            model.to(__snake_case						)
            model.eval()

            with torch.no_grad():
                  _lowerCAmelCase							  =							model(**__snake_case						)

                  if isinstance(__snake_case    ,    __snake_case						):
                        _lowerCAmelCase							  =							output.sample

            self.assertIsNotNone(__snake_case						)
            _lowerCAmelCase							  =							inputs_dict["""sample"""].shape
            self.assertEqual(output.shape    ,    __snake_case    ,    """Input and output shapes do not match"""						)




      def        lowercase__					(    self    :				Any						)       ->   Any:
            _lowerCAmelCase  ,							_lowerCAmelCase							  =							self.prepare_init_args_and_inputs_for_common()
            _lowerCAmelCase							  =							self.model_class(**__snake_case						)
            model.to(__snake_case						)
            model.eval()

            with torch.no_grad():
                  # Warmup pass when using mps (see #372)
                  if torch_device == "mps" and isinstance(__snake_case    ,    __snake_case						):
                        model(**self.dummy_input						)

                  _lowerCAmelCase							  =							model(**__snake_case						)
                  if isinstance(__snake_case    ,    __snake_case						):
                        _lowerCAmelCase							  =							first.sample

                  _lowerCAmelCase							  =							model(**__snake_case						)
                  if isinstance(__snake_case    ,    __snake_case						):
                        _lowerCAmelCase							  =							second.sample

            _lowerCAmelCase							  =							first.cpu().numpy()
            _lowerCAmelCase							  =							second.cpu().numpy()
            _lowerCAmelCase							  =							out_a[~np.isnan(__snake_case						)]
            _lowerCAmelCase							  =							out_a[~np.isnan(__snake_case						)]
            _lowerCAmelCase							  =							np.amax(np.abs(out_a - out_a						)						)
            self.assertLessEqual(__snake_case    ,    1E-5						)




      def        lowercase__					(    self    :				Any						)       ->   List[str]:
            _lowerCAmelCase  ,							_lowerCAmelCase							  =							self.prepare_init_args_and_inputs_for_common()

            _lowerCAmelCase							  =							8

            _lowerCAmelCase							  =							self.model_class(**__snake_case						)
            model.to(__snake_case						)
            model.eval()

            model.set_attention_slice("""auto"""						)
            with torch.no_grad():
                  _lowerCAmelCase							  =							model(**__snake_case						)
            assert output is not None

            model.set_attention_slice("""max"""						)
            with torch.no_grad():
                  _lowerCAmelCase							  =							model(**__snake_case						)
            assert output is not None

            model.set_attention_slice(2						)
            with torch.no_grad():
                  _lowerCAmelCase							  =							model(**__snake_case						)
            assert output is not None




      def        lowercase__					(    self    :				str						)       ->   Any:
            _lowerCAmelCase  ,							_lowerCAmelCase							  =							self.prepare_init_args_and_inputs_for_common()

            _lowerCAmelCase							  =							8

            _lowerCAmelCase							  =							self.model_class(**__snake_case						)
            model.to(__snake_case						)

            with torch.no_grad():
                  _lowerCAmelCase							  =							model(**__snake_case						).sample

            _lowerCAmelCase							  =							create_lora_layers(__snake_case						)

            # make sure we can set a list of attention processors
            model.set_attn_processor(__snake_case						)
            model.to(__snake_case						)

            # test that attn processors can be set to itself
            model.set_attn_processor(model.attn_processors						)

            with torch.no_grad():
                  _lowerCAmelCase							  =							model(**__snake_case    ,    cross_attention_kwargs={"""scale""": 0.0}						).sample
                  _lowerCAmelCase							  =							model(**__snake_case    ,    cross_attention_kwargs={"""scale""": 0.5}						).sample
                  _lowerCAmelCase							  =							model(**__snake_case    ,    cross_attention_kwargs={"""scale""": 0.5}						).sample

            assert (samplea - samplea).abs().max() < 3E-3
            assert (samplea - samplea).abs().max() < 3E-3

            # sample 2 and sample 3 should be different
            assert (samplea - samplea).abs().max() > 3E-3




      def        lowercase__					(    self    :				Tuple						)       ->   int:
            _lowerCAmelCase  ,							_lowerCAmelCase							  =							self.prepare_init_args_and_inputs_for_common()

            _lowerCAmelCase							  =							8

            torch.manual_seed(0						)
            _lowerCAmelCase							  =							self.model_class(**__snake_case						)
            model.to(__snake_case						)

            with torch.no_grad():
                  _lowerCAmelCase							  =							model(**__snake_case						).sample

            _lowerCAmelCase							  =							create_lora_layers(__snake_case						)
            model.set_attn_processor(__snake_case						)

            with torch.no_grad():
                  _lowerCAmelCase							  =							model(**__snake_case    ,    cross_attention_kwargs={"""scale""": 0.5}						).sample

            with tempfile.TemporaryDirectory() as tmpdirname:
                  model.save_attn_procs(__snake_case						)
                  self.assertTrue(os.path.isfile(os.path.join(__snake_case    ,    """pytorch_lora_weights.bin"""						)						)						)
                  torch.manual_seed(0						)
                  _lowerCAmelCase							  =							self.model_class(**__snake_case						)
                  new_model.to(__snake_case						)
                  new_model.load_attn_procs(__snake_case						)

            with torch.no_grad():
                  _lowerCAmelCase							  =							new_model(**__snake_case    ,    cross_attention_kwargs={"""scale""": 0.5}						).sample

            assert (sample - new_sample).abs().max() < 1E-3

            # LoRA and no LoRA should NOT be the same
            assert (sample - old_sample).abs().max() > 1E-4




      def        lowercase__					(    self    :				Optional[int]						)       ->   Optional[int]:
            _lowerCAmelCase  ,							_lowerCAmelCase							  =							self.prepare_init_args_and_inputs_for_common()

            _lowerCAmelCase							  =							8

            torch.manual_seed(0						)
            _lowerCAmelCase							  =							self.model_class(**__snake_case						)
            model.to(__snake_case						)

            with torch.no_grad():
                  _lowerCAmelCase							  =							model(**__snake_case						).sample

            _lowerCAmelCase							  =							create_lora_layers(__snake_case						)
            model.set_attn_processor(__snake_case						)

            with torch.no_grad():
                  _lowerCAmelCase							  =							model(**__snake_case    ,    cross_attention_kwargs={"""scale""": 0.5}						).sample

            with tempfile.TemporaryDirectory() as tmpdirname:
                  model.save_attn_procs(__snake_case    ,    safe_serialization=__snake_case						)
                  self.assertTrue(os.path.isfile(os.path.join(__snake_case    ,    """pytorch_lora_weights.safetensors"""						)						)						)
                  torch.manual_seed(0						)
                  _lowerCAmelCase							  =							self.model_class(**__snake_case						)
                  new_model.to(__snake_case						)
                  new_model.load_attn_procs(__snake_case						)

            with torch.no_grad():
                  _lowerCAmelCase							  =							new_model(**__snake_case    ,    cross_attention_kwargs={"""scale""": 0.5}						).sample

            assert (sample - new_sample).abs().max() < 3E-3

            # LoRA and no LoRA should NOT be the same
            assert (sample - old_sample).abs().max() > 1E-4




      def        lowercase__					(    self    :				Dict						)       ->   List[str]:
            # enable deterministic behavior for gradient checkpointing
            _lowerCAmelCase  ,							_lowerCAmelCase							  =							self.prepare_init_args_and_inputs_for_common()

            _lowerCAmelCase							  =							8

            torch.manual_seed(0						)
            _lowerCAmelCase							  =							self.model_class(**__snake_case						)
            model.to(__snake_case						)

            _lowerCAmelCase							  =							create_lora_layers(__snake_case    ,    mock_weights=__snake_case						)
            model.set_attn_processor(__snake_case						)
            # Saving as torch, properly reloads with directly filename
            with tempfile.TemporaryDirectory() as tmpdirname:
                  model.save_attn_procs(__snake_case						)
                  self.assertTrue(os.path.isfile(os.path.join(__snake_case    ,    """pytorch_lora_weights.bin"""						)						)						)
                  torch.manual_seed(0						)
                  _lowerCAmelCase							  =							self.model_class(**__snake_case						)
                  new_model.to(__snake_case						)
                  new_model.load_attn_procs(__snake_case    ,    weight_name="""pytorch_lora_weights.bin"""						)




      def        lowercase__					(    self    :				Optional[int]						)       ->   List[Any]:
            _lowerCAmelCase  ,							_lowerCAmelCase							  =							self.prepare_init_args_and_inputs_for_common()

            _lowerCAmelCase							  =							8

            torch.manual_seed(0						)
            _lowerCAmelCase							  =							self.model_class(**__snake_case						)
            model.to(__snake_case						)

            _lowerCAmelCase							  =							create_lora_layers(__snake_case    ,    mock_weights=__snake_case						)
            model.set_attn_processor(__snake_case						)
            # Saving as torch, properly reloads with directly filename
            with tempfile.TemporaryDirectory() as tmpdirname:
                  model.save_attn_procs(__snake_case						)
                  self.assertTrue(os.path.isfile(os.path.join(__snake_case    ,    """pytorch_lora_weights.bin"""						)						)						)
                  torch.manual_seed(0						)
                  _lowerCAmelCase							  =							self.model_class(**__snake_case						)
                  new_model.to(__snake_case						)
                  with self.assertRaises(__snake_case						) as e:
                        new_model.load_attn_procs(__snake_case    ,    use_safetensors=__snake_case						)
                  self.assertIn("""Error no file named pytorch_lora_weights.safetensors"""    ,    str(e.exception						)						)




      def        lowercase__					(    self    :				Any						)       ->   Any:
            _lowerCAmelCase  ,							_lowerCAmelCase							  =							self.prepare_init_args_and_inputs_for_common()

            _lowerCAmelCase							  =							8

            torch.manual_seed(0						)
            _lowerCAmelCase							  =							self.model_class(**__snake_case						)
            model.to(__snake_case						)

            with torch.no_grad():
                  _lowerCAmelCase							  =							model(**__snake_case						).sample

            _lowerCAmelCase							  =							create_lora_layers(__snake_case						)
            model.set_attn_processor(__snake_case						)

            with torch.no_grad():
                  _lowerCAmelCase							  =							model(**__snake_case    ,    cross_attention_kwargs={"""scale""": 0.0}						).sample

            model.set_attn_processor(AttnProcessor()						)

            with torch.no_grad():
                  _lowerCAmelCase							  =							model(**__snake_case						).sample

            assert (sample - new_sample).abs().max() < 1E-4
            assert (sample - old_sample).abs().max() < 3E-3




      @unittest.skipIf(
          torch_device != """cuda""" or not is_xformers_available()    ,    reason="""XFormers attention is only available with CUDA and `xformers` installed"""    ,    )
      def        lowercase__					(    self    :				Union[str, Any]						)       ->   Tuple:
            # enable deterministic behavior for gradient checkpointing
            _lowerCAmelCase  ,							_lowerCAmelCase							  =							self.prepare_init_args_and_inputs_for_common()

            _lowerCAmelCase							  =							4

            torch.manual_seed(0						)
            _lowerCAmelCase							  =							self.model_class(**__snake_case						)
            model.to(__snake_case						)
            _lowerCAmelCase							  =							create_lora_layers(__snake_case						)
            model.set_attn_processor(__snake_case						)

            # default
            with torch.no_grad():
                  _lowerCAmelCase							  =							model(**__snake_case						).sample

                  model.enable_xformers_memory_efficient_attention()
                  _lowerCAmelCase							  =							model(**__snake_case						).sample

                  model.disable_xformers_memory_efficient_attention()
                  _lowerCAmelCase							  =							model(**__snake_case						).sample

            assert (sample - on_sample).abs().max() < 1E-4
            assert (sample - off_sample).abs().max() < 1E-4




      def        lowercase__					(    self    :				Dict						)       ->   Optional[int]:
            _lowerCAmelCase  ,							_lowerCAmelCase							  =							self.prepare_init_args_and_inputs_for_common()
            _lowerCAmelCase							  =							32

            _lowerCAmelCase							  =							self.model_class(**__snake_case						)
            model.to(__snake_case						)
            model.eval()

            with torch.no_grad():
                  _lowerCAmelCase							  =							model(**__snake_case						)[0]

            model.enable_forward_chunking()
            with torch.no_grad():
                  _lowerCAmelCase							  =							model(**__snake_case						)[0]

            self.assertEqual(output.shape    ,    output_a.shape    ,    """Shape doesn't match"""						)
            assert np.abs(output.cpu() - output_a.cpu()						).max() < 1E-2


# (todo: sayakpaul) implement SLOW tests.






