'''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 	__SCREAMING_SNAKE_CASE		(			lowerCamelCase_		,      lowerCamelCase_		,      unittest.TestCase			):







  '''simple docstring'''






  lowerCamelCase_						:Optional[int]		       =						UNetaDModel
  lowerCamelCase_						:int		       =						'''sample'''


  @property
  def _UpperCamelCase	(					self		):




       '''simple docstring'''

       UpperCAmelCase_				:  Any  = 4
       UpperCAmelCase_				:  Any  = 1_4
       UpperCAmelCase_				:  Dict  = 1_6

       UpperCAmelCase_				:  Union[str, Any]  = floats_tensor((batch_size, num_features, seq_len)		).to(snake_case_		)
       UpperCAmelCase_				:  Union[str, Any]  = torch.tensor([1_0] * batch_size		).to(snake_case_		)

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


  @property
  def _UpperCamelCase	(					self		):




       '''simple docstring'''

       return (4, 1_4, 1_6)


  @property
  def _UpperCamelCase	(					self		):




       '''simple docstring'''

       return (4, 1_4, 1_6)


  def _UpperCamelCase	(					self		):




       '''simple docstring'''

       pass


  def _UpperCamelCase	(					self		):




       '''simple docstring'''

       pass


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




       '''simple docstring'''

       super().test_determinism()


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




       '''simple docstring'''

       super().test_outputs_equivalence()


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




       '''simple docstring'''

       super().test_from_save_pretrained()


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




       '''simple docstring'''

       super().test_from_save_pretrained_variant()


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




       '''simple docstring'''

       super().test_model_from_pretrained()


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




       '''simple docstring'''

       super().test_output()


  def _UpperCamelCase	(					self		):




       '''simple docstring'''

       UpperCAmelCase_				:  Optional[int]  = {
           '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',
       }
       UpperCAmelCase_				:  List[str]  = self.dummy_input
       return init_dict, inputs_dict


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




       '''simple docstring'''

       UpperCAmelCase_  ,     UpperCAmelCase_				:  Any  = UNetaDModel.from_pretrained(
           'bglick13/hopper-medium-v2-value-function-hor32' ,						output_loading_info=snake_case_ ,						subfolder='unet'		)
       self.assertIsNotNone(snake_case_		)
       self.assertEqual(len(loading_info['missing_keys']		) ,						0		)

       model.to(snake_case_		)
       UpperCAmelCase_				:  Union[str, Any]  = 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		):




       '''simple docstring'''

       UpperCAmelCase_				:  str  = 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		)

       UpperCAmelCase_				:  Optional[Any]  = model.config.in_channels
       UpperCAmelCase_				:  List[Any]  = 1_6
       UpperCAmelCase_				:  int  = torch.randn((1, seq_len, num_features)		).permute(
           0 ,						2 ,						1		)  # match original, we can update values and remove
       UpperCAmelCase_				:  List[str]  = torch.full((num_features,) ,						0		)

       with torch.no_grad():
            UpperCAmelCase_				:  Union[str, Any]  = model(snake_case_ ,						snake_case_		).sample.permute(0 ,						2 ,						1		)

       UpperCAmelCase_				:  Tuple  = output[0, -3:, -3:].flatten()
       # fmt: off
       UpperCAmelCase_				:  Any  = torch.tensor([-2.13_71_72, 1.1_42_60_16, 0.3_68_86_87, -0.76_69_22, 0.7_30_31_46, 0.11_03_88_64, -0.4_76_06_33, 0.13_27_01_72, 0.02_59_13_48]		)
       # fmt: on
       self.assertTrue(torch.allclose(snake_case_ ,						snake_case_ ,						rtol=1E-3		)		)


  def _UpperCamelCase	(					self		):




       '''simple docstring'''

       pass






  @slow
  def _UpperCamelCase	(					self		):




       '''simple docstring'''

       UpperCAmelCase_				:  str  = 'harmonai/maestro-150k'
       UpperCAmelCase_				:  Optional[Any]  = UNetaDModel.from_pretrained(snake_case_ ,						subfolder='unet'		)
       model.to(snake_case_		)

       UpperCAmelCase_				:  Optional[int]  = 6_5_5_3_6
       UpperCAmelCase_				:  Optional[Any]  = torch.sin(torch.arange(snake_case_		)[None, None, :].repeat(1 ,						2 ,						1		)		).to(snake_case_		)
       UpperCAmelCase_				:  Dict  = torch.tensor([1]		).to(snake_case_		)

       with torch.no_grad():
            UpperCAmelCase_				:  str  = model(snake_case_ ,						snake_case_		).sample

       UpperCAmelCase_				:  Any  = output.abs().sum()
       UpperCAmelCase_				:  Tuple  = output.abs().max()

       assert (output_sum - 2_24.08_96).abs() < 0.5
       assert (output_max - 0.06_07).abs() < 4E-4







class 	__SCREAMING_SNAKE_CASE		(			lowerCamelCase_		,      lowerCamelCase_		,      unittest.TestCase			):







  '''simple docstring'''






  lowerCamelCase_						:Tuple		       =						UNetaDModel
  lowerCamelCase_						:Any		       =						'''sample'''


  @property
  def _UpperCamelCase	(					self		):




       '''simple docstring'''

       UpperCAmelCase_				:  Dict  = 4
       UpperCAmelCase_				:  int  = 1_4
       UpperCAmelCase_				:  Optional[Any]  = 1_6

       UpperCAmelCase_				:  Dict  = floats_tensor((batch_size, num_features, seq_len)		).to(snake_case_		)
       UpperCAmelCase_				:  Any  = torch.tensor([1_0] * batch_size		).to(snake_case_		)

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


  @property
  def _UpperCamelCase	(					self		):




       '''simple docstring'''

       return (4, 1_4, 1_6)


  @property
  def _UpperCamelCase	(					self		):




       '''simple docstring'''

       return (4, 1_4, 1)


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




       '''simple docstring'''

       super().test_determinism()


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




       '''simple docstring'''

       super().test_outputs_equivalence()


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




       '''simple docstring'''

       super().test_from_save_pretrained()


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




       '''simple docstring'''

       super().test_from_save_pretrained_variant()


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




       '''simple docstring'''

       super().test_model_from_pretrained()


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




       '''simple docstring'''

       UpperCAmelCase_  ,     UpperCAmelCase_				:  Dict  = self.prepare_init_args_and_inputs_for_common()
       UpperCAmelCase_				:  Optional[Any]  = self.model_class(**snake_case_		)
       model.to(snake_case_		)
       model.eval()

       with torch.no_grad():
            UpperCAmelCase_				:  Optional[Any]  = model(**snake_case_		)

            if isinstance(snake_case_ ,						snake_case_		):
                 UpperCAmelCase_				:  Any  = output.sample

       self.assertIsNotNone(snake_case_		)
       UpperCAmelCase_				:  Optional[Any]  = torch.Size((inputs_dict['sample'].shape[0], 1)		)
       self.assertEqual(output.shape ,						snake_case_ ,						'Input and output shapes do not match'		)


  def _UpperCamelCase	(					self		):




       '''simple docstring'''

       pass


  def _UpperCamelCase	(					self		):




       '''simple docstring'''

       pass


  def _UpperCamelCase	(					self		):




       '''simple docstring'''

       UpperCAmelCase_				:  Dict  = {
           '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',
       }
       UpperCAmelCase_				:  Optional[Any]  = self.dummy_input
       return init_dict, inputs_dict


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




       '''simple docstring'''

       UpperCAmelCase_  ,     UpperCAmelCase_				:  int  = UNetaDModel.from_pretrained(
           'bglick13/hopper-medium-v2-value-function-hor32' ,						output_loading_info=snake_case_ ,						subfolder='value_function'		)
       self.assertIsNotNone(snake_case_		)
       self.assertEqual(len(vf_loading_info['missing_keys']		) ,						0		)

       value_function.to(snake_case_		)
       UpperCAmelCase_				:  List[str]  = 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		):




       '''simple docstring'''

       UpperCAmelCase_  ,     UpperCAmelCase_				:  Union[str, Any]  = UNetaDModel.from_pretrained(
           'bglick13/hopper-medium-v2-value-function-hor32' ,						output_loading_info=snake_case_ ,						subfolder='value_function'		)
       torch.manual_seed(0		)
       if torch.cuda.is_available():
            torch.cuda.manual_seed_all(0		)

       UpperCAmelCase_				:  List[Any]  = value_function.config.in_channels
       UpperCAmelCase_				:  Tuple  = 1_4
       UpperCAmelCase_				:  Any  = torch.randn((1, seq_len, num_features)		).permute(
           0 ,						2 ,						1		)  # match original, we can update values and remove
       UpperCAmelCase_				:  Union[str, Any]  = torch.full((num_features,) ,						0		)

       with torch.no_grad():
            UpperCAmelCase_				:  Tuple  = value_function(snake_case_ ,						snake_case_		).sample

       # fmt: off
       UpperCAmelCase_				:  Tuple  = torch.tensor([1_65.25] * seq_len		)
       # fmt: on
       self.assertTrue(torch.allclose(snake_case_ ,						snake_case_ ,						rtol=1E-3		)		)






  def _UpperCamelCase	(					self		):




       '''simple docstring'''

       pass











'''simple docstring'''



import inspect
from typing import Callable, List, Optional, Set, Tuple, Union

import torch
from packaging import version
from safetensors.torch import storage_ptr, storage_size
from torch import nn

from .utils import logging


_lowercase       :						int						=[nn.LayerNorm]

_lowercase       :						Tuple						=logging.get_logger(__name__)

_lowercase       :						List[Any]						=version.parse(version.parse(torch.__version__).base_version)

_lowercase       :						int						=parsed_torch_version_base >= version.parse("2.0")
_lowercase       :						Tuple						=parsed_torch_version_base >= version.parse("1.12")
_lowercase       :						int						=parsed_torch_version_base >= version.parse("1.11")
_lowercase       :						Dict						=parsed_torch_version_base < version.parse("1.11")


def 							__UpperCAmelCase				(				UpperCamelCase__				:Dict   ,				UpperCamelCase__				:List[str]   ,				UpperCamelCase__				:str   ,				UpperCamelCase__				:Any   ,				self				:Optional[int]			)		->						Union[str, Any]:
				from torch import _softmax_backward_data

				if is_torch_less_than_1_11:
								return _softmax_backward_data(UpperCamelCase__   ,				UpperCamelCase__   ,				parent.dim   ,				self			)
				else:
								return _softmax_backward_data(UpperCamelCase__   ,				UpperCamelCase__   ,				parent.dim   ,				self.dtype			)


def 							__UpperCAmelCase				(				UpperCamelCase__				:nn.Linear   ,				UpperCamelCase__				:torch.LongTensor   ,				UpperCamelCase__				:int = 0			)		->						nn.Linear:
				snake_case__      :   Optional[Any]      =    index.to(layer.weight.device			)
				snake_case__      :   int      =    layer.weight.index_select(UpperCamelCase__   ,				UpperCamelCase__			).clone().detach()
				if layer.bias is not None:
								if dim == 1:
												snake_case__      :   Dict      =    layer.bias.clone().detach()
								else:
												snake_case__      :   Optional[Any]      =    layer.bias[index].clone().detach()
				snake_case__      :   Union[str, Any]      =    list(layer.weight.size()			)
				snake_case__      :   int      =    len(UpperCamelCase__			)
				snake_case__      :   int      =    nn.Linear(new_size[1]   ,				new_size[0]   ,				bias=layer.bias is not None			).to(layer.weight.device			)
				snake_case__      :   Union[str, Any]      =    False
				new_layer.weight.copy_(W.contiguous()			)
				snake_case__      :   List[Any]      =    True
				if layer.bias is not None:
								snake_case__      :   Union[str, Any]      =    False
								new_layer.bias.copy_(b.contiguous()			)
								snake_case__      :   Dict      =    True
				return new_layer






class 			_SCREAMING_SNAKE_CASE				(nn.Module      ):



			def __init__(     self		:						int    ,			__UpperCamelCase		:						Tuple    ,			__UpperCamelCase		:						Optional[Any]			)      ->      Union[str, Any]:
							"""simple docstring"""





							super().__init__()
							snake_case__      :   str      =    nf
							snake_case__      :   Tuple      =    nn.Parameter(torch.empty(__UpperCamelCase    ,			__UpperCamelCase			)			)
							snake_case__      :   Optional[Any]      =    nn.Parameter(torch.zeros(__UpperCamelCase			)			)
							nn.init.normal_(self.weight    ,			std=0.02			)







			def 				lowerCAmelCase (     self		:						Dict    ,			__UpperCamelCase		:						Optional[int]			)      ->      Dict:
							"""simple docstring"""





							snake_case__      :   Dict      =    x.size()[:-1] + (self.nf,)
							snake_case__      :   List[str]      =    torch.addmm(self.bias    ,			x.view(-1    ,			x.size(-1			)			)    ,			self.weight			)
							snake_case__      :   Dict      =    x.view(__UpperCamelCase			)
							return x



def 							__UpperCAmelCase				(				UpperCamelCase__				:ConvaD   ,				UpperCamelCase__				:torch.LongTensor   ,				UpperCamelCase__				:int = 1			)		->						ConvaD:
				snake_case__      :   Optional[Any]      =    index.to(layer.weight.device			)
				snake_case__      :   List[Any]      =    layer.weight.index_select(UpperCamelCase__   ,				UpperCamelCase__			).clone().detach()
				if dim == 0:
								snake_case__      :   Any      =    layer.bias.clone().detach()
				else:
								snake_case__      :   Optional[Any]      =    layer.bias[index].clone().detach()
				snake_case__      :   Union[str, Any]      =    list(layer.weight.size()			)
				snake_case__      :   List[str]      =    len(UpperCamelCase__			)
				snake_case__      :   Optional[Any]      =    ConvaD(new_size[1]   ,				new_size[0]			).to(layer.weight.device			)
				snake_case__      :   List[str]      =    False
				new_layer.weight.copy_(W.contiguous()			)
				snake_case__      :   List[Any]      =    True
				snake_case__      :   List[str]      =    False
				new_layer.bias.copy_(b.contiguous()			)
				snake_case__      :   str      =    True
				return new_layer


def 							__UpperCAmelCase				(				UpperCamelCase__				:Union[nn.Linear, ConvaD]   ,				UpperCamelCase__				:torch.LongTensor   ,				UpperCamelCase__				:Optional[int] = None			)		->						Union[nn.Linear, ConvaD]:
				if isinstance(UpperCamelCase__   ,				nn.Linear			):
								return prune_linear_layer(UpperCamelCase__   ,				UpperCamelCase__   ,				dim=0 if dim is None else dim			)
				elif isinstance(UpperCamelCase__   ,				UpperCamelCase__			):
								return prune_convad_layer(UpperCamelCase__   ,				UpperCamelCase__   ,				dim=1 if dim is None else dim			)
				else:
								raise ValueError(F'''Can\'t prune layer of class {layer.__class__}'''			)


