

# Copyright 2023 The HuggingFace Inc. team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
import ast
import difflib
from collections.abc import Mapping
from typing import Any, Callable, Dict



class 					lowercase		( _UpperCamelCase      ):



 '''simple docstring'''


 pass







def 		UpperCamelCase_(    snake_case__:						str						,					snake_case__:						Dict[str, Callable]						,					snake_case__:						Dict=None						,					snake_case__:						int=False     ) ->				List[Any]:
 try:
  UpperCAmelCase__           =       ast.parse(snake_case__     )
 except SyntaxError as e:
  print('The code generated by the agent is not valid.\n'						,					snake_case__     )
  return
 if state is None:
  UpperCAmelCase__           =       {}
 UpperCAmelCase__           =       None
 for idx, node in enumerate(expression.body     ):
  try:
   UpperCAmelCase__           =       evaluate_ast(snake_case__						,					snake_case__						,					snake_case__     )
  except InterpretorError as e:
   UpperCAmelCase__           =       f"Evaluation of the code stopped at line {idx} before the end because of the following error"
   if chat_mode:
    msg += (
        f". Copy paste the following error message and send it back to the agent:\nI get an error: '{e}'"
    )
   else:
    msg += f":\n{e}"
   print(snake_case__     )
   break
  if line_result is not None:
   UpperCAmelCase__           =       line_result

 return result
def 		UpperCamelCase_(    snake_case__:						ast.AST						,					snake_case__:						Dict[str, Any]						,					snake_case__:						Dict[str, Callable]     ) ->				int:
 if isinstance(snake_case__						,					ast.Assign     ):
  # Assignement -> we evaluate the assignement which should update the state
  # We return the variable assigned as it may be used to determine the final result.
  return evaluate_assign(snake_case__						,					snake_case__						,					snake_case__     )
 elif isinstance(snake_case__						,					ast.Call     ):
  # Function call -> we return the value of the function call
  return evaluate_call(snake_case__						,					snake_case__						,					snake_case__     )
 elif isinstance(snake_case__						,					ast.Constant     ):
  # Constant -> just return the value
  return expression.value
 elif isinstance(snake_case__						,					ast.Dict     ):
  # Dict -> evaluate all keys and values
  UpperCAmelCase__           =       [evaluate_ast(snake_case__						,					snake_case__						,					snake_case__     ) for k in expression.keys]
  UpperCAmelCase__           =       [evaluate_ast(snake_case__						,					snake_case__						,					snake_case__     ) for v in expression.values]
  return dict(zip(snake_case__						,					snake_case__     )     )
 elif isinstance(snake_case__						,					ast.Expr     ):
  # Expression -> evaluate the content
  return evaluate_ast(expression.value						,					snake_case__						,					snake_case__     )
 elif isinstance(snake_case__						,					ast.For     ):
  # For loop -> execute the loop
  return evaluate_for(snake_case__						,					snake_case__						,					snake_case__     )
 elif isinstance(snake_case__						,					ast.FormattedValue     ):
  # Formatted value (part of f-string) -> evaluate the content and return
  return evaluate_ast(expression.value						,					snake_case__						,					snake_case__     )
 elif isinstance(snake_case__						,					ast.If     ):
  # If -> execute the right branch
  return evaluate_if(snake_case__						,					snake_case__						,					snake_case__     )
 elif hasattr(snake_case__						,					'Index'     ) and isinstance(snake_case__						,					ast.Index     ):
  return evaluate_ast(expression.value						,					snake_case__						,					snake_case__     )
 elif isinstance(snake_case__						,					ast.JoinedStr     ):
  return "".join([str(evaluate_ast(snake_case__						,					snake_case__						,					snake_case__     )     ) for v in expression.values]     )
 elif isinstance(snake_case__						,					ast.List     ):
  # List -> evaluate all elements
  return [evaluate_ast(snake_case__						,					snake_case__						,					snake_case__     ) for elt in expression.elts]
 elif isinstance(snake_case__						,					ast.Name     ):
  # Name -> pick up the value in the state
  return evaluate_name(snake_case__						,					snake_case__						,					snake_case__     )
 elif isinstance(snake_case__						,					ast.Subscript     ):
  # Subscript -> return the value of the indexing
  return evaluate_subscript(snake_case__						,					snake_case__						,					snake_case__     )
 else:
  # For now we refuse anything else. Let's add things as we need them.
  raise InterpretorError(f"{expression.__class__.__name__} is not supported."     )
def 		UpperCamelCase_(    snake_case__:						List[str]						,					snake_case__:						Optional[Any]						,					snake_case__:						Optional[Any]     ) ->				Dict:
 UpperCAmelCase__           =       assign.targets
 UpperCAmelCase__           =       evaluate_ast(assign.value						,					snake_case__						,					snake_case__     )

 if len(snake_case__     ) == 1:
  UpperCAmelCase__           =       result
 else:
  if len(snake_case__     ) != len(snake_case__     ):
   raise InterpretorError(f"Expected {len(snake_case__     )} values but got {len(snake_case__     )}."     )
  for var_name, r in zip(snake_case__						,					snake_case__     ):
   UpperCAmelCase__           =       r
 return result
