






"""simple docstring"""



import unittest

import torch

from diffusers import UNetaDModel
from diffusers.utils import floats_tensor, slow, torch_device

from .test_modeling_common import ModelTesterMixin, UNetTesterMixin

class 			a	(  a_,							a_,							unittest.TestCase      ):
    UpperCAmelCase_					:	Tuple      	=UNetaDModel
    UpperCAmelCase_					:	Optional[int]      	="sample"







    @property
    def      UpperCamelCase_ (     self			):
      lowercase				       =			4
      lowercase				       =			1_4
      lowercase				       =			1_6

      lowercase				       =			floats_tensor((batch_size, num_features, seq_len)			).to(_lowerCamelCase			)
      lowercase				       =			torch.tensor([1_0] * batch_size			).to(_lowerCamelCase			)

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







    @property
    def      UpperCamelCase_ (     self			):
      return (4, 1_4, 1_6)







    @property
    def      UpperCamelCase_ (     self			):
      return (4, 1_4, 1_6)







    def      UpperCamelCase_ (     self			):
      pass







    def      UpperCamelCase_ (     self			):
      pass







    @unittest.skipIf(torch_device == 'mps'				,						'mish op not supported in MPS'			)
    def      UpperCamelCase_ (     self			):
      super().test_determinism()







    @unittest.skipIf(torch_device == 'mps'				,						'mish op not supported in MPS'			)
    def      UpperCamelCase_ (     self			):
      super().test_outputs_equivalence()







    @unittest.skipIf(torch_device == 'mps'				,						'mish op not supported in MPS'			)
    def      UpperCamelCase_ (     self			):
      super().test_from_save_pretrained()







    @unittest.skipIf(torch_device == 'mps'				,						'mish op not supported in MPS'			)
    def      UpperCamelCase_ (     self			):
      super().test_from_save_pretrained_variant()







    @unittest.skipIf(torch_device == 'mps'				,						'mish op not supported in MPS'			)
    def      UpperCamelCase_ (     self			):
      super().test_model_from_pretrained()







    @unittest.skipIf(torch_device == 'mps'				,						'mish op not supported in MPS'			)
    def      UpperCamelCase_ (     self			):
      super().test_output()







    def      UpperCamelCase_ (     self			):
      lowercase				       =			{
          'block_out_channels': (3_2, 6_4, 1_2_8, 2_5_6),
          'in_channels': 1_4,
          'out_channels': 1_4,
          'time_embedding_type': 'positional',
          'use_timestep_embedding': True,
          'flip_sin_to_cos': False,
          'freq_shift': 1.0,
          'out_block_type': 'OutConv1DBlock',
          'mid_block_type': 'MidResTemporalBlock1D',
          'down_block_types': ('DownResnetBlock1D', 'DownResnetBlock1D', 'DownResnetBlock1D', 'DownResnetBlock1D'),
          'up_block_types': ('UpResnetBlock1D', 'UpResnetBlock1D', 'UpResnetBlock1D'),
          'act_fn': 'mish',
      }
      lowercase				       =			self.dummy_input
      return init_dict, inputs_dict







    @unittest.skipIf(torch_device == 'mps'				,						'mish op not supported in MPS'			)
    def      UpperCamelCase_ (     self			):
      lowercase       ,	lowercase				       =			UNetaDModel.from_pretrained(
          'bglick13/hopper-medium-v2-value-function-hor32'				,						output_loading_info=_lowerCamelCase				,						subfolder='unet'			)
      self.assertIsNotNone(_lowerCamelCase			)
      self.assertEqual(len(loading_info['missing_keys']			)				,						0			)

      model.to(_lowerCamelCase			)
      lowercase				       =			model(**self.dummy_input			)

      assert image is not None, "Make sure output is not None"







    @unittest.skipIf(torch_device == 'mps'				,						'mish op not supported in MPS'			)
    def      UpperCamelCase_ (     self			):
      lowercase				       =			UNetaDModel.from_pretrained('bglick13/hopper-medium-v2-value-function-hor32'				,						subfolder='unet'			)
      torch.manual_seed(0			)
      if torch.cuda.is_available():
        torch.cuda.manual_seed_all(0			)

      lowercase				       =			model.config.in_channels
      lowercase				       =			1_6
      lowercase				       =			torch.randn((1, seq_len, num_features)			).permute(
          0				,						2				,						1			)  # match original, we can update values and remove
      lowercase				       =			torch.full((num_features,)				,						0			)

      with torch.no_grad():
        lowercase				       =			model(_lowerCamelCase				,						_lowerCamelCase			).sample.permute(0				,						2				,						1			)

      lowercase				       =			output[0, -3:, -3:].flatten()
      # fmt: off
      lowercase				       =			torch.tensor([-2.1_3_7_1_7_2, 1.1_4_2_6_0_1_6, 0.3_6_8_8_6_8_7, -0.7_6_6_9_2_2, 0.7_3_0_3_1_4_6, 0.1_1_0_3_8_8_6_4, -0.4_7_6_0_6_3_3, 0.1_3_2_7_0_1_7_2, 0.0_2_5_9_1_3_4_8]			)
      # fmt: on
      self.assertTrue(torch.allclose(_lowerCamelCase				,						_lowerCamelCase				,						rtol=1e-3			)			)







    def      UpperCamelCase_ (     self			):
      # Not implemented yet for this UNet
      pass







    @slow
    def      UpperCamelCase_ (     self			):
      lowercase				       =			'harmonai/maestro-150k'
      lowercase				       =			UNetaDModel.from_pretrained(_lowerCamelCase				,						subfolder='unet'			)
      model.to(_lowerCamelCase			)

      lowercase				       =			6_5_5_3_6
      lowercase				       =			torch.sin(torch.arange(_lowerCamelCase			)[None, None, :].repeat(1				,						2				,						1			)			).to(_lowerCamelCase			)
      lowercase				       =			torch.tensor([1]			).to(_lowerCamelCase			)

      with torch.no_grad():
        lowercase				       =			model(_lowerCamelCase				,						_lowerCamelCase			).sample

      lowercase				       =			output.abs().sum()
      lowercase				       =			output.abs().max()

      assert (output_sum - 2_2_4.0_8_9_6).abs() < 0.5
      assert (output_max - 0.0_6_0_7).abs() < 4e-4