def 							__UpperCAmelCase				(				UpperCamelCase__				:Callable[..., torch.Tensor]   ,				UpperCamelCase__				:int   ,				UpperCamelCase__				:int   ,				*UpperCamelCase__				:List[str]			)		->						torch.Tensor:
				assert len(UpperCamelCase__			) > 0, F'''{input_tensors} has to be a tuple/list of tensors'''

				# inspect.signature exist since python 3.5 and is a python method -> no problem with backward compatibility
				snake_case__      :   str      =    len(inspect.signature(UpperCamelCase__			).parameters			)
				if num_args_in_forward_chunk_fn != len(UpperCamelCase__			):
								raise ValueError(
								    F'''forward_chunk_fn expects {num_args_in_forward_chunk_fn} arguments, but only {len(UpperCamelCase__			)} input '''
								    '''tensors are given'''			)

				if chunk_size > 0:
								snake_case__      :   Tuple      =    input_tensors[0].shape[chunk_dim]
								for input_tensor in input_tensors:
												if input_tensor.shape[chunk_dim] != tensor_shape:
																raise ValueError(
																    F'''All input tenors have to be of the same shape: {tensor_shape}, '''
																    F'''found shape {input_tensor.shape[chunk_dim]}'''			)

								if input_tensors[0].shape[chunk_dim] % chunk_size != 0:
												raise ValueError(
												    F'''The dimension to be chunked {input_tensors[0].shape[chunk_dim]} has to be a multiple of the chunk '''
												    F'''size {chunk_size}'''			)

								snake_case__      :   Optional[Any]      =    input_tensors[0].shape[chunk_dim] // chunk_size

								# chunk input tensor into tuples
								snake_case__      :   List[Any]      =    tuple(input_tensor.chunk(UpperCamelCase__   ,				dim=UpperCamelCase__			) for input_tensor in input_tensors			)
								# apply forward fn to every tuple
								snake_case__      :   List[Any]      =    tuple(forward_fn(*UpperCamelCase__			) for input_tensors_chunk in zip(*UpperCamelCase__			)			)
								# concatenate output at same dimension
								return torch.cat(UpperCamelCase__   ,				dim=UpperCamelCase__			)

				return forward_fn(*UpperCamelCase__			)