def 		UpperCamelCase_(    snake_case__:						Optional[Any]						,					snake_case__:						Optional[Any]						,					snake_case__:						Any     ) ->				List[Any]:
 if not isinstance(call.func						,					ast.Name     ):
  raise InterpretorError(
      f"It is not permitted to evaluate other functions than the provided tools (tried to execute {call.func} of "
      f"type {type(call.func     )}."     )
 UpperCAmelCase__           =       call.func.id
 if func_name not in tools:
  raise InterpretorError(
      f"It is not permitted to evaluate other functions than the provided tools (tried to execute {call.func.id})."     )

 UpperCAmelCase__           =       tools[func_name]
 # Todo deal with args
 UpperCAmelCase__           =       [evaluate_ast(snake_case__						,					snake_case__						,					snake_case__     ) for arg in call.args]
 UpperCAmelCase__           =       {keyword.arg: evaluate_ast(keyword.value						,					snake_case__						,					snake_case__     ) for keyword in call.keywords}
 return func(*snake_case__						,					**snake_case__     )
def 		UpperCamelCase_(    snake_case__:						Union[str, Any]						,					snake_case__:						Union[str, Any]						,					snake_case__:						Optional[int]     ) ->				str:
 UpperCAmelCase__           =       evaluate_ast(subscript.slice						,					snake_case__						,					snake_case__     )
 UpperCAmelCase__           =       evaluate_ast(subscript.value						,					snake_case__						,					snake_case__     )
 if isinstance(snake_case__						,					(list, tuple)     ):
  return value[int(snake_case__     )]
 if index in value:
  return value[index]
 if isinstance(snake_case__						,					snake_case__     ) and isinstance(snake_case__						,					snake_case__     ):
  UpperCAmelCase__           =       difflib.get_close_matches(snake_case__						,					list(value.keys()     )     )
  if len(snake_case__     ) > 0:
   return value[close_matches[0]]

 raise InterpretorError(f"Could not index {value} with '{index}'."     )
def 		UpperCamelCase_(    snake_case__:						Optional[int]						,					snake_case__:						Dict						,					snake_case__:						Dict     ) ->				Optional[Any]:
 if name.id in state:
  return state[name.id]
 UpperCAmelCase__           =       difflib.get_close_matches(name.id						,					list(state.keys()     )     )
 if len(snake_case__     ) > 0:
  return state[close_matches[0]]
 raise InterpretorError(f"The variable `{name.id}` is not defined."     )
def 		UpperCamelCase_(    snake_case__:						int						,					snake_case__:						List[str]						,					snake_case__:						Tuple     ) ->				int:
 if len(condition.ops     ) > 1:
  raise InterpretorError('Cannot evaluate conditions with multiple operators'     )

 UpperCAmelCase__           =       evaluate_ast(condition.left						,					snake_case__						,					snake_case__     )
 UpperCAmelCase__           =       condition.ops[0]
 UpperCAmelCase__           =       evaluate_ast(condition.comparators[0]						,					snake_case__						,					snake_case__     )

 if isinstance(snake_case__						,					ast.Eq     ):
  return left == right
 elif isinstance(snake_case__						,					ast.NotEq     ):
  return left != right
 elif isinstance(snake_case__						,					ast.Lt     ):
  return left < right
 elif isinstance(snake_case__						,					ast.LtE     ):
  return left <= right
 elif isinstance(snake_case__						,					ast.Gt     ):
  return left > right
 elif isinstance(snake_case__						,					ast.GtE     ):
  return left >= right
 elif isinstance(snake_case__						,					ast.Is     ):
  return left is right
 elif isinstance(snake_case__						,					ast.IsNot     ):
  return left is not right
 elif isinstance(snake_case__						,					ast.In     ):
  return left in right
 elif isinstance(snake_case__						,					ast.NotIn     ):
  return left not in right
 else:
  raise InterpretorError(f"Operator not supported: {comparator}"     )
def 		UpperCamelCase_(    snake_case__:						List[Any]						,					snake_case__:						Optional[Any]						,					snake_case__:						Optional[int]     ) ->				Union[str, Any]:
 UpperCAmelCase__           =       None
 if evaluate_condition(if_statement.test						,					snake_case__						,					snake_case__     ):
  for line in if_statement.body:
   UpperCAmelCase__           =       evaluate_ast(snake_case__						,					snake_case__						,					snake_case__     )
   if line_result is not None:
    UpperCAmelCase__           =       line_result
 else:
  for line in if_statement.orelse:
   UpperCAmelCase__           =       evaluate_ast(snake_case__						,					snake_case__						,					snake_case__     )
   if line_result is not None:
    UpperCAmelCase__           =       line_result
 return result







