"""simple docstring"""







import enum
import inspect
from typing import Iterable, List, Optional, Tuple, Union






class      _SCREAMING_SNAKE_CASE							(			enum.Enum				):



   '''simple docstring'''



   __UpperCAmelCase	:							Optional[Any]    =    "timm"
   __UpperCAmelCase	:							Optional[int]    =    "transformers"



def       A__    (	lowercase_     :	Optional[Iterable[str]]      , lowercase_     :	Optional[Iterable[int]]      , lowercase_     :	Optional[Iterable[str]]						)  ->							Dict:
   if stage_names is None:
      raise ValueError('Stage_names must be set for transformers backbones'						)

   if out_features is not None:
      if not isinstance(lowercase_      , (list,)						):
         raise ValueError(F'''out_features must be a list {type(lowercase_						)}'''						)
      if any(feat not in stage_names for feat in out_features						):
         raise ValueError(F'''out_features must be a subset of stage_names: {stage_names} got {out_features}'''						)

   if out_indices is not None:
      if not isinstance(lowercase_      , (list, tuple)						):
         raise ValueError(F'''out_indices must be a list or tuple, got {type(lowercase_						)}'''						)
      if any(idx >= len(lowercase_						) for idx in out_indices						):
         raise ValueError('out_indices must be valid indices for stage_names {stage_names}, got {out_indices}'						)

   if out_features is not None and out_indices is not None:
      if len(lowercase_						) != len(lowercase_						):
         raise ValueError('out_features and out_indices should have the same length if both are set'						)
      if out_features != [stage_names[idx] for idx in out_indices]:
         raise ValueError('out_features and out_indices should correspond to the same stages if both are set'						)



def       A__    (	lowercase_     :	Optional[List[str]]      , lowercase_     :	Optional[Union[List[int], Tuple[int]]]      , lowercase_     :	List[str]      , )  ->							Optional[int]:
   if out_indices is None and out_features is None:
      __lowerCamelCase    		=					[len(lowercase_						) - 1]
      __lowerCamelCase    		=					[stage_names[-1]]
   elif out_indices is None and out_features is not None:
      __lowerCamelCase    		=					[stage_names.index(lowercase_						) for layer in out_features]
   elif out_features is None and out_indices is not None:
      __lowerCamelCase    		=					[stage_names[idx] for idx in out_indices]
   return out_features, out_indices



def       A__    (	lowercase_     :	Optional[List[str]]      , lowercase_     :	Optional[Union[List[int], Tuple[int]]]      , lowercase_     :	List[str]      , )  ->							Tuple[List[str], List[int]]:
   verify_out_features_out_indices(out_features=lowercase_      , out_indices=lowercase_      , stage_names=lowercase_						)
   __lowerCamelCase		,__lowerCamelCase    		=					_align_output_features_output_indices(
       out_features=lowercase_      , out_indices=lowercase_      , stage_names=lowercase_						)
   # Verify that the aligned out_features and out_indices are valid
   verify_out_features_out_indices(out_features=lowercase_      , out_indices=lowercase_      , stage_names=lowercase_						)
   return output_features, output_indices