def 							__UpperCAmelCase				(				UpperCamelCase__				:List[int]   ,				UpperCamelCase__				:int   ,				UpperCamelCase__				:int   ,				UpperCamelCase__				:Set[int]			)		->						Tuple[Set[int], torch.LongTensor]:
				snake_case__      :   str      =    torch.ones(UpperCamelCase__   ,				UpperCamelCase__			)
				snake_case__      :   int      =    set(UpperCamelCase__			) - already_pruned_heads  # Convert to set and remove already pruned heads
				for head in heads:
								# Compute how many pruned heads are before the head and move the index accordingly
								snake_case__      :   str      =    head - sum(1 if h < head else 0 for h in already_pruned_heads			)
								snake_case__      :   Optional[int]      =    0
				snake_case__      :   Optional[Any]      =    mask.view(-1			).contiguous().eq(1			)
				snake_case__      :   torch.LongTensor      =    torch.arange(len(UpperCamelCase__			)			)[mask].long()
				return heads, index


def 							__UpperCAmelCase				(				*UpperCamelCase__				:Union[torch.Tensor, List[torch.Tensor]]   ,				UpperCamelCase__				:Optional[str] = None			)		->						Tuple[torch.Tensor, ...]:
				return torch.meshgrid(*UpperCamelCase__   ,				indexing=UpperCamelCase__			)



def 							__UpperCAmelCase				(				UpperCamelCase__				:torch.Tensor			)		->						Tuple[torch.device, int, int]:
				return tensor.device, storage_ptr(UpperCamelCase__			), storage_size(UpperCamelCase__			)




