






"""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				     =  logging.get_logger(__name__)


def __A   (				__lowerCAmelCase			,		__lowerCAmelCase = True     )->  str:





						"""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=__lowerCAmelCase			,		cross_attention_dim=__lowerCAmelCase     )
												_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			(				snake_case__					,	snake_case__					,	unittest.TestCase):





	"""simple docstring"""






	UpperCamelCase__								=				UNetaDConditionModel
	UpperCamelCase__								=				"sample"

	@property
	def 	UpperCamelCase   (							self   ):
							"""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 	UpperCamelCase   (							self   ):
							"""simple docstring"""







							return (4, 4, 32, 32)

	@property
	def 	UpperCamelCase   (							self   ):
							"""simple docstring"""







							return (4, 4, 32, 32)

	def 	UpperCamelCase   (							self   ):
							"""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 	UpperCamelCase   (							self   ):
							"""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 	UpperCamelCase   (							self   ):
							"""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 	UpperCamelCase   (							self   ):
							"""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 	UpperCamelCase   (							self   ):
							"""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 	UpperCamelCase   (							self   ):
							"""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 	UpperCamelCase   (							self   ):
							"""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 	UpperCamelCase   (							self   ):
							"""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 	UpperCamelCase   (							self   ):
							"""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 	UpperCamelCase   (							self   ):
							"""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 	UpperCamelCase   (							self   ):
							"""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 	UpperCamelCase   (							self   ):
							"""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 	UpperCamelCase   (							self   ):
							"""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.








import unittest

import torch

from diffusers import DDIMScheduler, TextToVideoZeroPipeline
from diffusers.utils import load_pt, require_torch_gpu, slow

from ..test_pipelines_common import assert_mean_pixel_difference





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





	"""simple docstring"""






	def 	UpperCamelCase   (							self   ):
							"""simple docstring"""







							_UpperCAmelCase     =      'runwayml/stable-diffusion-v1-5'
							_UpperCAmelCase     =      TextToVideoZeroPipeline.from_pretrained(UpperCAmelCase	,   torch_dtype=torch.floataa   ).to('cuda'   )
							_UpperCAmelCase     =      DDIMScheduler.from_config(pipe.scheduler.config   )
							_UpperCAmelCase     =      torch.Generator(device='cuda'   ).manual_seed(0   )

							_UpperCAmelCase     =      'A bear is playing a guitar on Times Square'
							_UpperCAmelCase     =      pipe(prompt=UpperCAmelCase	,   generator=UpperCAmelCase   ).images

							_UpperCAmelCase     =      load_pt(
							    'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/text-to-video/A bear is playing a guitar on Times Square.pt'   )

							assert_mean_pixel_difference(UpperCAmelCase	,   UpperCAmelCase   )
