





'''simple docstring'''






import gc
import random
import tempfile
import unittest

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

from diffusers import AutoencoderKL, DDIMScheduler, LMSDiscreteScheduler, PNDMScheduler, UNetaDConditionModel
from diffusers.pipelines.semantic_stable_diffusion import SemanticStableDiffusionPipeline as StableDiffusionPipeline
from diffusers.utils import floats_tensor, nightly, torch_device
from diffusers.utils.testing_utils import enable_full_determinism, require_torch_gpu


enable_full_determinism()




class 				UpperCAmelCase  (    unittest.TestCase     ):




      def        lowercase__					(    self    :				Optional[int]						)       ->   List[str]:
            # clean up the VRAM after each test
            super().tearDown()
            gc.collect()
            torch.cuda.empty_cache()




      @property
      def        lowercase__					(    self    :				int						)       ->   Optional[int]:
            _lowerCAmelCase							  =							1
            _lowerCAmelCase							  =							3
            _lowerCAmelCase							  =							(32, 32)

            _lowerCAmelCase							  =							floats_tensor((batch_size, num_channels) + sizes    ,    rng=random.Random(0						)						).to(__snake_case						)
            return image




      @property
      def        lowercase__					(    self    :				Tuple						)       ->   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""")    ,    cross_attention_dim=32    ,    )
            return model




      @property
      def        lowercase__					(    self    :				Dict						)       ->   Optional[Any]:
            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    ,    )
            return model




      @property
      def        lowercase__					(    self    :				str						)       ->   str:
            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    ,    )
            return CLIPTextModel(__snake_case						)




      @property
      def        lowercase__					(    self    :				List[Any]						)       ->   List[Any]:
            def extract(*__snake_case    :				str    ,    **__snake_case    :				Optional[Any]						):




                  class 				UpperCAmelCase  :




                        def __init__(    self    :				str						)       ->   int:
                              _lowerCAmelCase							  =							torch.ones([0]						)




                        def        lowercase__					(    self    :				Optional[int]    ,    __snake_case    :				Optional[Any]						)       ->   Optional[Any]:
                              self.pixel_values.to(__snake_case						)
                              return self




                  return Out()

            return extract




      def        lowercase__					(    self    :				int						)       ->   Tuple:
            _lowerCAmelCase							  =							"""cpu"""  # ensure determinism for the device-dependent torch.Generator
            _lowerCAmelCase							  =							self.dummy_cond_unet
            _lowerCAmelCase							  =							DDIMScheduler(
                beta_start=0.0_00_85    ,    beta_end=0.0_12    ,    beta_schedule="""scaled_linear"""    ,    clip_sample=__snake_case    ,    set_alpha_to_one=__snake_case    ,    )

            _lowerCAmelCase							  =							self.dummy_vae
            _lowerCAmelCase							  =							self.dummy_text_encoder
            _lowerCAmelCase							  =							CLIPTokenizer.from_pretrained("""hf-internal-testing/tiny-random-clip"""						)

            # make sure here that pndm scheduler skips prk
            _lowerCAmelCase							  =							StableDiffusionPipeline(
                unet=__snake_case    ,    scheduler=__snake_case    ,    vae=__snake_case    ,    text_encoder=__snake_case    ,    tokenizer=__snake_case    ,    safety_checker=__snake_case    ,    feature_extractor=self.dummy_extractor    ,    )
            _lowerCAmelCase							  =							sd_pipe.to(__snake_case						)
            sd_pipe.set_progress_bar_config(disable=__snake_case						)

            _lowerCAmelCase							  =							"""A painting of a squirrel eating a burger"""

            _lowerCAmelCase							  =							torch.Generator(device=__snake_case						).manual_seed(0						)
            _lowerCAmelCase							  =							sd_pipe([prompt]    ,    generator=__snake_case    ,    guidance_scale=6.0    ,    num_inference_steps=2    ,    output_type="""np"""						)
            _lowerCAmelCase							  =							output.images

            _lowerCAmelCase							  =							torch.Generator(device=__snake_case						).manual_seed(0						)
            _lowerCAmelCase							  =							sd_pipe(
                [prompt]    ,    generator=__snake_case    ,    guidance_scale=6.0    ,    num_inference_steps=2    ,    output_type="""np"""    ,    return_dict=__snake_case    ,    )[0]

            _lowerCAmelCase							  =							image[0, -3:, -3:, -1]
            _lowerCAmelCase							  =							image_from_tuple[0, -3:, -3:, -1]

            assert image.shape == (1, 64, 64, 3)
            _lowerCAmelCase							  =							np.array([0.57_53, 0.61_14, 0.50_01, 0.50_34, 0.54_70, 0.47_29, 0.49_71, 0.48_67, 0.48_67]						)

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




      def        lowercase__					(    self    :				str						)       ->   str:
            _lowerCAmelCase							  =							"""cpu"""  # ensure determinism for the device-dependent torch.Generator
            _lowerCAmelCase							  =							self.dummy_cond_unet
            _lowerCAmelCase							  =							PNDMScheduler(skip_prk_steps=__snake_case						)
            _lowerCAmelCase							  =							self.dummy_vae
            _lowerCAmelCase							  =							self.dummy_text_encoder
            _lowerCAmelCase							  =							CLIPTokenizer.from_pretrained("""hf-internal-testing/tiny-random-clip"""						)

            # make sure here that pndm scheduler skips prk
            _lowerCAmelCase							  =							StableDiffusionPipeline(
                unet=__snake_case    ,    scheduler=__snake_case    ,    vae=__snake_case    ,    text_encoder=__snake_case    ,    tokenizer=__snake_case    ,    safety_checker=__snake_case    ,    feature_extractor=self.dummy_extractor    ,    )
            _lowerCAmelCase							  =							sd_pipe.to(__snake_case						)
            sd_pipe.set_progress_bar_config(disable=__snake_case						)

            _lowerCAmelCase							  =							"""A painting of a squirrel eating a burger"""
            _lowerCAmelCase							  =							torch.Generator(device=__snake_case						).manual_seed(0						)
            _lowerCAmelCase							  =							sd_pipe([prompt]    ,    generator=__snake_case    ,    guidance_scale=6.0    ,    num_inference_steps=2    ,    output_type="""np"""						)

            _lowerCAmelCase							  =							output.images

            _lowerCAmelCase							  =							torch.Generator(device=__snake_case						).manual_seed(0						)
            _lowerCAmelCase							  =							sd_pipe(
                [prompt]    ,    generator=__snake_case    ,    guidance_scale=6.0    ,    num_inference_steps=2    ,    output_type="""np"""    ,    return_dict=__snake_case    ,    )[0]

            _lowerCAmelCase							  =							image[0, -3:, -3:, -1]
            _lowerCAmelCase							  =							image_from_tuple[0, -3:, -3:, -1]

            assert image.shape == (1, 64, 64, 3)
            _lowerCAmelCase							  =							np.array([0.51_22, 0.57_12, 0.48_25, 0.50_53, 0.56_46, 0.47_69, 0.51_79, 0.48_94, 0.49_94]						)

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




      def        lowercase__					(    self    :				List[str]						)       ->   int:
            _lowerCAmelCase							  =							StableDiffusionPipeline.from_pretrained(
                """hf-internal-testing/tiny-stable-diffusion-lms-pipe"""    ,    safety_checker=__snake_case						)
            assert isinstance(__snake_case    ,    __snake_case						)
            assert isinstance(pipe.scheduler    ,    __snake_case						)
            assert pipe.safety_checker is None

            _lowerCAmelCase							  =							pipe("""example prompt"""    ,    num_inference_steps=2						).images[0]
            assert image is not None

            # check that there's no error when saving a pipeline with one of the models being None
            with tempfile.TemporaryDirectory() as tmpdirname:
                  pipe.save_pretrained(__snake_case						)
                  _lowerCAmelCase							  =							StableDiffusionPipeline.from_pretrained(__snake_case						)

            # sanity check that the pipeline still works
            assert pipe.safety_checker is None
            _lowerCAmelCase							  =							pipe("""example prompt"""    ,    num_inference_steps=2						).images[0]
            assert image is not None




      @unittest.skipIf(torch_device != """cuda"""    ,    """This test requires a GPU"""						)
      def        lowercase__					(    self    :				Union[str, Any]						)       ->   List[Any]:
            _lowerCAmelCase							  =							self.dummy_cond_unet
            _lowerCAmelCase							  =							PNDMScheduler(skip_prk_steps=__snake_case						)
            _lowerCAmelCase							  =							self.dummy_vae
            _lowerCAmelCase							  =							self.dummy_text_encoder
            _lowerCAmelCase							  =							CLIPTokenizer.from_pretrained("""hf-internal-testing/tiny-random-clip"""						)

            # put models in fp16
            _lowerCAmelCase							  =							unet.half()
            _lowerCAmelCase							  =							vae.half()
            _lowerCAmelCase							  =							bert.half()

            # make sure here that pndm scheduler skips prk
            _lowerCAmelCase							  =							StableDiffusionPipeline(
                unet=__snake_case    ,    scheduler=__snake_case    ,    vae=__snake_case    ,    text_encoder=__snake_case    ,    tokenizer=__snake_case    ,    safety_checker=__snake_case    ,    feature_extractor=self.dummy_extractor    ,    )
            _lowerCAmelCase							  =							sd_pipe.to(__snake_case						)
            sd_pipe.set_progress_bar_config(disable=__snake_case						)

            _lowerCAmelCase							  =							"""A painting of a squirrel eating a burger"""
            _lowerCAmelCase							  =							sd_pipe([prompt]    ,    num_inference_steps=2    ,    output_type="""np"""						).images

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




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




      def        lowercase__					(    self    :				List[str]						)       ->   List[str]:
            # clean up the VRAM after each test
            super().tearDown()
            gc.collect()
            torch.cuda.empty_cache()




      def        lowercase__					(    self    :				Optional[int]						)       ->   Tuple:
            _lowerCAmelCase							  =							"""cuda"""
            _lowerCAmelCase							  =							StableDiffusionPipeline.from_pretrained("""runwayml/stable-diffusion-v1-5"""						)
            _lowerCAmelCase							  =							pipe.to(__snake_case						)
            pipe.set_progress_bar_config(disable=__snake_case						)

            _lowerCAmelCase							  =							"""a photo of a cat"""
            _lowerCAmelCase							  =							{
                """editing_prompt""": ["""sunglasses"""],
                """reverse_editing_direction""": [False],
                """edit_warmup_steps""": 10,
                """edit_guidance_scale""": 6,
                """edit_threshold""": 0.95,
                """edit_momentum_scale""": 0.5,
                """edit_mom_beta""": 0.6,
            }

            _lowerCAmelCase							  =							3
            _lowerCAmelCase							  =							7

            # no sega enabled
            _lowerCAmelCase							  =							torch.Generator(__snake_case						)
            generator.manual_seed(__snake_case						)
            _lowerCAmelCase							  =							pipe(
                [prompt]    ,    generator=__snake_case    ,    guidance_scale=__snake_case    ,    num_inference_steps=50    ,    output_type="""np"""    ,    width=5_12    ,    height=5_12    ,    )

            _lowerCAmelCase							  =							output.images
            _lowerCAmelCase							  =							image[0, -3:, -3:, -1]
            _lowerCAmelCase							  =							[
                0.34_67_31_13,
                0.38_49_27_33,
                0.37_59_73_52,
                0.34_08_63_35,
                0.35_65_07_48,
                0.35_57_92_05,
                0.3_38_47_63,
                0.34_34_02_36,
                0.3_57_32_71,
            ]

            assert image.shape == (1, 5_12, 5_12, 3)

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

            # with sega enabled
            # generator = torch.manual_seed(seed)
            generator.manual_seed(__snake_case						)
            _lowerCAmelCase							  =							pipe(
                [prompt]    ,    generator=__snake_case    ,    guidance_scale=__snake_case    ,    num_inference_steps=50    ,    output_type="""np"""    ,    width=5_12    ,    height=5_12    ,    **__snake_case    ,    )

            _lowerCAmelCase							  =							output.images
            _lowerCAmelCase							  =							image[0, -3:, -3:, -1]
            _lowerCAmelCase							  =							[
                0.41_88_78_26,
                0.37_72_87_66,
                0.30_13_82_72,
                0.41_41_63_35,
                0.41_66_49_85,
                0.36_28_33_92,
                0.36_19_12_46,
                0.43_36_44_65,
                0.43_00_17_32,
            ]

            assert image.shape == (1, 5_12, 5_12, 3)

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




      def        lowercase__					(    self    :				Optional[Any]						)       ->   Union[str, Any]:
            _lowerCAmelCase							  =							"""cuda"""
            _lowerCAmelCase							  =							StableDiffusionPipeline.from_pretrained("""runwayml/stable-diffusion-v1-5"""						)
            _lowerCAmelCase							  =							pipe.to(__snake_case						)
            pipe.set_progress_bar_config(disable=__snake_case						)

            _lowerCAmelCase							  =							"""an image of a crowded boulevard, realistic, 4k"""
            _lowerCAmelCase							  =							{
                """editing_prompt""": """crowd, crowded, people""",
                """reverse_editing_direction""": True,
                """edit_warmup_steps""": 10,
                """edit_guidance_scale""": 8.3,
                """edit_threshold""": 0.9,
                """edit_momentum_scale""": 0.5,
                """edit_mom_beta""": 0.6,
            }

            _lowerCAmelCase							  =							9
            _lowerCAmelCase							  =							7

            # no sega enabled
            _lowerCAmelCase							  =							torch.Generator(__snake_case						)
            generator.manual_seed(__snake_case						)
            _lowerCAmelCase							  =							pipe(
                [prompt]    ,    generator=__snake_case    ,    guidance_scale=__snake_case    ,    num_inference_steps=50    ,    output_type="""np"""    ,    width=5_12    ,    height=5_12    ,    )

            _lowerCAmelCase							  =							output.images
            _lowerCAmelCase							  =							image[0, -3:, -3:, -1]
            _lowerCAmelCase							  =							[
                0.43_49_79_98,
                0.91_81_40_65,
                0.7_54_07_39,
                0.55_58_02_05,
                0.8_46_72_65,
                0.5_38_96_91,
                0.62_57_45_06,
                0.58_89_77_63,
                0.50_92_67_57,
            ]

            assert image.shape == (1, 5_12, 5_12, 3)

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

            # with sega enabled
            # generator = torch.manual_seed(seed)
            generator.manual_seed(__snake_case						)
            _lowerCAmelCase							  =							pipe(
                [prompt]    ,    generator=__snake_case    ,    guidance_scale=__snake_case    ,    num_inference_steps=50    ,    output_type="""np"""    ,    width=5_12    ,    height=5_12    ,    **__snake_case    ,    )

            _lowerCAmelCase							  =							output.images
            _lowerCAmelCase							  =							image[0, -3:, -3:, -1]
            _lowerCAmelCase							  =							[
                0.3_08_97_19,
                0.30_50_01_44,
                0.29_01_60_42,
                0.30_63_09_64,
                0.32_56_87,
                0.29_41_92_25,
                0.2_90_80_91,
                0.28_72_35_98,
                0.27_69_62_94,
            ]

            assert image.shape == (1, 5_12, 5_12, 3)

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




      def        lowercase__					(    self    :				int						)       ->   str:
            _lowerCAmelCase							  =							"""cuda"""
            _lowerCAmelCase							  =							StableDiffusionPipeline.from_pretrained("""runwayml/stable-diffusion-v1-5"""						)
            _lowerCAmelCase							  =							pipe.to(__snake_case						)
            pipe.set_progress_bar_config(disable=__snake_case						)

            _lowerCAmelCase							  =							"""a castle next to a river"""
            _lowerCAmelCase							  =							{
                """editing_prompt""": ["""boat on a river, boat""", """monet, impression, sunrise"""],
                """reverse_editing_direction""": False,
                """edit_warmup_steps""": [15, 18],
                """edit_guidance_scale""": 6,
                """edit_threshold""": [0.9, 0.8],
                """edit_momentum_scale""": 0.5,
                """edit_mom_beta""": 0.6,
            }

            _lowerCAmelCase							  =							48
            _lowerCAmelCase							  =							7

            # no sega enabled
            _lowerCAmelCase							  =							torch.Generator(__snake_case						)
            generator.manual_seed(__snake_case						)
            _lowerCAmelCase							  =							pipe(
                [prompt]    ,    generator=__snake_case    ,    guidance_scale=__snake_case    ,    num_inference_steps=50    ,    output_type="""np"""    ,    width=5_12    ,    height=5_12    ,    )

            _lowerCAmelCase							  =							output.images
            _lowerCAmelCase							  =							image[0, -3:, -3:, -1]
            _lowerCAmelCase							  =							[
                0.75_16_35_55,
                0.76_03_71_45,
                0.6_17_85,
                0.9_18_96_73,
                0.8_62_77_01,
                0.85_18_96_94,
                0.8_51_28_13,
                0.87_01_20_76,
                0.8_31_28_57,
            ]

            assert image.shape == (1, 5_12, 5_12, 3)

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

            # with sega enabled
            # generator = torch.manual_seed(seed)
            generator.manual_seed(__snake_case						)
            _lowerCAmelCase							  =							pipe(
                [prompt]    ,    generator=__snake_case    ,    guidance_scale=__snake_case    ,    num_inference_steps=50    ,    output_type="""np"""    ,    width=5_12    ,    height=5_12    ,    **__snake_case    ,    )

            _lowerCAmelCase							  =							output.images
            _lowerCAmelCase							  =							image[0, -3:, -3:, -1]
            _lowerCAmelCase							  =							[
                0.73_55_33_65,
                0.7_53_72_71,
                0.74_34_19_05,
                0.66_48_03_56,
                0.6_47_29_25,
                0.63_03_94_16,
                0.64_81_29_05,
                0.6_74_97_17,
                0.6_51_71_02,
            ]

            assert image.shape == (1, 5_12, 5_12, 3)

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




      def        lowercase__					(    self    :				List[Any]						)       ->   Dict:
            _lowerCAmelCase							  =							"""cuda"""
            _lowerCAmelCase							  =							StableDiffusionPipeline.from_pretrained("""runwayml/stable-diffusion-v1-5"""    ,    torch_dtype=torch.floataa						)
            _lowerCAmelCase							  =							pipe.to(__snake_case						)
            pipe.set_progress_bar_config(disable=__snake_case						)

            _lowerCAmelCase							  =							"""a photo of a cat"""
            _lowerCAmelCase							  =							{
                """editing_prompt""": ["""sunglasses"""],
                """reverse_editing_direction""": [False],
                """edit_warmup_steps""": 10,
                """edit_guidance_scale""": 6,
                """edit_threshold""": 0.95,
                """edit_momentum_scale""": 0.5,
                """edit_mom_beta""": 0.6,
            }

            _lowerCAmelCase							  =							3
            _lowerCAmelCase							  =							7

            # no sega enabled
            _lowerCAmelCase							  =							torch.Generator(__snake_case						)
            generator.manual_seed(__snake_case						)
            _lowerCAmelCase							  =							pipe(
                [prompt]    ,    generator=__snake_case    ,    guidance_scale=__snake_case    ,    num_inference_steps=50    ,    output_type="""np"""    ,    width=5_12    ,    height=5_12    ,    )

            _lowerCAmelCase							  =							output.images
            _lowerCAmelCase							  =							image[0, -3:, -3:, -1]
            _lowerCAmelCase							  =							[
                0.34_88_76_95,
                0.3_87_69_53,
                0.3_75,
                0.34_42_38_28,
                0.3_58_15_43,
                0.35_71_77_73,
                0.3_38_37_89,
                0.34_57_03_12,
                0.35_93_75,
            ]

            assert image.shape == (1, 5_12, 5_12, 3)

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

            # with sega enabled
            # generator = torch.manual_seed(seed)
            generator.manual_seed(__snake_case						)
            _lowerCAmelCase							  =							pipe(
                [prompt]    ,    generator=__snake_case    ,    guidance_scale=__snake_case    ,    num_inference_steps=50    ,    output_type="""np"""    ,    width=5_12    ,    height=5_12    ,    **__snake_case    ,    )

            _lowerCAmelCase							  =							output.images
            _lowerCAmelCase							  =							image[0, -3:, -3:, -1]
            _lowerCAmelCase							  =							[
                0.42_28_51_56,
                0.36_91_40_62,
                0.29_07_71_48,
                0.42_04_10_16,
                0.41_91_89_45,
                0.35_49_80_47,
                0.3_61_81_64,
                0.4_42_38_28,
                0.43_11_52_34,
            ]

            assert image.shape == (1, 5_12, 5_12, 3)

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













'''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.