class      _SCREAMING_SNAKE_CASE							:



   '''simple docstring'''



   __UpperCAmelCase	:							Optional[BackboneType]    =    None







   def    _A				(	self ,				_SCREAMING_SNAKE_CASE	):


      '''simple docstring'''

      if getattr(self ,				'_backbone' ,				_SCREAMING_SNAKE_CASE	) is None:
         raise ValueError('self._backbone must be set before calling _init_timm_backbone'	)

      # These will diagree with the defaults for the transformers models e.g. for resnet50
      # the transformer model has out_features = ['stem', 'stage1', 'stage2', 'stage3', 'stage4']
      # the timm model has out_features = ['act', 'layer1', 'layer2', 'layer3', 'layer4']
      __lowerCamelCase    		=					[stage['module'] for stage in self._backbone.feature_info.info]
      __lowerCamelCase    		=					[stage['num_chs'] for stage in self._backbone.feature_info.info]
      __lowerCamelCase    		=					self._backbone.feature_info.out_indices
      __lowerCamelCase    		=					self._backbone.feature_info.module_name()

      # We verify the out indices and out features are valid
      verify_out_features_out_indices(
          out_features=_SCREAMING_SNAKE_CASE ,				out_indices=_SCREAMING_SNAKE_CASE ,				stage_names=self.stage_names	)
      __lowerCamelCase		,__lowerCamelCase    		=					out_features, out_indices







   def    _A				(	self ,				_SCREAMING_SNAKE_CASE	):


      '''simple docstring'''

      __lowerCamelCase    		=					getattr(_SCREAMING_SNAKE_CASE ,				'stage_names'	)
      __lowerCamelCase    		=					getattr(_SCREAMING_SNAKE_CASE ,				'out_features' ,				_SCREAMING_SNAKE_CASE	)
      __lowerCamelCase    		=					getattr(_SCREAMING_SNAKE_CASE ,				'out_indices' ,				_SCREAMING_SNAKE_CASE	)

      __lowerCamelCase    		=					stage_names
      __lowerCamelCase		,__lowerCamelCase    		=					get_aligned_output_features_output_indices(
          out_features=_SCREAMING_SNAKE_CASE ,				out_indices=_SCREAMING_SNAKE_CASE ,				stage_names=_SCREAMING_SNAKE_CASE	)
      # Number of channels for each stage. This is set in the transformer backbone model init
      __lowerCamelCase    		=					None







   def    _A				(	self ,				_SCREAMING_SNAKE_CASE	):


      '''simple docstring'''

      __lowerCamelCase    		=					config

      __lowerCamelCase    		=					getattr(_SCREAMING_SNAKE_CASE ,				'use_timm_backbone' ,				_SCREAMING_SNAKE_CASE	)
      __lowerCamelCase    		=					BackboneType.TIMM if self.use_timm_backbone else BackboneType.TRANSFORMERS

      if self.backbone_type == BackboneType.TIMM:
         self._init_timm_backbone(_SCREAMING_SNAKE_CASE	)
      elif self.backbone_type == BackboneType.TRANSFORMERS:
         self._init_transformers_backbone(_SCREAMING_SNAKE_CASE	)
      else:
         raise ValueError(F'''backbone_type {self.backbone_type} not supported.'''	)







   @property
   def    _A				(	self	):


      '''simple docstring'''

      return self._out_features







   @out_features.setter
   def    _A				(	self ,				_SCREAMING_SNAKE_CASE	):


      '''simple docstring'''

      __lowerCamelCase		,__lowerCamelCase    		=					get_aligned_output_features_output_indices(
          out_features=_SCREAMING_SNAKE_CASE ,				out_indices=_SCREAMING_SNAKE_CASE ,				stage_names=self.stage_names	)







   @property
   def    _A				(	self	):


      '''simple docstring'''

      return self._out_indices







   @out_indices.setter
   def    _A				(	self ,				_SCREAMING_SNAKE_CASE	):


      '''simple docstring'''

      __lowerCamelCase		,__lowerCamelCase    		=					get_aligned_output_features_output_indices(
          out_features=_SCREAMING_SNAKE_CASE ,				out_indices=_SCREAMING_SNAKE_CASE ,				stage_names=self.stage_names	)







   @property
   def    _A				(	self	):


      '''simple docstring'''

      return {stage: self.num_features[i] for i, stage in enumerate(self.stage_names	)}







   @property
   def    _A				(	self	):


      '''simple docstring'''

      return [self.out_feature_channels[name] for name in self.out_features]







   def    _A				(	self ,				*_SCREAMING_SNAKE_CASE ,				**_SCREAMING_SNAKE_CASE	):


      '''simple docstring'''

      __lowerCamelCase    		=					dict(inspect.signature(self.forward	).parameters	)
      __lowerCamelCase    		=					{k: v for k, v in kwargs.items() if k in signature}
      return self(*_SCREAMING_SNAKE_CASE ,				**_SCREAMING_SNAKE_CASE	)







   def    _A				(	self ,				_SCREAMING_SNAKE_CASE ,				_SCREAMING_SNAKE_CASE = None ,				_SCREAMING_SNAKE_CASE = None ,				_SCREAMING_SNAKE_CASE = None ,				):


      '''simple docstring'''

      raise NotImplementedError('This method should be implemented by the derived class.'	)







   def    _A				(	self	):


      '''simple docstring'''

      __lowerCamelCase    		=					super().to_dict()
      __lowerCamelCase    		=					output.pop('_out_features'	)
      __lowerCamelCase    		=					output.pop('_out_indices'	)
      return output






class      _SCREAMING_SNAKE_CASE							:



   '''simple docstring'''



   @property
   def    _A				(	self	):


      '''simple docstring'''

      return self._out_features







   @out_features.setter
   def    _A				(	self ,				_SCREAMING_SNAKE_CASE	):


      '''simple docstring'''

      __lowerCamelCase		,__lowerCamelCase    		=					get_aligned_output_features_output_indices(
          out_features=_SCREAMING_SNAKE_CASE ,				out_indices=_SCREAMING_SNAKE_CASE ,				stage_names=self.stage_names	)







   @property
   def    _A				(	self	):


      '''simple docstring'''

      return self._out_indices







   @out_indices.setter
   def    _A				(	self ,				_SCREAMING_SNAKE_CASE	):


      '''simple docstring'''

      __lowerCamelCase		,__lowerCamelCase    		=					get_aligned_output_features_output_indices(
          out_features=_SCREAMING_SNAKE_CASE ,				out_indices=_SCREAMING_SNAKE_CASE ,				stage_names=self.stage_names	)







   def    _A				(	self	):


      '''simple docstring'''

      __lowerCamelCase    		=					super().to_dict()
      __lowerCamelCase    		=					output.pop('_out_features'	)
      __lowerCamelCase    		=					output.pop('_out_indices'	)
      return output









