


import json
import os

import torch

from diffusers import UNetaDModel


os.makedirs("hub/hopper-medium-v2/unet/hor32", exist_ok=True)
os.makedirs("hub/hopper-medium-v2/unet/hor128", exist_ok=True)

os.makedirs("hub/hopper-medium-v2/value_function", exist_ok=True)






def __UpperCamelCase   (			__lowerCamelCase		: int				)      ->     Dict:







     '''simple docstring'''






     if hor == 128:
          _a						=		('DownResnetBlock1D', 'DownResnetBlock1D', 'DownResnetBlock1D')
          _a						=		(32, 128, 256)
          _a						=		('UpResnetBlock1D', 'UpResnetBlock1D')

     elif hor == 32:
          _a						=		('DownResnetBlock1D', 'DownResnetBlock1D', 'DownResnetBlock1D', 'DownResnetBlock1D')
          _a						=		(32, 64, 128, 256)
          _a						=		('UpResnetBlock1D', 'UpResnetBlock1D', 'UpResnetBlock1D')
     _a						=		torch.load(F"/Users/bglickenhaus/Documents/diffuser/temporal_unet-hopper-mediumv2-hor{hor}.torch"				)
     _a						=		model.state_dict()
     _a						=		{
         'down_block_types': down_block_types,
         'block_out_channels': block_out_channels,
         'up_block_types': up_block_types,
         'layers_per_block': 1,
         'use_timestep_embedding': True,
         'out_block_type': 'OutConv1DBlock',
         'norm_num_groups': 8,
         'downsample_each_block': False,
         'in_channels': 14,
         'out_channels': 14,
         'extra_in_channels': 0,
         'time_embedding_type': 'positional',
         'flip_sin_to_cos': False,
         'freq_shift': 1,
         'sample_size': 6_5536,
         'mid_block_type': 'MidResTemporalBlock1D',
         'act_fn': 'mish',
     }
     _a						=		UNetaDModel(**SCREAMING_SNAKE_CASE_				)
     print(F"length of state dict: {len(state_dict.keys()				)}"				)
     print(F"length of value function dict: {len(hf_value_function.state_dict().keys()				)}"				)
     _a						=		dict(zip(model.state_dict().keys()      ,  hf_value_function.state_dict().keys()				)				)
     for k, v in mapping.items():
          _a						=		state_dict.pop(SCREAMING_SNAKE_CASE_				)
     hf_value_function.load_state_dict(SCREAMING_SNAKE_CASE_				)

     torch.save(hf_value_function.state_dict()      ,  F"hub/hopper-medium-v2/unet/hor{hor}/diffusion_pytorch_model.bin"				)
     with open(F"hub/hopper-medium-v2/unet/hor{hor}/config.json"      ,  "w"				) as f:
          json.dump(SCREAMING_SNAKE_CASE_      ,  SCREAMING_SNAKE_CASE_				)






def __UpperCamelCase   (			)      ->     Optional[Any]:







     '''simple docstring'''






     _a						=		{
         'in_channels': 14,
         'down_block_types': ('DownResnetBlock1D', 'DownResnetBlock1D', 'DownResnetBlock1D', 'DownResnetBlock1D'),
         'up_block_types': (),
         'out_block_type': 'ValueFunction',
         'mid_block_type': 'ValueFunctionMidBlock1D',
         'block_out_channels': (32, 64, 128, 256),
         'layers_per_block': 1,
         'downsample_each_block': True,
         'sample_size': 6_5536,
         'out_channels': 14,
         'extra_in_channels': 0,
         'time_embedding_type': 'positional',
         'use_timestep_embedding': True,
         'flip_sin_to_cos': False,
         'freq_shift': 1,
         'norm_num_groups': 8,
         'act_fn': 'mish',
     }

     _a						=		torch.load("/Users/bglickenhaus/Documents/diffuser/value_function-hopper-mediumv2-hor32.torch"				)
     _a						=		model
     _a						=		UNetaDModel(**SCREAMING_SNAKE_CASE_				)
     print(F"length of state dict: {len(state_dict.keys()				)}"				)
     print(F"length of value function dict: {len(hf_value_function.state_dict().keys()				)}"				)

     _a						=		dict(zip(state_dict.keys()      ,  hf_value_function.state_dict().keys()				)				)
     for k, v in mapping.items():
          _a						=		state_dict.pop(SCREAMING_SNAKE_CASE_				)

     hf_value_function.load_state_dict(SCREAMING_SNAKE_CASE_				)

     torch.save(hf_value_function.state_dict()      ,  "hub/hopper-medium-v2/value_function/diffusion_pytorch_model.bin"				)
     with open("hub/hopper-medium-v2/value_function/config.json"      ,  "w"				) as f:
          json.dump(SCREAMING_SNAKE_CASE_      ,  SCREAMING_SNAKE_CASE_				)


