
"""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 						lowercase(snake_case__						,	snake_case__						,	unittest.TestCase					):
     __snake_case:  Tuple  				=    UNetaDModel
     __snake_case:  Dict  				=    'sample'



     @property
     def  lowercase__     (       self      )		->    Any:







         """simple docstring"""



         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  lowercase__     (       self      )		->    List[str]:







         """simple docstring"""



         return (4, 1_4, 1_6)



     @property
     def  lowercase__     (       self      )		->    Optional[Any]:







         """simple docstring"""



         return (4, 1_4, 1_6)



     def  lowercase__     (       self      )		->    Optional[Any]:







         """simple docstring"""



         pass



     def  lowercase__     (       self      )		->    List[str]:







         """simple docstring"""



         pass



     @unittest.skipIf(torch_device == 'mps'    ,							'mish op not supported in MPS'      )
     def  lowercase__     (       self      )		->    List[Any]:







         """simple docstring"""



         super().test_determinism()



     @unittest.skipIf(torch_device == 'mps'    ,							'mish op not supported in MPS'      )
     def  lowercase__     (       self      )		->    int:







         """simple docstring"""



         super().test_outputs_equivalence()



     @unittest.skipIf(torch_device == 'mps'    ,							'mish op not supported in MPS'      )
     def  lowercase__     (       self      )		->    Any:







         """simple docstring"""



         super().test_from_save_pretrained()



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







         """simple docstring"""



         super().test_from_save_pretrained_variant()



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







         """simple docstring"""



         super().test_model_from_pretrained()



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







         """simple docstring"""



         super().test_output()



     def  lowercase__     (       self      )		->    List[Any]:







         """simple docstring"""



         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  lowercase__     (       self      )		->    Tuple:







         """simple docstring"""



         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  lowercase__     (       self      )		->    str:







         """simple docstring"""



         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.13_71_72, 1.1_42_60_16, 0.3_68_86_87, -0.76_69_22, 0.7_30_31_46, 0.11_03_88_64, -0.4_76_06_33, 0.13_27_01_72, 0.02_59_13_48]      )
         # fmt: on
         self.assertTrue(torch.allclose(__SCREAMING_SNAKE_CASE    ,							__SCREAMING_SNAKE_CASE    ,							rtol=1e-3      )      )



     def  lowercase__     (       self      )		->    Any:







         """simple docstring"""



         pass




     @slow
     def  lowercase__     (       self      )		->    int:







         """simple docstring"""



         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_24.08_96).abs() < 0.5
         assert (output_max - 0.06_07).abs() < 4e-4






class 						lowercase(snake_case__						,	snake_case__						,	unittest.TestCase					):
     __snake_case:  Tuple  				=    UNetaDModel
     __snake_case:  Dict  				=    'sample'



     @property
     def  lowercase__     (       self      )		->    List[str]:







         """simple docstring"""



         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  lowercase__     (       self      )		->    List[Any]:







         """simple docstring"""



         return (4, 1_4, 1_6)



     @property
     def  lowercase__     (       self      )		->    str:







         """simple docstring"""



         return (4, 1_4, 1)



     @unittest.skipIf(torch_device == 'mps'    ,							'mish op not supported in MPS'      )
     def  lowercase__     (       self      )		->    int:







         """simple docstring"""



         super().test_determinism()



     @unittest.skipIf(torch_device == 'mps'    ,							'mish op not supported in MPS'      )
     def  lowercase__     (       self      )		->    List[Any]:







         """simple docstring"""



         super().test_outputs_equivalence()



     @unittest.skipIf(torch_device == 'mps'    ,							'mish op not supported in MPS'      )
     def  lowercase__     (       self      )		->    Any:







         """simple docstring"""



         super().test_from_save_pretrained()



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







         """simple docstring"""



         super().test_from_save_pretrained_variant()



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







         """simple docstring"""



         super().test_model_from_pretrained()



     @unittest.skipIf(torch_device == 'mps'    ,							'mish op not supported in MPS'      )
     def  lowercase__     (       self      )		->    List[Any]:







         """simple docstring"""



         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  lowercase__     (       self      )		->    Optional[int]:







         """simple docstring"""



         pass



     def  lowercase__     (       self      )		->    str:







         """simple docstring"""



         pass



     def  lowercase__     (       self      )		->    Tuple:







         """simple docstring"""



         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  lowercase__     (       self      )		->    List[Any]:







         """simple docstring"""



         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  lowercase__     (       self      )		->    Dict:







         """simple docstring"""



         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_65.25] * seq_len      )
         # fmt: on
         self.assertTrue(torch.allclose(__SCREAMING_SNAKE_CASE    ,							__SCREAMING_SNAKE_CASE    ,							rtol=1e-3      )      )




     def  lowercase__     (       self      )		->    Dict:







         """simple docstring"""



         pass




"""simple docstring"""




import gc
import tempfile
import unittest

import numpy as np
import torch

from diffusers import VersatileDiffusionDualGuidedPipeline
from diffusers.utils.testing_utils import load_image, nightly, require_torch_gpu, torch_device


a		:					str    	=					False







@nightly
@require_torch_gpu
class 						lowercase(unittest.TestCase					):



     def  lowercase__     (       self      )		->    Any:







         """simple docstring"""



         super().tearDown()
         gc.collect()
         torch.cuda.empty_cache()



     def  lowercase__     (       self      )		->    Optional[int]:







         """simple docstring"""



         a__    = VersatileDiffusionDualGuidedPipeline.from_pretrained('shi-labs/versatile-diffusion'      )
         # remove text_unet
         pipe.remove_unused_weights()
         pipe.to(__SCREAMING_SNAKE_CASE      )
         pipe.set_progress_bar_config(disable=__SCREAMING_SNAKE_CASE      )

         a__    = load_image(
             'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/versatile_diffusion/benz.jpg'      )

         a__    = torch.manual_seed(0      )
         a__    = pipe(
             prompt='first prompt'    ,							image=__SCREAMING_SNAKE_CASE    ,							text_to_image_strength=0.75    ,							generator=__SCREAMING_SNAKE_CASE    ,							guidance_scale=7.5    ,							num_inference_steps=2    ,							output_type='numpy'    ,							).images

         with tempfile.TemporaryDirectory() as tmpdirname:
             pipe.save_pretrained(__SCREAMING_SNAKE_CASE      )
             a__    = VersatileDiffusionDualGuidedPipeline.from_pretrained(__SCREAMING_SNAKE_CASE      )

         pipe.to(__SCREAMING_SNAKE_CASE      )
         pipe.set_progress_bar_config(disable=__SCREAMING_SNAKE_CASE      )

         a__    = generator.manual_seed(0      )
         a__    = pipe(
             prompt='first prompt'    ,							image=__SCREAMING_SNAKE_CASE    ,							text_to_image_strength=0.75    ,							generator=__SCREAMING_SNAKE_CASE    ,							guidance_scale=7.5    ,							num_inference_steps=2    ,							output_type='numpy'    ,							).images

         assert np.abs(image - new_image      ).sum() < 1e-5, "Models don't have the same forward pass"




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







         """simple docstring"""



         a__    = VersatileDiffusionDualGuidedPipeline.from_pretrained('shi-labs/versatile-diffusion'      )
         pipe.remove_unused_weights()
         pipe.to(__SCREAMING_SNAKE_CASE      )
         pipe.set_progress_bar_config(disable=__SCREAMING_SNAKE_CASE      )

         a__    = 'cyberpunk 2077'
         a__    = load_image(
             'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/versatile_diffusion/benz.jpg'      )
         a__    = torch.manual_seed(0      )
         a__    = pipe(
             prompt=__SCREAMING_SNAKE_CASE    ,							image=__SCREAMING_SNAKE_CASE    ,							text_to_image_strength=0.75    ,							generator=__SCREAMING_SNAKE_CASE    ,							guidance_scale=7.5    ,							num_inference_steps=5_0    ,							output_type='numpy'    ,							).images

         a__    = image[0, 2_5_3:2_5_6, 2_5_3:2_5_6, -1]

         assert image.shape == (1, 5_1_2, 5_1_2, 3)
         a__    = np.array([0.07_87, 0.08_49, 0.08_26, 0.08_12, 0.08_07, 0.07_95, 0.08_18, 0.07_98, 0.07_79]      )

         assert np.abs(image_slice.flatten() - expected_slice      ).max() < 1e-2


