


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 							a       (      _UpperCAmelCase  ,_UpperCAmelCase  ,unittest.TestCase ):
     UpperCAmelCase__			:	Dict       =						UNetaDModel
     UpperCAmelCase__			:	Optional[Any]       =						"sample"





     @property
     def 			SCREAMING_SNAKE_CASE__							(     self       :					Optional[int]				):
        __lowerCamelCase:					Optional[int]	 =						4
        __lowerCamelCase:					Dict	 =						14
        __lowerCamelCase:					Optional[int]	 =						16

        __lowerCamelCase:					Union[str, Any]	 =						floats_tensor((batch_size, num_features, seq_len)				).to(SCREAMING_SNAKE_CASE_				)
        __lowerCamelCase:					Dict	 =						torch.tensor([10] * batch_size				).to(SCREAMING_SNAKE_CASE_				)

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





     @property
     def 			SCREAMING_SNAKE_CASE__							(     self       :					int				):
        return (4, 14, 16)





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





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





     def 			SCREAMING_SNAKE_CASE__							(     self       :					int				):
        pass





     @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       :					Any				):
        super().test_outputs_equivalence()





     @unittest.skipIf(torch_device == """mps"""     ,						"""mish op not supported in MPS"""				)
     def 			SCREAMING_SNAKE_CASE__							(     self       :					int				):
        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       :					str				):
        super().test_model_from_pretrained()





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





     def 			SCREAMING_SNAKE_CASE__							(     self       :					Any				):
        __lowerCamelCase:					int	 =						{
            """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""",
        }
        __lowerCamelCase:					str	 =						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				):
        __lowerCamelCase ,  __lowerCamelCase:					int	 =						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_				)
        __lowerCamelCase:					Optional[int]	 =						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       :					Dict				):
        __lowerCamelCase:					Union[str, Any]	 =						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				)

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

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

        __lowerCamelCase:					Optional[int]	 =						output[0, -3:, -3:].flatten()
        # fmt: off
        __lowerCamelCase:					Tuple	 =						torch.tensor([-2.13_7172, 1.142_6016, 0.368_8687, -0.76_6922, 0.730_3146, 0.1103_8864, -0.476_0633, 0.1327_0172, 0.0259_1348]				)
        # fmt: on
        self.assertTrue(torch.allclose(SCREAMING_SNAKE_CASE_     ,						SCREAMING_SNAKE_CASE_     ,						rtol=1E-3				)				)





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







     @slow
     def 			SCREAMING_SNAKE_CASE__							(     self       :					List[str]				):
        __lowerCamelCase:					int	 =						"""harmonai/maestro-150k"""
        __lowerCamelCase:					List[Any]	 =						UNetaDModel.from_pretrained(SCREAMING_SNAKE_CASE_     ,						subfolder="""unet"""				)
        model.to(SCREAMING_SNAKE_CASE_				)

        __lowerCamelCase:					List[str]	 =						65536
        __lowerCamelCase:					Dict	 =						torch.sin(torch.arange(SCREAMING_SNAKE_CASE_				)[None, None, :].repeat(1     ,						2     ,						1				)				).to(SCREAMING_SNAKE_CASE_				)
        __lowerCamelCase:					str	 =						torch.tensor([1]				).to(SCREAMING_SNAKE_CASE_				)

        with torch.no_grad():
           __lowerCamelCase:					Union[str, Any]	 =						model(SCREAMING_SNAKE_CASE_     ,						SCREAMING_SNAKE_CASE_				).sample

        __lowerCamelCase:					Union[str, Any]	 =						output.abs().sum()
        __lowerCamelCase:					Tuple	 =						output.abs().max()

        assert (output_sum - 224.0896).abs() < 0.5
        assert (output_max - 0.0607).abs() < 4E-4