if __name__ == "__main__":
   unet(32)
   # unet(128)
   value_function()





'''simple docstring'''

import inspect
import unittest

from huggingface_hub import hf_hub_download

from transformers import ASTConfig
from transformers.testing_utils import require_torch, require_torchaudio, slow, torch_device
from transformers.utils import cached_property, is_torch_available, is_torchaudio_available

from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, floats_tensor, ids_tensor
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
   import torch
   from torch import nn

   from transformers import ASTForAudioClassification, ASTModel
   from transformers.models.audio_spectrogram_transformer.modeling_audio_spectrogram_transformer import (
       AUDIO_SPECTROGRAM_TRANSFORMER_PRETRAINED_MODEL_ARCHIVE_LIST,
   )


if is_torchaudio_available():
   import torchaudio

   from transformers import ASTFeatureExtractor





class 						__SCREAMING_SNAKE_CASE					:






     def __init__(			self     ,							__UpperCamelCase     ,							__UpperCamelCase=13     ,							__UpperCamelCase=2     ,							__UpperCamelCase=24     ,							__UpperCamelCase=16     ,							__UpperCamelCase=True     ,							__UpperCamelCase=True     ,							__UpperCamelCase=32     ,							__UpperCamelCase=5     ,							__UpperCamelCase=4     ,							__UpperCamelCase=37     ,							__UpperCamelCase="gelu"     ,							__UpperCamelCase=0.1     ,							__UpperCamelCase=0.1     ,							__UpperCamelCase=10     ,							__UpperCamelCase=0.02     ,							__UpperCamelCase=None     ,							__UpperCamelCase=2     ,							__UpperCamelCase=2     ,							)       ->    Optional[int]:
          _a						=		parent
          _a						=		batch_size
          _a						=		patch_size
          _a						=		max_length
          _a						=		num_mel_bins
          _a						=		is_training
          _a						=		use_labels
          _a						=		hidden_size
          _a						=		num_hidden_layers
          _a						=		num_attention_heads
          _a						=		intermediate_size
          _a						=		hidden_act
          _a						=		hidden_dropout_prob
          _a						=		attention_probs_dropout_prob
          _a						=		type_sequence_label_size
          _a						=		initializer_range
          _a						=		scope
          _a						=		frequency_stride
          _a						=		time_stride

          # in AST, the seq length equals the number of patches + 2 (we add 2 for the [CLS] and distillation tokens)
          _a						=		(self.num_mel_bins - self.patch_size) // self.frequency_stride + 1
          _a						=		(self.max_length - self.patch_size) // self.time_stride + 1
          _a						=		frequency_out_dimension * time_out_dimension
          _a						=		num_patches + 2






     def  a_					(			self		)       ->    List[str]:
          _a						=		floats_tensor([self.batch_size, self.max_length, self.num_mel_bins]		)

          _a						=		None
          if self.use_labels:
               _a						=		ids_tensor([self.batch_size]     ,							self.type_sequence_label_size		)

          _a						=		self.get_config()

          return config, input_values, labels






     def  a_					(			self		)       ->    Optional[int]:
          return ASTConfig(
              patch_size=self.patch_size     ,							max_length=self.max_length     ,							num_mel_bins=self.num_mel_bins     ,							hidden_size=self.hidden_size     ,							num_hidden_layers=self.num_hidden_layers     ,							num_attention_heads=self.num_attention_heads     ,							intermediate_size=self.intermediate_size     ,							hidden_act=self.hidden_act     ,							hidden_dropout_prob=self.hidden_dropout_prob     ,							attention_probs_dropout_prob=self.attention_probs_dropout_prob     ,							is_decoder=__UpperCamelCase     ,							initializer_range=self.initializer_range     ,							frequency_stride=self.frequency_stride     ,							time_stride=self.time_stride     ,							)






     def  a_					(			self     ,							__UpperCamelCase     ,							__UpperCamelCase     ,							__UpperCamelCase		)       ->    Optional[Any]:
          _a						=		ASTModel(config=__UpperCamelCase		)
          model.to(__UpperCamelCase		)
          model.eval()
          _a						=		model(__UpperCamelCase		)
          self.parent.assertEqual(result.last_hidden_state.shape     ,							(self.batch_size, self.seq_length, self.hidden_size)		)






     def  a_					(			self		)       ->    str:
          _a						=		self.prepare_config_and_inputs()
          (
              (
              _a
          )			,			(
              _a
          )			,			(
              _a
          )			,			
          )						=		config_and_inputs
          _a						=		{"input_values": input_values}
          return config, inputs_dict







