

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_      (	_snake_case		,  _snake_case		,  unittest.TestCase      ):


  """simple docstring"""
  SCREAMING_SNAKE_CASE__					:			Union[str, Any]					    =   UNetaDModel
  SCREAMING_SNAKE_CASE__					:			List[str]					    =   """sample"""



  @property
  def   _UpperCamelCase     (						self     :					str						)       ->			Optional[Any]:
      SCREAMING_SNAKE_CASE__         =    4
      SCREAMING_SNAKE_CASE__         =    14
      SCREAMING_SNAKE_CASE__         =    16

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

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



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



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



  def   _UpperCamelCase     (						self     :					Optional[Any]						)       ->			Tuple:
      pass



  def   _UpperCamelCase     (						self     :					Optional[Any]						)       ->			Any:
      pass



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



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



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



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



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



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



  def   _UpperCamelCase     (						self     :					Dict						)       ->			List[Any]:
      SCREAMING_SNAKE_CASE__         =    {
          "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",
      }
      SCREAMING_SNAKE_CASE__         =    self.dummy_input
      return init_dict, inputs_dict



  @unittest.skipIf(torch_device == "mps"  ,	"mish op not supported in MPS"						)
  def   _UpperCamelCase     (						self     :					List[str]						)       ->			Tuple:
      SCREAMING_SNAKE_CASE__  ,		SCREAMING_SNAKE_CASE__         =    UNetaDModel.from_pretrained(
          "bglick13/hopper-medium-v2-value-function-hor32"  ,	output_loading_info=A_  ,	subfolder="unet"						)
      self.assertIsNotNone(A_						)
      self.assertEqual(len(loading_info["missing_keys"]						)  ,	0						)

      model.to(A_						)
      SCREAMING_SNAKE_CASE__         =    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   _UpperCamelCase     (						self     :					Optional[Any]						)       ->			Optional[int]:
      SCREAMING_SNAKE_CASE__         =    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						)

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

      with torch.no_grad():
          SCREAMING_SNAKE_CASE__         =    model(A_  ,	A_						).sample.permute(0  ,	2  ,	1						)

      SCREAMING_SNAKE_CASE__         =    output[0, -3:, -3:].flatten()
      # fmt: off
      SCREAMING_SNAKE_CASE__         =    torch.tensor([-2.137_172, 1.1_426_016, 0.3_688_687, -0.766_922, 0.7_303_146, 0.11_038_864, -0.4_760_633, 0.13_270_172, 0.02_591_348]						)
      # fmt: on
      self.assertTrue(torch.allclose(A_  ,	A_  ,	rtol=1E-3						)						)



  def   _UpperCamelCase     (						self     :					Optional[Any]						)       ->			int:
      # Not implemented yet for this UNet
      pass







  @slow
  def   _UpperCamelCase     (						self     :					Tuple						)       ->			int:
      SCREAMING_SNAKE_CASE__         =    "harmonai/maestro-150k"
      SCREAMING_SNAKE_CASE__         =    UNetaDModel.from_pretrained(A_  ,	subfolder="unet"						)
      model.to(A_						)

      SCREAMING_SNAKE_CASE__         =    6_5536
      SCREAMING_SNAKE_CASE__         =    torch.sin(torch.arange(A_						)[None, None, :].repeat(1  ,	2  ,	1						)						).to(A_						)
      SCREAMING_SNAKE_CASE__         =    torch.tensor([1]						).to(A_						)

      with torch.no_grad():
          SCREAMING_SNAKE_CASE__         =    model(A_  ,	A_						).sample

      SCREAMING_SNAKE_CASE__         =    output.abs().sum()
      SCREAMING_SNAKE_CASE__         =    output.abs().max()

      assert (output_sum - 224.0_896).abs() < 0.5
      assert (output_max - 0.0_607).abs() < 4E-4







