






"""simple docstring"""





from typing import SupportsIndex

import numpy as np
from scipy.ndimage import convolve






def 							A(    snake_case_		,  snake_case_		,  snake_case_						):
       """simple docstring"""





       lowercase__:							List[Any]         =     np.stack((horizontal_flow, vertical_flow)		,  2						)

       # Create a grid of all pixel coordinates and subtract the flow to get the
       # target pixels coordinates
       lowercase__:							Union[str, Any]         =     np.stack(
           np.meshgrid(np.arange(0		,  image.shape[1]						)		,  np.arange(0		,  image.shape[0]						)						)		,  2						)
       lowercase__:							Dict         =     np.round(grid - flow						).astype(np.intaa						)

       # Find the locations outside of the original image
       lowercase__:							Optional[int]         =     (grid < 0) | (grid >= np.array([image.shape[1], image.shape[0]]						))
       lowercase__:							List[Any]         =     0

       lowercase__:							int         =     image[grid[:, :, 1], grid[:, :, 0]]

       # Set pixels at invalid locations to 0
       lowercase__:							Tuple         =     0

       return warped






def 							A(    snake_case_		,  snake_case_		,  snake_case_		,  snake_case_ = None		,  ):
       """simple docstring"""





       if alpha is None:
              lowercase__:							Optional[int]         =     0.1

       # Initialize flow
       lowercase__:							Union[str, Any]         =     np.zeros_like(snake_case_						)
       lowercase__:							List[Any]         =     np.zeros_like(snake_case_						)

       # Prepare kernels for the calculation of the derivatives and the average velocity
       lowercase__:							int         =     np.array([[-1, 1], [-1, 1]]						) * 0.25
       lowercase__:							Tuple         =     np.array([[-1, -1], [1, 1]]						) * 0.25
       lowercase__:							Dict         =     np.array([[1, 1], [1, 1]]						) * 0.25
       lowercase__:							List[str]         =     np.array(
           [[1 / 12, 1 / 6, 1 / 12], [1 / 6, 0, 1 / 6], [1 / 12, 1 / 6, 1 / 12]]						)

       # Iteratively refine the flow
       for _ in range(snake_case_						):
              lowercase__:							Any         =     warp(snake_case_		,  snake_case_		,  snake_case_						)
              lowercase__:							Dict         =     convolve(snake_case_		,  snake_case_						) + convolve(snake_case_		,  snake_case_						)
              lowercase__:							Any         =     convolve(snake_case_		,  snake_case_						) + convolve(snake_case_		,  snake_case_						)
              lowercase__:							Dict         =     convolve(snake_case_		,  snake_case_						) + convolve(snake_case_		,  -kernel_t						)

              lowercase__:							str         =     convolve(snake_case_		,  snake_case_						)
              lowercase__:							List[Any]         =     convolve(snake_case_		,  snake_case_						)

              # This updates the flow as proposed in the paper (Step 12)
              lowercase__:							Optional[Any]         =     (
                  derivative_x * avg_horizontal_velocity
                  + derivative_y * avg_vertical_velocity
                  + derivative_t
              )
              lowercase__:							Any         =     update / (alpha**2 + derivative_x**2 + derivative_y**2)

              lowercase__:							Optional[int]         =     avg_horizontal_velocity - derivative_x * update
              lowercase__:							List[str]         =     avg_vertical_velocity - derivative_y * update

       return horizontal_flow, vertical_flow