class 							a       (      _UpperCAmelCase  ,_UpperCAmelCase  ,unittest.TestCase ):
     UpperCAmelCase__			:	List[str]       =						UNetaDModel
     UpperCAmelCase__			:	List[Any]       =						"sample"





     @property
     def 			SCREAMING_SNAKE_CASE__							(     self       :					List[Any]				):
        __lowerCamelCase:					List[str]	 =						4
        __lowerCamelCase:					str	 =						14
        __lowerCamelCase:					Dict	 =						16

        __lowerCamelCase:					Union[str, Any]	 =						floats_tensor((batch_size, num_features, seq_len)				).to(SCREAMING_SNAKE_CASE_				)
        __lowerCamelCase:					Union[str, Any]	 =						torch.tensor([10] * batch_size				).to(SCREAMING_SNAKE_CASE_				)

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





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





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





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





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

        with torch.no_grad():
           __lowerCamelCase:					Any	 =						model(**SCREAMING_SNAKE_CASE_				)

           if isinstance(SCREAMING_SNAKE_CASE_     ,						SCREAMING_SNAKE_CASE_				):
              __lowerCamelCase:					Any	 =						output.sample

        self.assertIsNotNone(SCREAMING_SNAKE_CASE_				)
        __lowerCamelCase:					Tuple	 =						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       :					Any				):
        pass





     def 			SCREAMING_SNAKE_CASE__							(     self       :					Optional[Any]				):
        __lowerCamelCase:					List[str]	 =						{
            """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""",
        }
        __lowerCamelCase:					Union[str, Any]	 =						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				):
        __lowerCamelCase ,  __lowerCamelCase:					Tuple	 =						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_				)
        __lowerCamelCase:					List[Any]	 =						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]				):
        __lowerCamelCase ,  __lowerCamelCase:					Dict	 =						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				)

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

        with torch.no_grad():
           __lowerCamelCase:					int	 =						value_function(SCREAMING_SNAKE_CASE_     ,						SCREAMING_SNAKE_CASE_				).sample

        # fmt: off
        __lowerCamelCase:					Union[str, Any]	 =						torch.tensor([165.25] * seq_len				)
        # fmt: on
        self.assertTrue(torch.allclose(SCREAMING_SNAKE_CASE_     ,						SCREAMING_SNAKE_CASE_     ,						rtol=1E-3				)				)







     def 			SCREAMING_SNAKE_CASE__							(     self       :					int				):
        # Not implemented yet for this UNet
        pass






import unittest

import numpy as np

from transformers.testing_utils import require_torch, require_vision
from transformers.utils import is_torch_available, is_vision_available

from ...test_image_processing_common import ImageProcessingSavingTestMixin, prepare_image_inputs


if is_torch_available():
   import torch

if is_vision_available():
   from PIL import Image

   from transformers import MobileNetVaImageProcessor



class 							a       (      unittest.TestCase ):





     def __init__(     self       :					Optional[int]     ,						SCREAMING_SNAKE_CASE_       :					str     ,						SCREAMING_SNAKE_CASE_       :					List[str]=7     ,						SCREAMING_SNAKE_CASE_       :					Union[str, Any]=3     ,						SCREAMING_SNAKE_CASE_       :					List[Any]=18     ,						SCREAMING_SNAKE_CASE_       :					List[str]=30     ,						SCREAMING_SNAKE_CASE_       :					List[Any]=400     ,						SCREAMING_SNAKE_CASE_       :					Dict=True     ,						SCREAMING_SNAKE_CASE_       :					Any=None     ,						SCREAMING_SNAKE_CASE_       :					List[str]=True     ,						SCREAMING_SNAKE_CASE_       :					Any=None     ,						):
        __lowerCamelCase:					List[Any]	 =						size if size is not None else {"""shortest_edge""": 20}
        __lowerCamelCase:					Tuple	 =						crop_size if crop_size is not None else {"""height""": 18, """width""": 18}
        __lowerCamelCase:					List[Any]	 =						parent
        __lowerCamelCase:					Optional[int]	 =						batch_size
        __lowerCamelCase:					str	 =						num_channels
        __lowerCamelCase:					Dict	 =						image_size
        __lowerCamelCase:					Tuple	 =						min_resolution
        __lowerCamelCase:					Union[str, Any]	 =						max_resolution
        __lowerCamelCase:					Tuple	 =						do_resize
        __lowerCamelCase:					List[str]	 =						size
        __lowerCamelCase:					int	 =						do_center_crop
        __lowerCamelCase:					Tuple	 =						crop_size







     def 			SCREAMING_SNAKE_CASE__							(     self       :					List[Any]				):
        return {
            "do_resize": self.do_resize,
            "size": self.size,
            "do_center_crop": self.do_center_crop,
            "crop_size": self.crop_size,
        }



