






"""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__       :			Union[str, Any]=						logging.get_logger(__name__)







def 	lowerCAmelCase_(			SCREAMING_SNAKE_CASE			,						SCREAMING_SNAKE_CASE = True       )       ->						Tuple:
    """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     (	_a				,						_a				,						unittest.TestCase	):
   a   :					Union[str, Any]			    =UNetaDConditionModel
   a   :					Dict			    ="""sample"""

   @property
   def     SCREAMING_SNAKE_CASE__			(       self   )		->							Tuple:
       UpperCamelCase__								=     4
       UpperCamelCase__								=     4
       UpperCamelCase__								=     4
       UpperCamelCase__								=     (32, 32)

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

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

   @property
   def     SCREAMING_SNAKE_CASE__			(       self   )		->							int:
       return (4, 4, 32, 32)

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

   def     SCREAMING_SNAKE_CASE__			(       self   )		->							Any:
       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   )		->							Tuple:
       UpperCamelCase__							,    UpperCamelCase__								=     self.prepare_init_args_and_inputs_for_common()
       UpperCamelCase__								=     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     SCREAMING_SNAKE_CASE__			(       self   )		->							Union[str, Any]:
       UpperCamelCase__							,    UpperCamelCase__								=     self.prepare_init_args_and_inputs_for_common()

       UpperCamelCase__								=     32

       UpperCamelCase__								=     self.model_class(**snake_case_   )
       model.to(snake_case_   )
       model.eval()

       with torch.no_grad():
           UpperCamelCase__								=     model(**snake_case_   )

           if isinstance(snake_case_			,					snake_case_   ):
               UpperCamelCase__								=     output.sample

       self.assertIsNotNone(snake_case_   )
       UpperCamelCase__								=     inputs_dict['sample'].shape
       self.assertEqual(output.shape			,					snake_case_			,					'Input and output shapes do not match'   )

   def     SCREAMING_SNAKE_CASE__			(       self   )		->							Any:
       UpperCamelCase__							,    UpperCamelCase__								=     self.prepare_init_args_and_inputs_for_common()
       UpperCamelCase__								=     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   )

           UpperCamelCase__								=     model(**snake_case_   )
           if isinstance(snake_case_			,					snake_case_   ):
               UpperCamelCase__								=     first.sample

           UpperCamelCase__								=     model(**snake_case_   )
           if isinstance(snake_case_			,					snake_case_   ):
               UpperCamelCase__								=     second.sample

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

   def     SCREAMING_SNAKE_CASE__			(       self   )		->							List[str]:
       UpperCamelCase__							,    UpperCamelCase__								=     self.prepare_init_args_and_inputs_for_common()

       UpperCamelCase__								=     8

       UpperCamelCase__								=     self.model_class(**snake_case_   )
       model.to(snake_case_   )
       model.eval()

       model.set_attention_slice('auto'   )
       with torch.no_grad():
           UpperCamelCase__								=     model(**snake_case_   )
       assert output is not None

       model.set_attention_slice('max'   )
       with torch.no_grad():
           UpperCamelCase__								=     model(**snake_case_   )
       assert output is not None

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

   def     SCREAMING_SNAKE_CASE__			(       self   )		->							int:
       UpperCamelCase__							,    UpperCamelCase__								=     self.prepare_init_args_and_inputs_for_common()

       UpperCamelCase__								=     8

       UpperCamelCase__								=     self.model_class(**snake_case_   )
       model.to(snake_case_   )

       with torch.no_grad():
           UpperCamelCase__								=     model(**snake_case_   ).sample

       UpperCamelCase__								=     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():
           UpperCamelCase__								=     model(**snake_case_			,					cross_attention_kwargs={'scale': 0.0}   ).sample
           UpperCamelCase__								=     model(**snake_case_			,					cross_attention_kwargs={'scale': 0.5}   ).sample
           UpperCamelCase__								=     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     SCREAMING_SNAKE_CASE__			(       self   )		->							Optional[int]:
       UpperCamelCase__							,    UpperCamelCase__								=     self.prepare_init_args_and_inputs_for_common()

       UpperCamelCase__								=     8

       torch.manual_seed(0   )
       UpperCamelCase__								=     self.model_class(**snake_case_   )
       model.to(snake_case_   )

       with torch.no_grad():
           UpperCamelCase__								=     model(**snake_case_   ).sample

       UpperCamelCase__								=     create_lora_layers(snake_case_   )
       model.set_attn_processor(snake_case_   )

       with torch.no_grad():
           UpperCamelCase__								=     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   )
           UpperCamelCase__								=     self.model_class(**snake_case_   )
           new_model.to(snake_case_   )
           new_model.load_attn_procs(snake_case_   )

       with torch.no_grad():
           UpperCamelCase__								=     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     SCREAMING_SNAKE_CASE__			(       self   )		->							Union[str, Any]:
       UpperCamelCase__							,    UpperCamelCase__								=     self.prepare_init_args_and_inputs_for_common()

       UpperCamelCase__								=     8

       torch.manual_seed(0   )
       UpperCamelCase__								=     self.model_class(**snake_case_   )
       model.to(snake_case_   )

       with torch.no_grad():
           UpperCamelCase__								=     model(**snake_case_   ).sample

       UpperCamelCase__								=     create_lora_layers(snake_case_   )
       model.set_attn_processor(snake_case_   )

       with torch.no_grad():
           UpperCamelCase__								=     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   )
           UpperCamelCase__								=     self.model_class(**snake_case_   )
           new_model.to(snake_case_   )
           new_model.load_attn_procs(snake_case_   )

       with torch.no_grad():
           UpperCamelCase__								=     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     SCREAMING_SNAKE_CASE__			(       self   )		->							List[Any]:
       # enable deterministic behavior for gradient checkpointing
       UpperCamelCase__							,    UpperCamelCase__								=     self.prepare_init_args_and_inputs_for_common()

       UpperCamelCase__								=     8

       torch.manual_seed(0   )
       UpperCamelCase__								=     self.model_class(**snake_case_   )
       model.to(snake_case_   )

       UpperCamelCase__								=     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   )
           UpperCamelCase__								=     self.model_class(**snake_case_   )
           new_model.to(snake_case_   )
           new_model.load_attn_procs(snake_case_			,					weight_name='pytorch_lora_weights.bin'   )

   def     SCREAMING_SNAKE_CASE__			(       self   )		->							Optional[Any]:
       UpperCamelCase__							,    UpperCamelCase__								=     self.prepare_init_args_and_inputs_for_common()

       UpperCamelCase__								=     8

       torch.manual_seed(0   )
       UpperCamelCase__								=     self.model_class(**snake_case_   )
       model.to(snake_case_   )

       UpperCamelCase__								=     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   )
           UpperCamelCase__								=     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     SCREAMING_SNAKE_CASE__			(       self   )		->							Any:
       UpperCamelCase__							,    UpperCamelCase__								=     self.prepare_init_args_and_inputs_for_common()

       UpperCamelCase__								=     8

       torch.manual_seed(0   )
       UpperCamelCase__								=     self.model_class(**snake_case_   )
       model.to(snake_case_   )

       with torch.no_grad():
           UpperCamelCase__								=     model(**snake_case_   ).sample

       UpperCamelCase__								=     create_lora_layers(snake_case_   )
       model.set_attn_processor(snake_case_   )

       with torch.no_grad():
           UpperCamelCase__								=     model(**snake_case_			,					cross_attention_kwargs={'scale': 0.0}   ).sample

       model.set_attn_processor(AttnProcessor()   )

       with torch.no_grad():
           UpperCamelCase__								=     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     SCREAMING_SNAKE_CASE__			(       self   )		->							Optional[Any]:
       # enable deterministic behavior for gradient checkpointing
       UpperCamelCase__							,    UpperCamelCase__								=     self.prepare_init_args_and_inputs_for_common()

       UpperCamelCase__								=     4

       torch.manual_seed(0   )
       UpperCamelCase__								=     self.model_class(**snake_case_   )
       model.to(snake_case_   )
       UpperCamelCase__								=     create_lora_layers(snake_case_   )
       model.set_attn_processor(snake_case_   )

       # default
       with torch.no_grad():
           UpperCamelCase__								=     model(**snake_case_   ).sample

           model.enable_xformers_memory_efficient_attention()
           UpperCamelCase__								=     model(**snake_case_   ).sample

           model.disable_xformers_memory_efficient_attention()
           UpperCamelCase__								=     model(**snake_case_   ).sample

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







   def     SCREAMING_SNAKE_CASE__			(       self   )		->							Union[str, Any]:
       UpperCamelCase__							,    UpperCamelCase__								=     self.prepare_init_args_and_inputs_for_common()
       UpperCamelCase__								=     32

       UpperCamelCase__								=     self.model_class(**snake_case_   )
       model.to(snake_case_   )
       model.eval()

       with torch.no_grad():
           UpperCamelCase__								=     model(**snake_case_   )[0]

       model.enable_forward_chunking()
       with torch.no_grad():
           UpperCamelCase__								=     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.












