



'''simple docstring'''




from tempfile import TemporaryDirectory
from unittest import TestCase
from unittest.mock import MagicMock, patch

from transformers import AutoModel, TFAutoModel
from transformers.onnx import FeaturesManager
from transformers.testing_utils import SMALL_MODEL_IDENTIFIER, require_tf, require_torch






@require_torch
@require_tf
class 		_UpperCamelCase						(							lowerCamelCase__					):


			'''simple docstring'''

			def UpperCamelCase__      (							self  :				str		):





					"""simple docstring"""






					__SCREAMING_SNAKE_CASE :	Union[str, Any]					 =	SMALL_MODEL_IDENTIFIER
					__SCREAMING_SNAKE_CASE :	int					 =	"""pt"""
					__SCREAMING_SNAKE_CASE :	str					 =	"""tf"""

			def UpperCamelCase__      (							self  :				Tuple							,  lowerCAmelCase__  :				Tuple		):





					"""simple docstring"""






					__SCREAMING_SNAKE_CASE :	Dict					 =	AutoModel.from_pretrained(self.test_model		)
					model_pt.save_pretrained(lowerCAmelCase__		)

			def UpperCamelCase__      (							self  :				Optional[int]							,  lowerCAmelCase__  :				Optional[Any]		):





					"""simple docstring"""






					__SCREAMING_SNAKE_CASE :	Any					 =	TFAutoModel.from_pretrained(self.test_model							,  from_pt=lowerCAmelCase__		)
					model_tf.save_pretrained(lowerCAmelCase__		)

			def UpperCamelCase__      (							self  :				Any		):





					"""simple docstring"""






					__SCREAMING_SNAKE_CASE :	Optional[int]					 =	"""mock_framework"""

					# Framework provided - return whatever the user provides
					__SCREAMING_SNAKE_CASE :	Dict					 =	FeaturesManager.determine_framework(self.test_model							,  lowerCAmelCase__		)
					self.assertEqual(lowerCAmelCase__							,  lowerCAmelCase__		)

					# Local checkpoint and framework provided - return provided framework
					# PyTorch checkpoint
					with TemporaryDirectory() as local_pt_ckpt:
							self._setup_pt_ckpt(lowerCAmelCase__		)
							__SCREAMING_SNAKE_CASE :	Tuple					 =	FeaturesManager.determine_framework(lowerCAmelCase__							,  lowerCAmelCase__		)
							self.assertEqual(lowerCAmelCase__							,  lowerCAmelCase__		)

					# TensorFlow checkpoint
					with TemporaryDirectory() as local_tf_ckpt:
							self._setup_tf_ckpt(lowerCAmelCase__		)
							__SCREAMING_SNAKE_CASE :	str					 =	FeaturesManager.determine_framework(lowerCAmelCase__							,  lowerCAmelCase__		)
							self.assertEqual(lowerCAmelCase__							,  lowerCAmelCase__		)

			def UpperCamelCase__      (							self  :				Optional[int]		):





					"""simple docstring"""






					with TemporaryDirectory() as local_pt_ckpt:
							self._setup_pt_ckpt(lowerCAmelCase__		)
							__SCREAMING_SNAKE_CASE :	Tuple					 =	FeaturesManager.determine_framework(lowerCAmelCase__		)
							self.assertEqual(lowerCAmelCase__							,  self.framework_pt		)

					# TensorFlow checkpoint
					with TemporaryDirectory() as local_tf_ckpt:
							self._setup_tf_ckpt(lowerCAmelCase__		)
							__SCREAMING_SNAKE_CASE :	Union[str, Any]					 =	FeaturesManager.determine_framework(lowerCAmelCase__		)
							self.assertEqual(lowerCAmelCase__							,  self.framework_tf		)

					# Invalid local checkpoint
					with TemporaryDirectory() as local_invalid_ckpt:
							with self.assertRaises(lowerCAmelCase__		):
									__SCREAMING_SNAKE_CASE :	List[Any]					 =	FeaturesManager.determine_framework(lowerCAmelCase__		)



			def UpperCamelCase__      (							self  :				int		):





					"""simple docstring"""






					__SCREAMING_SNAKE_CASE :	int					 =	MagicMock(return_value=lowerCAmelCase__		)
					with patch("""transformers.onnx.features.is_tf_available"""							,  lowerCAmelCase__		):
							__SCREAMING_SNAKE_CASE :	List[Any]					 =	FeaturesManager.determine_framework(self.test_model		)
							self.assertEqual(lowerCAmelCase__							,  self.framework_pt		)

					# PyTorch not in environment -> use TensorFlow
					__SCREAMING_SNAKE_CASE :	Dict					 =	MagicMock(return_value=lowerCAmelCase__		)
					with patch("""transformers.onnx.features.is_torch_available"""							,  lowerCAmelCase__		):
							__SCREAMING_SNAKE_CASE :	Any					 =	FeaturesManager.determine_framework(self.test_model		)
							self.assertEqual(lowerCAmelCase__							,  self.framework_tf		)

					# Both in environment -> use PyTorch
					__SCREAMING_SNAKE_CASE :	str					 =	MagicMock(return_value=lowerCAmelCase__		)
					__SCREAMING_SNAKE_CASE :	Optional[Any]					 =	MagicMock(return_value=lowerCAmelCase__		)
					with patch("""transformers.onnx.features.is_tf_available"""							,  lowerCAmelCase__		), patch(
					    """transformers.onnx.features.is_torch_available"""							,  lowerCAmelCase__		):
							__SCREAMING_SNAKE_CASE :	Tuple					 =	FeaturesManager.determine_framework(self.test_model		)
							self.assertEqual(lowerCAmelCase__							,  self.framework_pt		)

					# Both not in environment -> raise error
					__SCREAMING_SNAKE_CASE :	List[Any]					 =	MagicMock(return_value=lowerCAmelCase__		)
					__SCREAMING_SNAKE_CASE :	int					 =	MagicMock(return_value=lowerCAmelCase__		)
					with patch("""transformers.onnx.features.is_tf_available"""							,  lowerCAmelCase__		), patch(
					    """transformers.onnx.features.is_torch_available"""							,  lowerCAmelCase__		):
							with self.assertRaises(lowerCAmelCase__		):
									__SCREAMING_SNAKE_CASE :	str					 =	FeaturesManager.determine_framework(self.test_model		)




