


import argparse
import os


UpperCamelCase			:					Union[str, Any]         =							{
    '''pytorch''': {
        '''1.13''': {
            '''torch''': '''1.13.1''',
            '''torchvision''': '''0.14.1''',
            '''torchaudio''': '''0.13.1''',
            '''python''': 3.9,
            '''cuda''': '''cu116''',
            '''install''': (
                '''python3 -m pip install --no-cache-dir -U torch==1.13.1 torchvision==0.14.1 torchaudio==0.13.1'''
                ''' --extra-index-url https://download.pytorch.org/whl/cu116'''
            ),
            '''base_image''': '''nvidia/cuda:11.6.2-cudnn8-devel-ubuntu20.04''',
        },
        '''1.12''': {
            '''torch''': '''1.12.1''',
            '''torchvision''': '''0.13.1''',
            '''torchaudio''': '''0.12.1''',
            '''python''': 3.9,
            '''cuda''': '''cu113''',
            '''install''': (
                '''python3 -m pip install --no-cache-dir -U torch==1.12.1 torchvision==0.13.1 torchaudio==0.12.1'''
                ''' --extra-index-url https://download.pytorch.org/whl/cu113'''
            ),
            '''base_image''': '''nvidia/cuda:11.2.2-cudnn8-devel-ubuntu20.04''',
        },
        '''1.11''': {
            '''torch''': '''1.11.0''',
            '''torchvision''': '''0.12.0''',
            '''torchaudio''': '''0.11.0''',
            '''python''': 3.9,
            '''cuda''': '''cu113''',
            '''install''': (
                '''python3 -m pip install --no-cache-dir -U torch==1.11.0 torchvision==0.12.0 torchaudio==0.11.0'''
                ''' --extra-index-url https://download.pytorch.org/whl/cu113'''
            ),
            '''base_image''': '''nvidia/cuda:11.2.2-cudnn8-devel-ubuntu20.04''',
        },
        '''1.10''': {
            '''torch''': '''1.10.2''',
            '''torchvision''': '''0.11.3''',
            '''torchaudio''': '''0.10.2''',
            '''python''': 3.9,
            '''cuda''': '''cu113''',
            '''install''': (
                '''python3 -m pip install --no-cache-dir -U torch==1.10.2 torchvision==0.11.3 torchaudio==0.10.2'''
                ''' --extra-index-url https://download.pytorch.org/whl/cu113'''
            ),
            '''base_image''': '''nvidia/cuda:11.2.2-cudnn8-devel-ubuntu20.04''',
        },
        # torchaudio < 0.10 has no CUDA-enabled binary distributions
        '''1.9''': {
            '''torch''': '''1.9.1''',
            '''torchvision''': '''0.10.1''',
            '''torchaudio''': '''0.9.1''',
            '''python''': 3.9,
            '''cuda''': '''cu111''',
            '''install''': (
                '''python3 -m pip install --no-cache-dir -U torch==1.9.1 torchvision==0.10.1 torchaudio==0.9.1'''
                ''' --extra-index-url https://download.pytorch.org/whl/cu111'''
            ),
            '''base_image''': '''nvidia/cuda:11.2.2-cudnn8-devel-ubuntu20.04''',
        },
    },
    '''tensorflow''': {
        '''2.11''': {
            '''tensorflow''': '''2.11.1''',
            '''install''': '''python3 -m pip install --no-cache-dir -U tensorflow==2.11.1''',
            '''base_image''': '''nvidia/cuda:11.2.2-cudnn8-devel-ubuntu20.04''',
        },
        '''2.10''': {
            '''tensorflow''': '''2.10.1''',
            '''install''': '''python3 -m pip install --no-cache-dir -U tensorflow==2.10.1''',
            '''base_image''': '''nvidia/cuda:11.2.2-cudnn8-devel-ubuntu20.04''',
        },
        '''2.9''': {
            '''tensorflow''': '''2.9.3''',
            '''install''': '''python3 -m pip install --no-cache-dir -U tensorflow==2.9.3''',
            '''base_image''': '''nvidia/cuda:11.2.2-cudnn8-devel-ubuntu20.04''',
        },
        '''2.8''': {
            '''tensorflow''': '''2.8.2''',
            '''install''': '''python3 -m pip install --no-cache-dir -U tensorflow==2.8.2''',
            '''base_image''': '''nvidia/cuda:11.2.2-cudnn8-devel-ubuntu20.04''',
        },
        '''2.7''': {
            '''tensorflow''': '''2.7.3''',
            '''install''': '''python3 -m pip install --no-cache-dir -U tensorflow==2.7.3''',
            '''base_image''': '''nvidia/cuda:11.2.2-cudnn8-devel-ubuntu20.04''',
        },
        '''2.6''': {
            '''tensorflow''': '''2.6.5''',
            '''install''': '''python3 -m pip install --no-cache-dir -U tensorflow==2.6.5''',
            '''base_image''': '''nvidia/cuda:11.2.2-cudnn8-devel-ubuntu20.04''',
        },
        '''2.5''': {
            '''tensorflow''': '''2.5.3''',
            '''install''': '''python3 -m pip install --no-cache-dir -U tensorflow==2.5.3''',
            '''base_image''': '''nvidia/cuda:11.2.2-cudnn8-devel-ubuntu20.04''',
        },
    },
}