"""simple docstring"""





def 	lowerCAmelCase_(			SCREAMING_SNAKE_CASE			,						SCREAMING_SNAKE_CASE = 10_00       )       ->						str:
    """simple docstring"""

    if (not isinstance(SCREAMING_SNAKE_CASE			,						SCREAMING_SNAKE_CASE       )) or (digit_position <= 0):
        raise ValueError('Digit position must be a positive integer'       )
    elif (not isinstance(SCREAMING_SNAKE_CASE			,						SCREAMING_SNAKE_CASE       )) or (precision < 0):
        raise ValueError('Precision must be a nonnegative integer'       )

    # compute an approximation of (16 ** (n - 1)) * pi whose fractional part is mostly
    # accurate
    UpperCamelCase__								=     (
        4 * _subsum(SCREAMING_SNAKE_CASE			,						1			,						SCREAMING_SNAKE_CASE       )
        - 2 * _subsum(SCREAMING_SNAKE_CASE			,						4			,						SCREAMING_SNAKE_CASE       )
        - _subsum(SCREAMING_SNAKE_CASE			,						5			,						SCREAMING_SNAKE_CASE       )
        - _subsum(SCREAMING_SNAKE_CASE			,						6			,						SCREAMING_SNAKE_CASE       )
    )

    # return the first hex digit of the fractional part of the result
    return hex(int((sum_result % 1) * 16       )       )[2:]







def 	lowerCAmelCase_(			SCREAMING_SNAKE_CASE			,						SCREAMING_SNAKE_CASE			,						SCREAMING_SNAKE_CASE       )       ->						float:
    """simple docstring"""

    UpperCamelCase__								=     0.0
    for sum_index in range(digit_pos_to_extract + precision       ):
        UpperCamelCase__								=     8 * sum_index + denominator_addend
        if sum_index < digit_pos_to_extract:
            # if the exponential term is an integer and we mod it by the denominator
            # before dividing, only the integer part of the sum will change;
            # the fractional part will not
            UpperCamelCase__								=     pow(
                16			,						digit_pos_to_extract - 1 - sum_index			,						SCREAMING_SNAKE_CASE       )
        else:
            UpperCamelCase__								=     pow(16			,						digit_pos_to_extract - 1 - sum_index       )
        total += exponential_term / denominator
    return total


if __name__ == "__main__":
  import doctest

  doctest.testmod()