'''simple docstring'''




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        lowerCAmelCase_     (			_lowerCamelCase:     int      ):
		if hor == 1_28:
				__SCREAMING_SNAKE_CASE :	Any					 =	("""DownResnetBlock1D""", """DownResnetBlock1D""", """DownResnetBlock1D""")
				__SCREAMING_SNAKE_CASE :	List[Any]					 =	(32, 1_28, 2_56)
				__SCREAMING_SNAKE_CASE :	str					 =	("""UpResnetBlock1D""", """UpResnetBlock1D""")

		elif hor == 32:
				__SCREAMING_SNAKE_CASE :	Union[str, Any]					 =	("""DownResnetBlock1D""", """DownResnetBlock1D""", """DownResnetBlock1D""", """DownResnetBlock1D""")
				__SCREAMING_SNAKE_CASE :	str					 =	(32, 64, 1_28, 2_56)
				__SCREAMING_SNAKE_CASE :	Tuple					 =	("""UpResnetBlock1D""", """UpResnetBlock1D""", """UpResnetBlock1D""")
		__SCREAMING_SNAKE_CASE :	Optional[Any]					 =	torch.load(F"/Users/bglickenhaus/Documents/diffuser/temporal_unet-hopper-mediumv2-hor{hor}.torch"      )
		__SCREAMING_SNAKE_CASE :	Any					 =	model.state_dict()
		__SCREAMING_SNAKE_CASE :	Optional[Any]					 =	{
		    """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_55_36,
		    """mid_block_type""": """MidResTemporalBlock1D""",
		    """act_fn""": """mish""",
		}
		__SCREAMING_SNAKE_CASE :	int					 =	UNetaDModel(**_lowerCamelCase      )
		print(F"length of state dict: {len(state_dict.keys()      )}"      )
		print(F"length of value function dict: {len(hf_value_function.state_dict().keys()      )}"      )
		__SCREAMING_SNAKE_CASE :	Optional[int]					 =	dict(zip(model.state_dict().keys() ,   hf_value_function.state_dict().keys()      )      )
		for k, v in mapping.items():
				__SCREAMING_SNAKE_CASE :	Union[str, Any]					 =	state_dict.pop(_lowerCamelCase      )
		hf_value_function.load_state_dict(_lowerCamelCase      )

		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(_lowerCamelCase ,   _lowerCamelCase      )


def        lowerCAmelCase_     (			):
		__SCREAMING_SNAKE_CASE :	Dict					 =	{
		    """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, 1_28, 2_56),
		    """layers_per_block""": 1,
		    """downsample_each_block""": True,
		    """sample_size""": 6_55_36,
		    """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""",
		}

		__SCREAMING_SNAKE_CASE :	Optional[int]					 =	torch.load("""/Users/bglickenhaus/Documents/diffuser/value_function-hopper-mediumv2-hor32.torch"""      )
		__SCREAMING_SNAKE_CASE :	Dict					 =	model
		__SCREAMING_SNAKE_CASE :	List[Any]					 =	UNetaDModel(**_lowerCamelCase      )
		print(F"length of state dict: {len(state_dict.keys()      )}"      )
		print(F"length of value function dict: {len(hf_value_function.state_dict().keys()      )}"      )

		__SCREAMING_SNAKE_CASE :	Optional[Any]					 =	dict(zip(state_dict.keys() ,   hf_value_function.state_dict().keys()      )      )
		for k, v in mapping.items():
				__SCREAMING_SNAKE_CASE :	str					 =	state_dict.pop(_lowerCamelCase      )

		hf_value_function.load_state_dict(_lowerCamelCase      )

		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(_lowerCamelCase ,   _lowerCamelCase      )


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