if __name__ == "__main__":
  import doctest

  doctest.testmod()








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 		__lowerCamelCase     (   lowercase							, lowercase							, unittest.TestCase		):
							lowerCamelCase__:				int      =							UNetaDModel
							lowerCamelCase__:				Any      =							'''sample'''



							@property
							def     A__ (     self							)						->     Tuple:






								"""simple docstring"""



								UpperCAmelCase:					Union[str, Any] 						=     4
								UpperCAmelCase:					Optional[Any] 						=     1_4
								UpperCAmelCase:					str 						=     1_6

								UpperCAmelCase:					Optional[int] 						=     floats_tensor((batch_size, num_features, seq_len)							).to(__snake_case							)
								UpperCAmelCase:					Any 						=     torch.tensor([1_0] * batch_size							).to(__snake_case							)

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



							@property
							def     A__ (     self							)						->     Tuple:






								"""simple docstring"""



								return (4, 1_4, 1_6)



							@property
							def     A__ (     self							)						->     Union[str, Any]:






								"""simple docstring"""



								return (4, 1_4, 1_6)



							def     A__ (     self							)						->     int:






								"""simple docstring"""



								pass



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






								"""simple docstring"""



								pass



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






								"""simple docstring"""



								super().test_determinism()



							@unittest.skipIf(torch_device == "mps"	,				"mish op not supported in MPS"							)
							def     A__ (     self							)						->     Tuple:






								"""simple docstring"""



								super().test_outputs_equivalence()



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






								"""simple docstring"""



								super().test_from_save_pretrained()



							@unittest.skipIf(torch_device == "mps"	,				"mish op not supported in MPS"							)
							def     A__ (     self							)						->     str:






								"""simple docstring"""



								super().test_from_save_pretrained_variant()



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






								"""simple docstring"""



								super().test_model_from_pretrained()



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






								"""simple docstring"""



								super().test_output()



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






								"""simple docstring"""



								UpperCAmelCase:					Tuple 						=     {
								    "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",
								}
								UpperCAmelCase:					Tuple 						=     self.dummy_input
								return init_dict, inputs_dict



							@unittest.skipIf(torch_device == "mps"	,				"mish op not supported in MPS"							)
							def     A__ (     self							)						->     Tuple:






								"""simple docstring"""



								UpperCAmelCase						,   UpperCAmelCase:					Any 						=     UNetaDModel.from_pretrained(
								    "bglick13/hopper-medium-v2-value-function-hor32"	,				output_loading_info=__snake_case	,				subfolder="unet"							)
								self.assertIsNotNone(__snake_case							)
								self.assertEqual(len(loading_info["missing_keys"]							)	,				0							)

								model.to(__snake_case							)
								UpperCAmelCase:					Tuple 						=     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     A__ (     self							)						->     List[str]:






								"""simple docstring"""



								UpperCAmelCase:					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							)

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

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

								UpperCAmelCase:					str 						=     output[0, -3:, -3:].flatten()
								# fmt: off
								UpperCAmelCase:					Dict 						=     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(__snake_case	,				__snake_case	,				rtol=1e-3							)							)



							def     A__ (     self							)						->     Any:






								"""simple docstring"""



								pass




							@slow
							def     A__ (     self							)						->     Tuple:






								"""simple docstring"""



								UpperCAmelCase:					List[str] 						=     "harmonai/maestro-150k"
								UpperCAmelCase:					Optional[Any] 						=     UNetaDModel.from_pretrained(__snake_case	,				subfolder="unet"							)
								model.to(__snake_case							)

								UpperCAmelCase:					Any 						=     6_5_5_3_6
								UpperCAmelCase:					str 						=     torch.sin(torch.arange(__snake_case							)[None, None, :].repeat(1	,				2	,				1							)							).to(__snake_case							)
								UpperCAmelCase:					Dict 						=     torch.tensor([1]							).to(__snake_case							)

								with torch.no_grad():
									UpperCAmelCase:					Optional[Any] 						=     model(__snake_case	,				__snake_case							).sample

								UpperCAmelCase:					List[Any] 						=     output.abs().sum()
								UpperCAmelCase:					Dict 						=     output.abs().max()

								assert (output_sum - 2_24.08_96).abs() < 0.5
								assert (output_max - 0.06_07).abs() < 4e-4