"""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      __UpperCamelCase					(  _A	,						_A	,						unittest.TestCase		):
		SCREAMING_SNAKE_CASE														=       UNetaDModel
		SCREAMING_SNAKE_CASE														=       "sample"


		@property
		def    SCREAMING_SNAKE_CASE__  (self					:						Tuple):
							A     =   4
							A     =   1_4
							A     =   1_6

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

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


		@property
		def    SCREAMING_SNAKE_CASE__  (self					:						Optional[int]):
							return (4, 1_4, 1_6)


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


		def    SCREAMING_SNAKE_CASE__  (self					:						List[Any]):
							pass


		def    SCREAMING_SNAKE_CASE__  (self					:						str):
							pass


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


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


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


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


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


		@unittest.skipIf(torch_device == "mps"	,	"mish op not supported in MPS")
		def    SCREAMING_SNAKE_CASE__  (self					:						List[Any]):
							super().test_output()


		def    SCREAMING_SNAKE_CASE__  (self					:						Union[str, Any]):
							A     =   {
							    "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",
							}
							A     =   self.dummy_input
							return init_dict, inputs_dict


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

							model.to(__SCREAMING_SNAKE_CASE)
							A     =   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    SCREAMING_SNAKE_CASE__  (self					:						Optional[Any]):
							A     =   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)

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

							with torch.no_grad():
												A     =   model(__SCREAMING_SNAKE_CASE	,	__SCREAMING_SNAKE_CASE).sample.permute(0	,	2	,	1)

							A     =   output[0, -3:, -3:].flatten()
							# fmt: off
							A     =   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(__SCREAMING_SNAKE_CASE	,	__SCREAMING_SNAKE_CASE	,	rtol=1E-3))


		def    SCREAMING_SNAKE_CASE__  (self					:						Union[str, Any]):
							# Not implemented yet for this UNet
							pass


		@slow
		def    SCREAMING_SNAKE_CASE__  (self					:						Dict):
							A     =   "harmonai/maestro-150k"
							A     =   UNetaDModel.from_pretrained(__SCREAMING_SNAKE_CASE	,	subfolder="unet")
							model.to(__SCREAMING_SNAKE_CASE)

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

							with torch.no_grad():
												A     =   model(__SCREAMING_SNAKE_CASE	,	__SCREAMING_SNAKE_CASE).sample

							A     =   output.abs().sum()
							A     =   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      __UpperCamelCase					(  _A	,						_A	,						unittest.TestCase		):
		SCREAMING_SNAKE_CASE														=       UNetaDModel
		SCREAMING_SNAKE_CASE														=       "sample"


		@property
		def    SCREAMING_SNAKE_CASE__  (self					:						Any):
							A     =   4
							A     =   1_4
							A     =   1_6

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

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


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


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


		@unittest.skipIf(torch_device == "mps"	,	"mish op not supported in MPS")
		def    SCREAMING_SNAKE_CASE__  (self					:						Optional[Any]):
							super().test_determinism()


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


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


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


		@unittest.skipIf(torch_device == "mps"	,	"mish op not supported in MPS")
		def    SCREAMING_SNAKE_CASE__  (self					:						List[Any]):
							super().test_model_from_pretrained()


		@unittest.skipIf(torch_device == "mps"	,	"mish op not supported in MPS")
		def    SCREAMING_SNAKE_CASE__  (self					:						Any):
							# UNetRL is a value-function is different output shape
							A     ,       A     =   self.prepare_init_args_and_inputs_for_common()
							A     =   self.model_class(**__SCREAMING_SNAKE_CASE)
							model.to(__SCREAMING_SNAKE_CASE)
							model.eval()

							with torch.no_grad():
												A     =   model(**__SCREAMING_SNAKE_CASE)

												if isinstance(__SCREAMING_SNAKE_CASE	,	__SCREAMING_SNAKE_CASE):
																	A     =   output.sample

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


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


		def    SCREAMING_SNAKE_CASE__  (self					:						Optional[int]):
							pass


		def    SCREAMING_SNAKE_CASE__  (self					:						Any):
							A     =   {
							    "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",
							}
							A     =   self.dummy_input
							return init_dict, inputs_dict


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

							value_function.to(__SCREAMING_SNAKE_CASE)
							A     =   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    SCREAMING_SNAKE_CASE__  (self					:						Union[str, Any]):
							A     ,       A     =   UNetaDModel.from_pretrained(
							    "bglick13/hopper-medium-v2-value-function-hor32"	,	output_loading_info=__SCREAMING_SNAKE_CASE	,	subfolder="value_function")
							torch.manual_seed(0)
							if torch.cuda.is_available():
												torch.cuda.manual_seed_all(0)

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

							with torch.no_grad():
												A     =   value_function(__SCREAMING_SNAKE_CASE	,	__SCREAMING_SNAKE_CASE).sample

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


		def    SCREAMING_SNAKE_CASE__  (self					:						Optional[Any]):
							# Not implemented yet for this UNet
							pass





