




'''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()

UpperCAmelCase_       				=      logging.get_logger(__name__)

def    _UpperCamelCase      (				SCREAMING_SNAKE_CASE__			:      Optional[Any]       ,						SCREAMING_SNAKE_CASE__			:      bool = True       ):
 '''simple docstring'''





 UpperCAmelCase__										=    {}
 for name in model.attn_processors.keys():
  UpperCAmelCase__										=    name.endswith("""attn2.processor"""       ) and not (
      name.startswith("""transformer_in"""       ) or """temp_attentions""" in name.split("""."""       )
  )
  UpperCAmelCase__										=    model.config.cross_attention_dim if has_cross_attention else None
  if name.startswith("""mid_block"""       ):
   UpperCAmelCase__										=    model.config.block_out_channels[-1]
  elif name.startswith("""up_blocks"""       ):
   UpperCAmelCase__										=    int(name[len("""up_blocks."""       )]       )
   UpperCAmelCase__										=    list(reversed(model.config.block_out_channels       )       )[block_id]
  elif name.startswith("""down_blocks"""       ):
   UpperCAmelCase__										=    int(name[len("""down_blocks."""       )]       )
   UpperCAmelCase__										=    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
   UpperCAmelCase__										=    8 * model.config.attention_head_dim

  UpperCAmelCase__										=    LoRAAttnProcessor(hidden_size=SCREAMING_SNAKE_CASE__       ,						cross_attention_dim=SCREAMING_SNAKE_CASE__       )
  UpperCAmelCase__										=    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 						lowerCAmelCase_   ( lowerCamelCase_			,  lowerCamelCase_			,  unittest.TestCase   ):







  '''simple docstring'''







  lowerCAmelCase_				:					Union[str, Any]				     =						UNetaDConditionModel
  lowerCAmelCase_				:					Union[str, Any]				     =						"""sample"""


  @property
  def 							SCREAMING_SNAKE_CASE__			(						self    :	List[Any] ):






   """simple docstring"""






   UpperCAmelCase__										=    4
   UpperCAmelCase__										=    4
   UpperCAmelCase__										=    4
   UpperCAmelCase__										=    (32, 32)

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

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


  @property
  def 							SCREAMING_SNAKE_CASE__			(						self    :	Union[str, Any] ):






   """simple docstring"""






   return (4, 4, 32, 32)


  @property
  def 							SCREAMING_SNAKE_CASE__			(						self    :	str ):






   """simple docstring"""






   return (4, 4, 32, 32)


  def 							SCREAMING_SNAKE_CASE__			(						self    :	Optional[int] ):






   """simple docstring"""






   UpperCAmelCase__										=    {
       """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,
   }
   UpperCAmelCase__										=    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 							SCREAMING_SNAKE_CASE__			(						self    :	List[Any] ):






   """simple docstring"""






   UpperCAmelCase__							,			UpperCAmelCase__										=    self.prepare_init_args_and_inputs_for_common()
   UpperCAmelCase__										=    self.model_class(**_UpperCAmelCase )

   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 							SCREAMING_SNAKE_CASE__			(						self    :	Optional[Any] ):






   """simple docstring"""






   UpperCAmelCase__							,			UpperCAmelCase__										=    self.prepare_init_args_and_inputs_for_common()

   UpperCAmelCase__										=    32

   UpperCAmelCase__										=    self.model_class(**_UpperCAmelCase )
   model.to(_UpperCAmelCase )
   model.eval()

   with torch.no_grad():
    UpperCAmelCase__										=    model(**_UpperCAmelCase )

    if isinstance(_UpperCAmelCase						,		_UpperCAmelCase ):
     UpperCAmelCase__										=    output.sample

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


  def 							SCREAMING_SNAKE_CASE__			(						self    :	List[str] ):






   """simple docstring"""






   UpperCAmelCase__							,			UpperCAmelCase__										=    self.prepare_init_args_and_inputs_for_common()
   UpperCAmelCase__										=    self.model_class(**_UpperCAmelCase )
   model.to(_UpperCAmelCase )
   model.eval()

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

    UpperCAmelCase__										=    model(**_UpperCAmelCase )
    if isinstance(_UpperCAmelCase						,		_UpperCAmelCase ):
     UpperCAmelCase__										=    first.sample

    UpperCAmelCase__										=    model(**_UpperCAmelCase )
    if isinstance(_UpperCAmelCase						,		_UpperCAmelCase ):
     UpperCAmelCase__										=    second.sample

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


  def 							SCREAMING_SNAKE_CASE__			(						self    :	str ):






   """simple docstring"""






   UpperCAmelCase__							,			UpperCAmelCase__										=    self.prepare_init_args_and_inputs_for_common()

   UpperCAmelCase__										=    8

   UpperCAmelCase__										=    self.model_class(**_UpperCAmelCase )
   model.to(_UpperCAmelCase )
   model.eval()

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

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

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


  def 							SCREAMING_SNAKE_CASE__			(						self    :	int ):






   """simple docstring"""






   UpperCAmelCase__							,			UpperCAmelCase__										=    self.prepare_init_args_and_inputs_for_common()

   UpperCAmelCase__										=    8

   UpperCAmelCase__										=    self.model_class(**_UpperCAmelCase )
   model.to(_UpperCAmelCase )

   with torch.no_grad():
    UpperCAmelCase__										=    model(**_UpperCAmelCase ).sample

   UpperCAmelCase__										=    create_lora_layers(_UpperCAmelCase )

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

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

   with torch.no_grad():
    UpperCAmelCase__										=    model(**_UpperCAmelCase						,		cross_attention_kwargs={"""scale""": 0.0} ).sample
    UpperCAmelCase__										=    model(**_UpperCAmelCase						,		cross_attention_kwargs={"""scale""": 0.5} ).sample
    UpperCAmelCase__										=    model(**_UpperCAmelCase						,		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 							SCREAMING_SNAKE_CASE__			(						self    :	List[Any] ):






   """simple docstring"""






   UpperCAmelCase__							,			UpperCAmelCase__										=    self.prepare_init_args_and_inputs_for_common()

   UpperCAmelCase__										=    8

   torch.manual_seed(0 )
   UpperCAmelCase__										=    self.model_class(**_UpperCAmelCase )
   model.to(_UpperCAmelCase )

   with torch.no_grad():
    UpperCAmelCase__										=    model(**_UpperCAmelCase ).sample

   UpperCAmelCase__										=    create_lora_layers(_UpperCAmelCase )
   model.set_attn_processor(_UpperCAmelCase )

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

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

   with torch.no_grad():
    UpperCAmelCase__										=    new_model(**_UpperCAmelCase						,		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 							SCREAMING_SNAKE_CASE__			(						self    :	Tuple ):






   """simple docstring"""






   UpperCAmelCase__							,			UpperCAmelCase__										=    self.prepare_init_args_and_inputs_for_common()

   UpperCAmelCase__										=    8

   torch.manual_seed(0 )
   UpperCAmelCase__										=    self.model_class(**_UpperCAmelCase )
   model.to(_UpperCAmelCase )

   with torch.no_grad():
    UpperCAmelCase__										=    model(**_UpperCAmelCase ).sample

   UpperCAmelCase__										=    create_lora_layers(_UpperCAmelCase )
   model.set_attn_processor(_UpperCAmelCase )

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

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

   with torch.no_grad():
    UpperCAmelCase__										=    new_model(**_UpperCAmelCase						,		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 							SCREAMING_SNAKE_CASE__			(						self    :	List[str] ):






   """simple docstring"""






   UpperCAmelCase__							,			UpperCAmelCase__										=    self.prepare_init_args_and_inputs_for_common()

   UpperCAmelCase__										=    8

   torch.manual_seed(0 )
   UpperCAmelCase__										=    self.model_class(**_UpperCAmelCase )
   model.to(_UpperCAmelCase )

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


  def 							SCREAMING_SNAKE_CASE__			(						self    :	List[Any] ):






   """simple docstring"""






   UpperCAmelCase__							,			UpperCAmelCase__										=    self.prepare_init_args_and_inputs_for_common()

   UpperCAmelCase__										=    8

   torch.manual_seed(0 )
   UpperCAmelCase__										=    self.model_class(**_UpperCAmelCase )
   model.to(_UpperCAmelCase )

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


  def 							SCREAMING_SNAKE_CASE__			(						self    :	Optional[int] ):






   """simple docstring"""






   UpperCAmelCase__							,			UpperCAmelCase__										=    self.prepare_init_args_and_inputs_for_common()

   UpperCAmelCase__										=    8

   torch.manual_seed(0 )
   UpperCAmelCase__										=    self.model_class(**_UpperCAmelCase )
   model.to(_UpperCAmelCase )

   with torch.no_grad():
    UpperCAmelCase__										=    model(**_UpperCAmelCase ).sample

   UpperCAmelCase__										=    create_lora_layers(_UpperCAmelCase )
   model.set_attn_processor(_UpperCAmelCase )

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

   model.set_attn_processor(AttnProcessor() )

   with torch.no_grad():
    UpperCAmelCase__										=    model(**_UpperCAmelCase ).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 							SCREAMING_SNAKE_CASE__			(						self    :	Optional[Any] ):






   """simple docstring"""






   UpperCAmelCase__							,			UpperCAmelCase__										=    self.prepare_init_args_and_inputs_for_common()

   UpperCAmelCase__										=    4

   torch.manual_seed(0 )
   UpperCAmelCase__										=    self.model_class(**_UpperCAmelCase )
   model.to(_UpperCAmelCase )
   UpperCAmelCase__										=    create_lora_layers(_UpperCAmelCase )
   model.set_attn_processor(_UpperCAmelCase )

   # default
   with torch.no_grad():
    UpperCAmelCase__										=    model(**_UpperCAmelCase ).sample

    model.enable_xformers_memory_efficient_attention()
    UpperCAmelCase__										=    model(**_UpperCAmelCase ).sample

    model.disable_xformers_memory_efficient_attention()
    UpperCAmelCase__										=    model(**_UpperCAmelCase ).sample

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



  def 							SCREAMING_SNAKE_CASE__			(						self    :	Optional[int] ):






   """simple docstring"""






   UpperCAmelCase__							,			UpperCAmelCase__										=    self.prepare_init_args_and_inputs_for_common()
   UpperCAmelCase__										=    32

   UpperCAmelCase__										=    self.model_class(**_UpperCAmelCase )
   model.to(_UpperCAmelCase )
   model.eval()

   with torch.no_grad():
    UpperCAmelCase__										=    model(**_UpperCAmelCase )[0]

   model.enable_forward_chunking()
   with torch.no_grad():
    UpperCAmelCase__										=    model(**_UpperCAmelCase )[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.






'''simple docstring'''




import random
import unittest

import torch

from diffusers import IFImgaImgPipeline
from diffusers.utils import floats_tensor
from diffusers.utils.import_utils import is_xformers_available
from diffusers.utils.testing_utils import skip_mps, torch_device

from ..pipeline_params import (
    TEXT_GUIDED_IMAGE_VARIATION_BATCH_PARAMS,
    TEXT_GUIDED_IMAGE_VARIATION_PARAMS,
)
from ..test_pipelines_common import PipelineTesterMixin
from . import IFPipelineTesterMixin

@skip_mps
class 						lowerCAmelCase_   ( lowerCamelCase_			,  lowerCamelCase_			,  unittest.TestCase   ):







  '''simple docstring'''







  lowerCAmelCase_				:					Any				     =						IFImgaImgPipeline
  lowerCAmelCase_				:					Dict				     =						TEXT_GUIDED_IMAGE_VARIATION_PARAMS - {"""width""", """height"""}
  lowerCAmelCase_				:					Union[str, Any]				     =						TEXT_GUIDED_IMAGE_VARIATION_BATCH_PARAMS
  lowerCAmelCase_				:					Dict				     =						PipelineTesterMixin.required_optional_params - {"""latents"""}


  def 							SCREAMING_SNAKE_CASE__			(						self    :	Optional[int] ):






   """simple docstring"""






   return self._get_dummy_components()


  def 							SCREAMING_SNAKE_CASE__			(						self    :	Tuple						,		_UpperCAmelCase    :	Any						,		_UpperCAmelCase    :	Any=0 ):






   """simple docstring"""






   if str(_UpperCAmelCase ).startswith("""mps""" ):
    UpperCAmelCase__										=    torch.manual_seed(_UpperCAmelCase )
   else:
    UpperCAmelCase__										=    torch.Generator(device=_UpperCAmelCase ).manual_seed(_UpperCAmelCase )

   UpperCAmelCase__										=    floats_tensor((1, 3, 32, 32)						,		rng=random.Random(_UpperCAmelCase ) ).to(_UpperCAmelCase )

   UpperCAmelCase__										=    {
       """prompt""": """A painting of a squirrel eating a burger""",
       """image""": image,
       """generator""": generator,
       """num_inference_steps""": 2,
       """output_type""": """numpy""",
   }

   return inputs


  def 							SCREAMING_SNAKE_CASE__			(						self    :	List[Any] ):






   """simple docstring"""






   self._test_save_load_optional_components()


  @unittest.skipIf(
      torch_device != """cuda""" or not is_xformers_available()						,		reason="""XFormers attention is only available with CUDA and `xformers` installed"""						,		)
  def 							SCREAMING_SNAKE_CASE__			(						self    :	int ):






   """simple docstring"""






   self._test_xformers_attention_forwardGenerator_pass(expected_max_diff=1E-3 )


  @unittest.skipIf(torch_device != """cuda"""						,		reason="""float16 requires CUDA""" )
  def 							SCREAMING_SNAKE_CASE__			(						self    :	int ):






   """simple docstring"""






   super().test_save_load_floataa(expected_max_diff=1E-1 )


  @unittest.skipIf(torch_device != """cuda"""						,		reason="""float16 requires CUDA""" )
  def 							SCREAMING_SNAKE_CASE__			(						self    :	Optional[int] ):






   """simple docstring"""






   super().test_floataa_inference(expected_max_diff=1E-1 )


  def 							SCREAMING_SNAKE_CASE__			(						self    :	List[str] ):






   """simple docstring"""






   self._test_attention_slicing_forward_pass(expected_max_diff=1E-2 )


  def 							SCREAMING_SNAKE_CASE__			(						self    :	int ):






   """simple docstring"""






   self._test_save_load_local()



  def 							SCREAMING_SNAKE_CASE__			(						self    :	Any ):






   """simple docstring"""






   self._test_inference_batch_single_identical(
       expected_max_diff=1E-2						,		)