class 		__lowerCamelCase     (   lowercase							, lowercase							, unittest.TestCase		):
							lowerCamelCase__:				Dict      =							UNetaDModel
							lowerCamelCase__:				int      =							'''sample'''



							@property
							def     A__ (     self							)						->     int:






								"""simple docstring"""



								UpperCAmelCase:					List[Any] 						=     4
								UpperCAmelCase:					int 						=     1_4
								UpperCAmelCase:					Optional[int] 						=     1_6

								UpperCAmelCase:					List[str] 						=     floats_tensor((batch_size, num_features, seq_len)							).to(__snake_case							)
								UpperCAmelCase:					Any 						=     torch.tensor([1_0] * batch_size							).to(__snake_case							)

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



							@property
							def     A__ (     self							)						->     List[Any]:






								"""simple docstring"""



								return (4, 1_4, 1_6)



							@property
							def     A__ (     self							)						->     List[Any]:






								"""simple docstring"""



								return (4, 1_4, 1)



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






								"""simple docstring"""



								super().test_determinism()



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






								"""simple docstring"""



								super().test_outputs_equivalence()



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






								"""simple docstring"""



								super().test_from_save_pretrained()



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






								"""simple docstring"""



								super().test_from_save_pretrained_variant()



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






								"""simple docstring"""



								super().test_model_from_pretrained()



							@unittest.skipIf(torch_device == "mps"	,				"mish op not supported in MPS"							)
							def     A__ (     self							)						->     str:






								"""simple docstring"""



								UpperCAmelCase						,   UpperCAmelCase:					Optional[int] 						=     self.prepare_init_args_and_inputs_for_common()
								UpperCAmelCase:					Tuple 						=     self.model_class(**__snake_case							)
								model.to(__snake_case							)
								model.eval()

								with torch.no_grad():
									UpperCAmelCase:					Union[str, Any] 						=     model(**__snake_case							)

									if isinstance(__snake_case	,				__snake_case							):
										UpperCAmelCase:					List[str] 						=     output.sample

								self.assertIsNotNone(__snake_case							)
								UpperCAmelCase:					Optional[Any] 						=     torch.Size((inputs_dict["sample"].shape[0], 1)							)
								self.assertEqual(output.shape	,				__snake_case	,				"Input and output shapes do not match"							)



							def     A__ (     self							)						->     Tuple:






								"""simple docstring"""



								pass



							def     A__ (     self							)						->     Any:






								"""simple docstring"""



								pass



							def     A__ (     self							)						->     int:






								"""simple docstring"""



								UpperCAmelCase:					Tuple 						=     {
								    "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",
								}
								UpperCAmelCase:					Union[str, Any] 						=     self.dummy_input
								return init_dict, inputs_dict



							@unittest.skipIf(torch_device == "mps"	,				"mish op not supported in MPS"							)
							def     A__ (     self							)						->     Tuple:






								"""simple docstring"""



								UpperCAmelCase						,   UpperCAmelCase:					Optional[Any] 						=     UNetaDModel.from_pretrained(
								    "bglick13/hopper-medium-v2-value-function-hor32"	,				output_loading_info=__snake_case	,				subfolder="value_function"							)
								self.assertIsNotNone(__snake_case							)
								self.assertEqual(len(vf_loading_info["missing_keys"]							)	,				0							)

								value_function.to(__snake_case							)
								UpperCAmelCase:					Union[str, 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     A__ (     self							)						->     Any:






								"""simple docstring"""



								UpperCAmelCase						,   UpperCAmelCase:					Dict 						=     UNetaDModel.from_pretrained(
								    "bglick13/hopper-medium-v2-value-function-hor32"	,				output_loading_info=__snake_case	,				subfolder="value_function"							)
								torch.manual_seed(0							)
								if torch.cuda.is_available():
									torch.cuda.manual_seed_all(0							)

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

								with torch.no_grad():
									UpperCAmelCase:					Dict 						=     value_function(__snake_case	,				__snake_case							).sample

								# fmt: off
								UpperCAmelCase:					Tuple 						=     torch.tensor([1_65.25] * seq_len							)
								# fmt: on
								self.assertTrue(torch.allclose(__snake_case	,				__snake_case	,				rtol=1e-3							)							)




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






								"""simple docstring"""



								pass