class 			a	(  a_,							a_,							unittest.TestCase      ):
    UpperCAmelCase_					:	Dict      	=UNetaDModel
    UpperCAmelCase_					:	str      	="sample"







    @property
    def      UpperCamelCase_ (     self			):
      lowercase				       =			4
      lowercase				       =			1_4
      lowercase				       =			1_6

      lowercase				       =			floats_tensor((batch_size, num_features, seq_len)			).to(_lowerCamelCase			)
      lowercase				       =			torch.tensor([1_0] * batch_size			).to(_lowerCamelCase			)

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







    @property
    def      UpperCamelCase_ (     self			):
      return (4, 1_4, 1_6)







    @property
    def      UpperCamelCase_ (     self			):
      return (4, 1_4, 1)







    @unittest.skipIf(torch_device == 'mps'				,						'mish op not supported in MPS'			)
    def      UpperCamelCase_ (     self			):
      super().test_determinism()







    @unittest.skipIf(torch_device == 'mps'				,						'mish op not supported in MPS'			)
    def      UpperCamelCase_ (     self			):
      super().test_outputs_equivalence()







    @unittest.skipIf(torch_device == 'mps'				,						'mish op not supported in MPS'			)
    def      UpperCamelCase_ (     self			):
      super().test_from_save_pretrained()







    @unittest.skipIf(torch_device == 'mps'				,						'mish op not supported in MPS'			)
    def      UpperCamelCase_ (     self			):
      super().test_from_save_pretrained_variant()







    @unittest.skipIf(torch_device == 'mps'				,						'mish op not supported in MPS'			)
    def      UpperCamelCase_ (     self			):
      super().test_model_from_pretrained()







    @unittest.skipIf(torch_device == 'mps'				,						'mish op not supported in MPS'			)
    def      UpperCamelCase_ (     self			):
      # UNetRL is a value-function is different output shape
      lowercase       ,	lowercase				       =			self.prepare_init_args_and_inputs_for_common()
      lowercase				       =			self.model_class(**_lowerCamelCase			)
      model.to(_lowerCamelCase			)
      model.eval()

      with torch.no_grad():
        lowercase				       =			model(**_lowerCamelCase			)

        if isinstance(_lowerCamelCase				,						_lowerCamelCase			):
          lowercase				       =			output.sample

      self.assertIsNotNone(_lowerCamelCase			)
      lowercase				       =			torch.Size((inputs_dict['sample'].shape[0], 1)			)
      self.assertEqual(output.shape				,						_lowerCamelCase				,						'Input and output shapes do not match'			)







    def      UpperCamelCase_ (     self			):
      pass







    def      UpperCamelCase_ (     self			):
      pass







    def      UpperCamelCase_ (     self			):
      lowercase				       =			{
          'in_channels': 1_4,
          'out_channels': 1_4,
          'down_block_types': ['DownResnetBlock1D', 'DownResnetBlock1D', 'DownResnetBlock1D', 'DownResnetBlock1D'],
          'up_block_types': [],
          'out_block_type': 'ValueFunction',
          'mid_block_type': 'ValueFunctionMidBlock1D',
          'block_out_channels': [3_2, 6_4, 1_2_8, 2_5_6],
          'layers_per_block': 1,
          'downsample_each_block': True,
          'use_timestep_embedding': True,
          'freq_shift': 1.0,
          'flip_sin_to_cos': False,
          'time_embedding_type': 'positional',
          'act_fn': 'mish',
      }
      lowercase				       =			self.dummy_input
      return init_dict, inputs_dict







    @unittest.skipIf(torch_device == 'mps'				,						'mish op not supported in MPS'			)
    def      UpperCamelCase_ (     self			):
      lowercase       ,	lowercase				       =			UNetaDModel.from_pretrained(
          'bglick13/hopper-medium-v2-value-function-hor32'				,						output_loading_info=_lowerCamelCase				,						subfolder='value_function'			)
      self.assertIsNotNone(_lowerCamelCase			)
      self.assertEqual(len(vf_loading_info['missing_keys']			)				,						0			)

      value_function.to(_lowerCamelCase			)
      lowercase				       =			value_function(**self.dummy_input			)

      assert image is not None, "Make sure output is not None"







    @unittest.skipIf(torch_device == 'mps'				,						'mish op not supported in MPS'			)
    def      UpperCamelCase_ (     self			):
      lowercase       ,	lowercase				       =			UNetaDModel.from_pretrained(
          'bglick13/hopper-medium-v2-value-function-hor32'				,						output_loading_info=_lowerCamelCase				,						subfolder='value_function'			)
      torch.manual_seed(0			)
      if torch.cuda.is_available():
        torch.cuda.manual_seed_all(0			)

      lowercase				       =			value_function.config.in_channels
      lowercase				       =			1_4
      lowercase				       =			torch.randn((1, seq_len, num_features)			).permute(
          0				,						2				,						1			)  # match original, we can update values and remove
      lowercase				       =			torch.full((num_features,)				,						0			)

      with torch.no_grad():
        lowercase				       =			value_function(_lowerCamelCase				,						_lowerCamelCase			).sample

      # fmt: off
      lowercase				       =			torch.tensor([1_6_5.2_5] * seq_len			)
      # fmt: on
      self.assertTrue(torch.allclose(_lowerCamelCase				,						_lowerCamelCase				,						rtol=1e-3			)			)







    def      UpperCamelCase_ (     self			):
      # Not implemented yet for this UNet
      pass