@require_torch
@require_vision
class 							a       (      _UpperCAmelCase  ,unittest.TestCase ):
     UpperCAmelCase__			:	int       =						MobileNetVaImageProcessor if is_vision_available() else None





     def 			SCREAMING_SNAKE_CASE__							(     self       :					List[str]				):
        __lowerCamelCase:					Tuple	 =						MobileNetVaImageProcessingTester(self				)





     @property
     def 			SCREAMING_SNAKE_CASE__							(     self       :					str				):
        return self.image_processor_tester.prepare_image_processor_dict()





     def 			SCREAMING_SNAKE_CASE__							(     self       :					List[str]				):
        __lowerCamelCase:					List[Any]	 =						self.image_processing_class(**self.image_processor_dict				)
        self.assertTrue(hasattr(SCREAMING_SNAKE_CASE_     ,						"""do_resize"""				)				)
        self.assertTrue(hasattr(SCREAMING_SNAKE_CASE_     ,						"""size"""				)				)
        self.assertTrue(hasattr(SCREAMING_SNAKE_CASE_     ,						"""do_center_crop"""				)				)
        self.assertTrue(hasattr(SCREAMING_SNAKE_CASE_     ,						"""center_crop"""				)				)





     def 			SCREAMING_SNAKE_CASE__							(     self       :					Any				):
        __lowerCamelCase:					str	 =						self.image_processing_class.from_dict(self.image_processor_dict				)
        self.assertEqual(image_processor.size     ,						{"""shortest_edge""": 20}				)
        self.assertEqual(image_processor.crop_size     ,						{"""height""": 18, """width""": 18}				)

        __lowerCamelCase:					Optional[int]	 =						self.image_processing_class.from_dict(self.image_processor_dict     ,						size=42     ,						crop_size=84				)
        self.assertEqual(image_processor.size     ,						{"""shortest_edge""": 42}				)
        self.assertEqual(image_processor.crop_size     ,						{"""height""": 84, """width""": 84}				)





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





     def 			SCREAMING_SNAKE_CASE__							(     self       :					Optional[int]				):
        # Initialize image_processing
        __lowerCamelCase:					List[str]	 =						self.image_processing_class(**self.image_processor_dict				)
        # create random PIL images
        __lowerCamelCase:					str	 =						prepare_image_inputs(self.image_processor_tester     ,						equal_resolution=SCREAMING_SNAKE_CASE_				)
        for image in image_inputs:
           self.assertIsInstance(SCREAMING_SNAKE_CASE_     ,						Image.Image				)

        # Test not batched input
        __lowerCamelCase:					Optional[Any]	 =						image_processing(image_inputs[0]     ,						return_tensors="""pt"""				).pixel_values
        self.assertEqual(
            encoded_images.shape     ,						(
                1,
                self.image_processor_tester.num_channels,
                self.image_processor_tester.crop_size["""height"""],
                self.image_processor_tester.crop_size["""width"""],
            )     ,						)

        # Test batched
        __lowerCamelCase:					Tuple	 =						image_processing(SCREAMING_SNAKE_CASE_     ,						return_tensors="""pt"""				).pixel_values
        self.assertEqual(
            encoded_images.shape     ,						(
                self.image_processor_tester.batch_size,
                self.image_processor_tester.num_channels,
                self.image_processor_tester.crop_size["""height"""],
                self.image_processor_tester.crop_size["""width"""],
            )     ,						)





     def 			SCREAMING_SNAKE_CASE__							(     self       :					Optional[int]				):
        # Initialize image_processing
        __lowerCamelCase:					Union[str, Any]	 =						self.image_processing_class(**self.image_processor_dict				)
        # create random numpy tensors
        __lowerCamelCase:					Tuple	 =						prepare_image_inputs(self.image_processor_tester     ,						equal_resolution=SCREAMING_SNAKE_CASE_     ,						numpify=SCREAMING_SNAKE_CASE_				)
        for image in image_inputs:
           self.assertIsInstance(SCREAMING_SNAKE_CASE_     ,						np.ndarray				)

        # Test not batched input
        __lowerCamelCase:					List[Any]	 =						image_processing(image_inputs[0]     ,						return_tensors="""pt"""				).pixel_values
        self.assertEqual(
            encoded_images.shape     ,						(
                1,
                self.image_processor_tester.num_channels,
                self.image_processor_tester.crop_size["""height"""],
                self.image_processor_tester.crop_size["""width"""],
            )     ,						)

        # Test batched
        __lowerCamelCase:					Optional[Any]	 =						image_processing(SCREAMING_SNAKE_CASE_     ,						return_tensors="""pt"""				).pixel_values
        self.assertEqual(
            encoded_images.shape     ,						(
                self.image_processor_tester.batch_size,
                self.image_processor_tester.num_channels,
                self.image_processor_tester.crop_size["""height"""],
                self.image_processor_tester.crop_size["""width"""],
            )     ,						)







     def 			SCREAMING_SNAKE_CASE__							(     self       :					Optional[int]				):
        # Initialize image_processing
        __lowerCamelCase:					str	 =						self.image_processing_class(**self.image_processor_dict				)
        # create random PyTorch tensors
        __lowerCamelCase:					List[str]	 =						prepare_image_inputs(self.image_processor_tester     ,						equal_resolution=SCREAMING_SNAKE_CASE_     ,						torchify=SCREAMING_SNAKE_CASE_				)
        for image in image_inputs:
           self.assertIsInstance(SCREAMING_SNAKE_CASE_     ,						torch.Tensor				)

        # Test not batched input
        __lowerCamelCase:					Union[str, Any]	 =						image_processing(image_inputs[0]     ,						return_tensors="""pt"""				).pixel_values
        self.assertEqual(
            encoded_images.shape     ,						(
                1,
                self.image_processor_tester.num_channels,
                self.image_processor_tester.crop_size["""height"""],
                self.image_processor_tester.crop_size["""width"""],
            )     ,						)

        # Test batched
        __lowerCamelCase:					Optional[int]	 =						image_processing(SCREAMING_SNAKE_CASE_     ,						return_tensors="""pt"""				).pixel_values
        self.assertEqual(
            encoded_images.shape     ,						(
                self.image_processor_tester.batch_size,
                self.image_processor_tester.num_channels,
                self.image_processor_tester.crop_size["""height"""],
                self.image_processor_tester.crop_size["""width"""],
            )     ,						)