if __name__ == "__main__":
					UpperCamelCase			:					Optional[int]         =							argparse.ArgumentParser('''Choose the framework and version to install''')
					parser.add_argument(
					    '''--framework''', help='''The framework to install. Should be `torch` or `tensorflow`''', type=str, required=True
					)
					parser.add_argument('''--version''', help='''The version of the framework to install.''', type=str, required=True)
					UpperCamelCase			:					Optional[int]         =							parser.parse_args()

					UpperCamelCase			:					int         =							past_versions_testing[args.framework][args.version]

					os.system(f'''echo \"export INSTALL_CMD=\'{info['install']}\'\" >> ~/.profile''')
					print(f'''echo \"export INSTALL_CMD=\'{info['install']}\'\" >> ~/.profile''')

					UpperCamelCase			:					Optional[Any]         =							''''''
					if args.framework == "pytorch":
										UpperCamelCase			:					Optional[Any]         =							info['''cuda''']
					os.system(f'''echo \"export CUDA=\'{cuda}\'\" >> ~/.profile''')
					print(f'''echo \"export CUDA=\'{cuda}\'\" >> ~/.profile''')








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						( _lowercase			,							_lowercase			,							unittest.TestCase       ):
					__A        = UNetaDModel
					__A        = "sample"




					@property
					def 						snake_case_		(    self							)    ->	Any:
							UpperCamelCase__								=						4
							UpperCamelCase__								=						14
							UpperCamelCase__								=						16

							UpperCamelCase__								=						floats_tensor((batch_size, num_features, seq_len)							).to(__lowerCAmelCase							)
							UpperCamelCase__								=						torch.tensor([10] * batch_size							).to(__lowerCAmelCase							)

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




					@property
					def 						snake_case_		(    self							)    ->	Optional[int]:
							return (4, 14, 16)




					@property
					def 						snake_case_		(    self							)    ->	Any:
							return (4, 14, 16)




					def 						snake_case_		(    self							)    ->	Any:
							pass




					def 						snake_case_		(    self							)    ->	Tuple:
							pass




					@unittest.skipIf(torch_device == """mps"""     ,   """mish op not supported in MPS"""							)
					def 						snake_case_		(    self							)    ->	Any:
							super().test_determinism()




					@unittest.skipIf(torch_device == """mps"""     ,   """mish op not supported in MPS"""							)
					def 						snake_case_		(    self							)    ->	int:
							super().test_outputs_equivalence()




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




					@unittest.skipIf(torch_device == """mps"""     ,   """mish op not supported in MPS"""							)
					def 						snake_case_		(    self							)    ->	Union[str, Any]:
							super().test_from_save_pretrained_variant()




					@unittest.skipIf(torch_device == """mps"""     ,   """mish op not supported in MPS"""							)
					def 						snake_case_		(    self							)    ->	Optional[int]:
							super().test_model_from_pretrained()




					@unittest.skipIf(torch_device == """mps"""     ,   """mish op not supported in MPS"""							)
					def 						snake_case_		(    self							)    ->	Dict:
							super().test_output()




					def 						snake_case_		(    self							)    ->	int:
							UpperCamelCase__								=						{
							    """block_out_channels""": (32, 64, 128, 256),
							    """in_channels""": 14,
							    """out_channels""": 14,
							    """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__								=						self.dummy_input
							return init_dict, inputs_dict




					@unittest.skipIf(torch_device == """mps"""     ,   """mish op not supported in MPS"""							)
					def 						snake_case_		(    self							)    ->	List[Any]:
							UpperCamelCase__	,      UpperCamelCase__								=						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							)
							UpperCamelCase__								=						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 						snake_case_		(    self							)    ->	Optional[Any]:
							UpperCamelCase__								=						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__								=						model.config.in_channels
							UpperCamelCase__								=						16
							UpperCamelCase__								=						torch.randn((1, seq_len, num_features)							).permute(
							    0     ,   2     ,   1							)  # match original, we can update values and remove
							UpperCamelCase__								=						torch.full((num_features,)     ,   0							)

							with torch.no_grad():
									UpperCamelCase__								=						model(__lowerCAmelCase     ,   __lowerCAmelCase							).sample.permute(0     ,   2     ,   1							)

							UpperCamelCase__								=						output[0, -3:, -3:].flatten()
							# fmt: off
							UpperCamelCase__								=						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 						snake_case_		(    self							)    ->	int:
							# Not implemented yet for this UNet
							pass




					@slow
					def 						snake_case_		(    self							)    ->	Tuple:
							UpperCamelCase__								=						"""harmonai/maestro-150k"""
							UpperCamelCase__								=						UNetaDModel.from_pretrained(__lowerCAmelCase     ,   subfolder="""unet"""							)
							model.to(__lowerCAmelCase							)

							UpperCamelCase__								=						65_536
							UpperCamelCase__								=						torch.sin(torch.arange(__lowerCAmelCase							)[None, None, :].repeat(1     ,   2     ,   1							)							).to(__lowerCAmelCase							)
							UpperCamelCase__								=						torch.tensor([1]							).to(__lowerCAmelCase							)

							with torch.no_grad():
									UpperCamelCase__								=						model(__lowerCAmelCase     ,   __lowerCAmelCase							).sample

							UpperCamelCase__								=						output.abs().sum()
							UpperCamelCase__								=						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						( _lowercase			,							_lowercase			,							unittest.TestCase       ):
					__A        = UNetaDModel
					__A        = "sample"




					@property
					def 						snake_case_		(    self							)    ->	Any:
							UpperCamelCase__								=						4
							UpperCamelCase__								=						14
							UpperCamelCase__								=						16

							UpperCamelCase__								=						floats_tensor((batch_size, num_features, seq_len)							).to(__lowerCAmelCase							)
							UpperCamelCase__								=						torch.tensor([10] * batch_size							).to(__lowerCAmelCase							)

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




					@property
					def 						snake_case_		(    self							)    ->	Union[str, Any]:
							return (4, 14, 16)




					@property
					def 						snake_case_		(    self							)    ->	List[Any]:
							return (4, 14, 1)




					@unittest.skipIf(torch_device == """mps"""     ,   """mish op not supported in MPS"""							)
					def 						snake_case_		(    self							)    ->	Dict:
							super().test_determinism()




					@unittest.skipIf(torch_device == """mps"""     ,   """mish op not supported in MPS"""							)
					def 						snake_case_		(    self							)    ->	Optional[int]:
							super().test_outputs_equivalence()




					@unittest.skipIf(torch_device == """mps"""     ,   """mish op not supported in MPS"""							)
					def 						snake_case_		(    self							)    ->	Optional[int]:
							super().test_from_save_pretrained()




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




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




					@unittest.skipIf(torch_device == """mps"""     ,   """mish op not supported in MPS"""							)
					def 						snake_case_		(    self							)    ->	List[str]:
							# UNetRL is a value-function is different output shape
							UpperCamelCase__	,      UpperCamelCase__								=						self.prepare_init_args_and_inputs_for_common()
							UpperCamelCase__								=						self.model_class(**__lowerCAmelCase							)
							model.to(__lowerCAmelCase							)
							model.eval()

							with torch.no_grad():
									UpperCamelCase__								=						model(**__lowerCAmelCase							)

									if isinstance(__lowerCAmelCase     ,   __lowerCAmelCase							):
											UpperCamelCase__								=						output.sample

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




					def 						snake_case_		(    self							)    ->	Optional[int]:
							pass




					def 						snake_case_		(    self							)    ->	Union[str, Any]:
							pass




					def 						snake_case_		(    self							)    ->	Dict:
							UpperCamelCase__								=						{
							    """in_channels""": 14,
							    """out_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,
							    """use_timestep_embedding""": True,
							    """freq_shift""": 1.0,
							    """flip_sin_to_cos""": False,
							    """time_embedding_type""": """positional""",
							    """act_fn""": """mish""",
							}
							UpperCamelCase__								=						self.dummy_input
							return init_dict, inputs_dict




					@unittest.skipIf(torch_device == """mps"""     ,   """mish op not supported in MPS"""							)
					def 						snake_case_		(    self							)    ->	Tuple:
							UpperCamelCase__	,      UpperCamelCase__								=						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							)
							UpperCamelCase__								=						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 						snake_case_		(    self							)    ->	List[Any]:
							UpperCamelCase__	,      UpperCamelCase__								=						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							)

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

							with torch.no_grad():
									UpperCamelCase__								=						value_function(__lowerCAmelCase     ,   __lowerCAmelCase							).sample

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




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