import tempfile
import unittest

import torch

from diffusers import UNetaDConditionModel
from diffusers.training_utils import EMAModel
from diffusers.utils.testing_utils import enable_full_determinism, skip_mps, torch_device


enable_full_determinism()



class      snake_case__	(   unittest.TestCase	):
				__lowerCAmelCase		=   '''hf-internal-testing/tiny-stable-diffusion-pipe'''
				__lowerCAmelCase		=   1
				__lowerCAmelCase		=   77
				__lowerCAmelCase		=   32
				__lowerCAmelCase		=   4
				__lowerCAmelCase		=   __lowerCAmelCase		=   64
				__lowerCAmelCase		=   torch.manual_seed(0	)
				def 			UpperCamelCase_    (	self  ,   _UpperCAmelCase=0.9999   ):




										"""simple docstring"""




										__magic_name__					:				List[Any]				=   UNetaDConditionModel.from_pretrained(self.model_id  ,   subfolder='unet'   )
										__magic_name__					:				int				=   unet.to(_UpperCAmelCase   )
										__magic_name__					:				Optional[int]				=   EMAModel(unet.parameters()  ,   decay=_UpperCAmelCase  ,   model_cls=_UpperCAmelCase  ,   model_config=unet.config   )
										return unet, ema_unet
				def 			UpperCamelCase_    (	self   ):




										"""simple docstring"""




										__magic_name__					:				List[str]				=   torch.randn(
										    self.batch_size  ,   self.num_in_channels  ,   self.latent_height  ,   self.latent_width  ,   generator=self.generator   ).to(_UpperCAmelCase   )
										__magic_name__					:				List[str]				=   torch.randint(0  ,   1000  ,   size=(self.batch_size,)  ,   generator=self.generator   ).to(_UpperCAmelCase   )
										__magic_name__					:				Any				=   torch.randn(
										    self.batch_size  ,   self.prompt_length  ,   self.text_encoder_hidden_dim  ,   generator=self.generator   ).to(_UpperCAmelCase   )
										return noisy_latents, timesteps, encoder_hidden_states
				def 			UpperCamelCase_    (	self  ,   _UpperCAmelCase   ):




										"""simple docstring"""




										__magic_name__					:				Any				=   {}
										for k, param in unet.state_dict().items():
																__magic_name__					:				str				=   torch.randn_like(_UpperCAmelCase   ) + (param * torch.randn_like(_UpperCAmelCase   ))
																updated_state_dict.update({k: updated_param}   )
										unet.load_state_dict(_UpperCAmelCase   )
										return unet
				def 			UpperCamelCase_    (	self   ):




										"""simple docstring"""




										__magic_name__		,			__magic_name__					:				int				=   self.get_models()
										# Take the first (hypothetical) EMA step.
										ema_unet.step(unet.parameters()   )
										assert ema_unet.optimization_step == 1

										# Take two more.
										for _ in range(2   ):
																ema_unet.step(unet.parameters()   )
										assert ema_unet.optimization_step == 3
				def 			UpperCamelCase_    (	self   ):




										"""simple docstring"""




										__magic_name__		,			__magic_name__					:				Union[str, Any]				=   self.get_models()
										# Since the `unet` is not being updated (i.e., backprop'd)
										# there won't be any difference between the `params` of `unet`
										# and `ema_unet` even if we call `ema_unet.step(unet.parameters())`.
										ema_unet.step(unet.parameters()   )
										__magic_name__					:				Any				=   list(unet.parameters()   )
										for s_param, param in zip(ema_unet.shadow_params  ,   _UpperCAmelCase   ):
																assert torch.allclose(_UpperCAmelCase  ,   _UpperCAmelCase   )

										# The above holds true even if we call `ema.step()` multiple times since
										# `unet` params are still not being updated.
										for _ in range(4   ):
																ema_unet.step(unet.parameters()   )
										for s_param, param in zip(ema_unet.shadow_params  ,   _UpperCAmelCase   ):
																assert torch.allclose(_UpperCAmelCase  ,   _UpperCAmelCase   )
				def 			UpperCamelCase_    (	self   ):




										"""simple docstring"""




										__magic_name__		,			__magic_name__					:				Optional[Any]				=   self.get_models()
										# Here we simulate the parameter updates for `unet`. Since there might
										# be some parameters which are initialized to zero we take extra care to
										# initialize their values to something non-zero before the multiplication.
										__magic_name__					:				Any				=   self.simulate_backprop(_UpperCAmelCase   )

										# Take the EMA step.
										ema_unet.step(unet_pseudo_updated_step_one.parameters()   )

										# Now the EMA'd parameters won't be equal to the original model parameters.
										__magic_name__					:				Any				=   list(unet_pseudo_updated_step_one.parameters()   )
										for s_param, param in zip(ema_unet.shadow_params  ,   _UpperCAmelCase   ):
																assert ~torch.allclose(_UpperCAmelCase  ,   _UpperCAmelCase   )

										# Ensure this is the case when we take multiple EMA steps.
										for _ in range(4   ):
																ema_unet.step(unet.parameters()   )
										for s_param, param in zip(ema_unet.shadow_params  ,   _UpperCAmelCase   ):
																assert ~torch.allclose(_UpperCAmelCase  ,   _UpperCAmelCase   )
				def 			UpperCamelCase_    (	self   ):




										"""simple docstring"""




										__magic_name__		,			__magic_name__					:				List[str]				=   self.get_models()

										# First backprop + EMA
										__magic_name__					:				int				=   self.simulate_backprop(_UpperCAmelCase   )
										ema_unet.step(unet_step_one.parameters()   )
										__magic_name__					:				Any				=   ema_unet.shadow_params

										# Second backprop + EMA
										__magic_name__					:				Any				=   self.simulate_backprop(_UpperCAmelCase   )
										ema_unet.step(unet_step_two.parameters()   )
										__magic_name__					:				Union[str, Any]				=   ema_unet.shadow_params

										for step_one, step_two in zip(_UpperCAmelCase  ,   _UpperCAmelCase   ):
																assert ~torch.allclose(_UpperCAmelCase  ,   _UpperCAmelCase   )
				def 			UpperCamelCase_    (	self   ):




										"""simple docstring"""




										__magic_name__		,			__magic_name__					:				Optional[Any]				=   self.get_models(decay=0.0   )
										__magic_name__					:				str				=   self.simulate_backprop(_UpperCAmelCase   )
										ema_unet.step(unet_step_one.parameters()   )
										__magic_name__					:				List[Any]				=   ema_unet.shadow_params

										__magic_name__					:				Optional[Any]				=   self.simulate_backprop(_UpperCAmelCase   )
										ema_unet.step(unet_step_two.parameters()   )
										__magic_name__					:				Optional[Any]				=   ema_unet.shadow_params

										for step_one, step_two in zip(_UpperCAmelCase  ,   _UpperCAmelCase   ):
																assert torch.allclose(_UpperCAmelCase  ,   _UpperCAmelCase   )





				@skip_mps
				def 			UpperCamelCase_    (	self   ):




										"""simple docstring"""




										__magic_name__		,			__magic_name__					:				Optional[int]				=   self.get_models()
										__magic_name__		,			__magic_name__		,			__magic_name__					:				str				=   self.get_dummy_inputs()

										with tempfile.TemporaryDirectory() as tmpdir:
																ema_unet.save_pretrained(_UpperCAmelCase   )
																__magic_name__					:				Optional[int]				=   UNetaDConditionModel.from_pretrained(_UpperCAmelCase  ,   model_cls=_UpperCAmelCase   )
																__magic_name__					:				str				=   loaded_unet.to(unet.device   )

										# Since no EMA step has been performed the outputs should match.
										__magic_name__					:				Union[str, Any]				=   unet(_UpperCAmelCase  ,   _UpperCAmelCase  ,   _UpperCAmelCase   ).sample
										__magic_name__					:				List[Any]				=   loaded_unet(_UpperCAmelCase  ,   _UpperCAmelCase  ,   _UpperCAmelCase   ).sample

										assert torch.allclose(_UpperCAmelCase  ,   _UpperCAmelCase  ,   atol=1E-4   )