@require_torch
class 						__SCREAMING_SNAKE_CASE					(      lowerCamelCase__ ,			lowerCamelCase__ ,			unittest.TestCase ):
     UpperCAmelCase 					=				(
         (
             ASTModel,
             ASTForAudioClassification,
         )
         if is_torch_available()
         else ()
     )
     UpperCAmelCase 					=				(
         {'''audio-classification''': ASTForAudioClassification, '''feature-extraction''': ASTModel}
         if is_torch_available()
         else {}
     )
     UpperCAmelCase 					=				False
     UpperCAmelCase 					=				False
     UpperCAmelCase 					=				False
     UpperCAmelCase 					=				False






     def  a_					(			self     ,							__UpperCamelCase     ,							__UpperCamelCase     ,							__UpperCamelCase     ,							__UpperCamelCase     ,							__UpperCamelCase		)       ->    List[str]:
          if pipeline_test_casse_name == "AudioClassificationPipelineTests":
               return True

          return False






     def  a_					(			self		)       ->    int:
          _a						=		ASTModelTester(self		)
          _a						=		ConfigTester(self     ,							config_class=__UpperCamelCase     ,							has_text_modality=__UpperCamelCase     ,							hidden_size=37		)






     def  a_					(			self		)       ->    str:
          self.config_tester.run_common_tests()






     @unittest.skip(reason="AST does not use inputs_embeds"		)
     def  a_					(			self		)       ->    Union[str, Any]:
          pass






     def  a_					(			self		)       ->    List[Any]:
          _a			,			_a						=		self.model_tester.prepare_config_and_inputs_for_common()

          for model_class in self.all_model_classes:
               _a						=		model_class(__UpperCamelCase		)
               self.assertIsInstance(model.get_input_embeddings()     ,							(nn.Module)		)
               _a						=		model.get_output_embeddings()
               self.assertTrue(x is None or isinstance(__UpperCamelCase     ,							nn.Linear		)		)






     def  a_					(			self		)       ->    Any:
          _a			,			_a						=		self.model_tester.prepare_config_and_inputs_for_common()

          for model_class in self.all_model_classes:
               _a						=		model_class(__UpperCamelCase		)
               _a						=		inspect.signature(model.forward		)
               # signature.parameters is an OrderedDict => so arg_names order is deterministic
               _a						=		[*signature.parameters.keys()]

               _a						=		["input_values"]
               self.assertListEqual(arg_names[:1]     ,							__UpperCamelCase		)






     def  a_					(			self		)       ->    int:
          _a						=		self.model_tester.prepare_config_and_inputs()
          self.model_tester.create_and_check_model(*__UpperCamelCase		)






     @slow
     def  a_					(			self		)       ->    Optional[Any]:
          for model_name in AUDIO_SPECTROGRAM_TRANSFORMER_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
               _a						=		ASTModel.from_pretrained(__UpperCamelCase		)
               self.assertIsNotNone(__UpperCamelCase		)







def __UpperCamelCase   (			)      ->     Dict:







     '''simple docstring'''






     _a						=		hf_hub_download(
         repo_id="nielsr/audio-spectogram-transformer-checkpoint"      ,  filename="sample_audio.flac"      ,  repo_type="dataset"				)

     _a			,			_a						=		torchaudio.load(__lowerCamelCase				)

     return audio, sampling_rate






@require_torch
@require_torchaudio
class 						__SCREAMING_SNAKE_CASE					(      unittest.TestCase ):






     @cached_property
     def  a_					(			self		)       ->    int:
          return (
              ASTFeatureExtractor.from_pretrained("MIT/ast-finetuned-audioset-10-10-0.4593"		)
              if is_torchaudio_available()
              else None
          )






     @slow
     def  a_					(			self		)       ->    Dict:
          _a						=		self.default_feature_extractor
          _a						=		ASTForAudioClassification.from_pretrained("MIT/ast-finetuned-audioset-10-10-0.4593"		).to(__UpperCamelCase		)

          _a						=		self.default_feature_extractor
          _a			,			_a						=		prepare_audio()
          _a						=		audio.squeeze().numpy()
          _a						=		feature_extractor(__UpperCamelCase     ,							sampling_rate=__UpperCamelCase     ,							return_tensors="pt"		).to(__UpperCamelCase		)

          # forward pass
          with torch.no_grad():
               _a						=		model(**__UpperCamelCase		)

          # verify the logits
          _a						=		torch.Size((1, 527)		)
          self.assertEqual(outputs.logits.shape     ,							__UpperCamelCase		)

          _a						=		torch.tensor([-0.87_60, -7.00_42, -8.66_02]		).to(__UpperCamelCase		)

          self.assertTrue(torch.allclose(outputs.logits[0, :3]     ,							__UpperCamelCase     ,							atol=1e-4		)		)