def 		UpperCamelCase_(    snake_case__:						Any						,					snake_case__:						List[Any]						,					snake_case__:						Dict     ) ->				Optional[int]:
 UpperCAmelCase__           =       None
 UpperCAmelCase__           =       evaluate_ast(for_loop.iter						,					snake_case__						,					snake_case__     )
 for counter in iterator:
  UpperCAmelCase__           =       counter
  for expression in for_loop.body:
   UpperCAmelCase__           =       evaluate_ast(snake_case__						,					snake_case__						,					snake_case__     )
   if line_result is not None:
    UpperCAmelCase__           =       line_result
 return result






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



 '''simple docstring'''


 __SCREAMING_SNAKE_CASE      							=       UNetaDModel
 __SCREAMING_SNAKE_CASE      							=       """sample"""




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






  """simple docstring"""





  UpperCAmelCase__           =       4
  UpperCAmelCase__           =       14
  UpperCAmelCase__           =       16

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

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




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






  """simple docstring"""





  return (4, 14, 16)




 @property
 def 						UpperCamelCase__							(self				)				->      Dict:






  """simple docstring"""





  return (4, 14, 16)




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






  """simple docstring"""





  pass




 def 						UpperCamelCase__							(self				)				->      Any:






  """simple docstring"""





  pass




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






  """simple docstring"""





  super().test_determinism()




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






  """simple docstring"""





  super().test_outputs_equivalence()




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






  """simple docstring"""





  super().test_from_save_pretrained()




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






  """simple docstring"""





  super().test_from_save_pretrained_variant()




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






  """simple docstring"""





  super().test_model_from_pretrained()




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






  """simple docstring"""





  super().test_output()




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






  """simple docstring"""





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






  """simple docstring"""





  UpperCAmelCase__       ,					UpperCAmelCase__           =       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				)
  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 						UpperCamelCase__							(self				)				->      int:






  """simple docstring"""





  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(__a , __a				).sample.permute(0 , 2 , 1				)

  UpperCAmelCase__           =       output[0, -3:, -3:].flatten()
  # fmt: off
  UpperCAmelCase__           =       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(__a , __a , rtol=1E-3				)				)




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






  """simple docstring"""





  pass




 @slow
 def 						UpperCamelCase__							(self				)				->      Tuple:






  """simple docstring"""





  UpperCAmelCase__           =       'harmonai/maestro-150k'
  UpperCAmelCase__           =       UNetaDModel.from_pretrained(__a , subfolder='unet'				)
  model.to(__a				)

  UpperCAmelCase__           =       65536
  UpperCAmelCase__           =       torch.sin(torch.arange(__a				)[None, None, :].repeat(1 , 2 , 1				)				).to(__a				)
  UpperCAmelCase__           =       torch.tensor([1]				).to(__a				)

  with torch.no_grad():
   UpperCAmelCase__           =       model(__a , __a				).sample

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

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







class 					lowercase		( _UpperCamelCase				,       _UpperCamelCase				,       unittest.TestCase      ):



 '''simple docstring'''


 __SCREAMING_SNAKE_CASE      							=       UNetaDModel
 __SCREAMING_SNAKE_CASE      							=       """sample"""




 @property
 def 						UpperCamelCase__							(self				)				->      Optional[int]:






  """simple docstring"""





  UpperCAmelCase__           =       4
  UpperCAmelCase__           =       14
  UpperCAmelCase__           =       16

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

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




 @property
 def 						UpperCamelCase__							(self				)				->      int:






  """simple docstring"""





  return (4, 14, 16)




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






  """simple docstring"""





  return (4, 14, 1)




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






  """simple docstring"""





  super().test_determinism()




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






  """simple docstring"""





  super().test_outputs_equivalence()




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






  """simple docstring"""





  super().test_from_save_pretrained()




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






  """simple docstring"""





  super().test_from_save_pretrained_variant()




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






  """simple docstring"""





  super().test_model_from_pretrained()




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






  """simple docstring"""





  UpperCAmelCase__       ,					UpperCAmelCase__           =       self.prepare_init_args_and_inputs_for_common()
  UpperCAmelCase__           =       self.model_class(**__a				)
  model.to(__a				)
  model.eval()

  with torch.no_grad():
   UpperCAmelCase__           =       model(**__a				)

   if isinstance(__a , __a				):
    UpperCAmelCase__           =       output.sample

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




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






  """simple docstring"""





  pass




 def 						UpperCamelCase__							(self				)				->      Any:






  """simple docstring"""





  pass




 def 						UpperCamelCase__							(self				)				->      Any:






  """simple docstring"""





  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 						UpperCamelCase__							(self				)				->      Union[str, Any]:






  """simple docstring"""





  UpperCAmelCase__       ,					UpperCAmelCase__           =       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				)
  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 						UpperCamelCase__							(self				)				->      int:






  """simple docstring"""





  UpperCAmelCase__       ,					UpperCAmelCase__           =       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				)

  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(__a , __a				).sample

  # fmt: off
  UpperCAmelCase__           =       torch.tensor([1_65.25] * seq_len				)
  # fmt: on
  self.assertTrue(torch.allclose(__a , __a , rtol=1E-3				)				)




 def 						UpperCamelCase__							(self				)				->      Tuple:






  """simple docstring"""





  pass