class  UpperCamelCase_      (	_snake_case		,  _snake_case		,  unittest.TestCase      ):


  """simple docstring"""
  SCREAMING_SNAKE_CASE__					:			List[Any]					    =   UNetaDModel
  SCREAMING_SNAKE_CASE__					:			Optional[Any]					    =   """sample"""



  @property
  def   _UpperCamelCase     (						self     :					Union[str, Any]						)       ->			List[Any]:
      SCREAMING_SNAKE_CASE__         =    4
      SCREAMING_SNAKE_CASE__         =    14
      SCREAMING_SNAKE_CASE__         =    16

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

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



  @property
  def   _UpperCamelCase     (						self     :					List[str]						)       ->			List[str]:
      return (4, 14, 16)



  @property
  def   _UpperCamelCase     (						self     :					Tuple						)       ->			Optional[Any]:
      return (4, 14, 1)



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



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



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



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



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



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

      with torch.no_grad():
          SCREAMING_SNAKE_CASE__         =    model(**A_						)

          if isinstance(A_  ,	A_						):
              SCREAMING_SNAKE_CASE__         =    output.sample

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



  def   _UpperCamelCase     (						self     :					Any						)       ->			Tuple:
      pass



  def   _UpperCamelCase     (						self     :					Optional[Any]						)       ->			Tuple:
      pass



  def   _UpperCamelCase     (						self     :					str						)       ->			List[Any]:
      SCREAMING_SNAKE_CASE__         =    {
          "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",
      }
      SCREAMING_SNAKE_CASE__         =    self.dummy_input
      return init_dict, inputs_dict



  @unittest.skipIf(torch_device == "mps"  ,	"mish op not supported in MPS"						)
  def   _UpperCamelCase     (						self     :					Optional[int]						)       ->			Dict:
      SCREAMING_SNAKE_CASE__  ,		SCREAMING_SNAKE_CASE__         =    UNetaDModel.from_pretrained(
          "bglick13/hopper-medium-v2-value-function-hor32"  ,	output_loading_info=A_  ,	subfolder="value_function"						)
      self.assertIsNotNone(A_						)
      self.assertEqual(len(vf_loading_info["missing_keys"]						)  ,	0						)

      value_function.to(A_						)
      SCREAMING_SNAKE_CASE__         =    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   _UpperCamelCase     (						self     :					int						)       ->			Tuple:
      SCREAMING_SNAKE_CASE__  ,		SCREAMING_SNAKE_CASE__         =    UNetaDModel.from_pretrained(
          "bglick13/hopper-medium-v2-value-function-hor32"  ,	output_loading_info=A_  ,	subfolder="value_function"						)
      torch.manual_seed(0						)
      if torch.cuda.is_available():
          torch.cuda.manual_seed_all(0						)

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

      with torch.no_grad():
          SCREAMING_SNAKE_CASE__         =    value_function(A_  ,	A_						).sample

      # fmt: off
      SCREAMING_SNAKE_CASE__         =    torch.tensor([165.25] * seq_len						)
      # fmt: on
      self.assertTrue(torch.allclose(A_  ,	A_  ,	rtol=1E-3						)						)







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





from __future__ import annotations





def 				__a							(       __A						):
    SCREAMING_SNAKE_CASE__         =    {1}
    SCREAMING_SNAKE_CASE__         =    [1]
    for _ in range(2			,							depth + 1						):
        SCREAMING_SNAKE_CASE__         =    [*previous_coefficients, 0]
        SCREAMING_SNAKE_CASE__         =    [0, *previous_coefficients]
        SCREAMING_SNAKE_CASE__         =    []
        for x, y in zip(__A			,							__A						):
            coefficients.add(x + y						)
            previous_coefficients.append(x + y						)
    return coefficients





def 				__a							(       __A						):
    SCREAMING_SNAKE_CASE__         =    set()
    for number in unique_coefficients:
        SCREAMING_SNAKE_CASE__         =    2
        SCREAMING_SNAKE_CASE__         =    number
        while divisor**2 <= copy_number:
            SCREAMING_SNAKE_CASE__         =    0
            while copy_number % divisor == 0:
                copy_number //= divisor
                multiplicity += 1
            if multiplicity >= 2:
                non_squarefrees.add(__A						)
                break
            divisor += 1

    return unique_coefficients.difference(__A						)





def 				__a							(       __A = 51						):
    SCREAMING_SNAKE_CASE__         =    get_pascal_triangle_unique_coefficients(__A						)
    SCREAMING_SNAKE_CASE__         =    get_squarefrees(__A						)
    return sum(__A						)


if __name__ == "__main__":
  print(f'''{solution() = }''')


