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






 '''simple docstring'''






 UpperCAmelCase_	 =					UNetaDModel
 UpperCAmelCase_	 =					"""sample"""




 @property
 def 		__snake_case						(			self					:  str):







        '''simple docstring'''



        lowerCAmelCase__     	=					4
        lowerCAmelCase__     	=					14
        lowerCAmelCase__     	=					16

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

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




 @property
 def 		__snake_case						(			self					:  Optional[int]):







        '''simple docstring'''



        return (4, 14, 16)




 @property
 def 		__snake_case						(			self					:  List[str]):







        '''simple docstring'''



        return (4, 14, 16)




 def 		__snake_case						(			self					:  Dict):







        '''simple docstring'''



        pass




 def 		__snake_case						(			self					:  Optional[Any]):







        '''simple docstring'''



        pass




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







        '''simple docstring'''



        super().test_determinism()




 @unittest.skipIf(torch_device == 'mps'			,    'mish op not supported in MPS')
 def 		__snake_case						(			self					:  str):







        '''simple docstring'''



        super().test_outputs_equivalence()




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







        '''simple docstring'''



        super().test_from_save_pretrained()




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







        '''simple docstring'''



        super().test_from_save_pretrained_variant()




 @unittest.skipIf(torch_device == 'mps'			,    'mish op not supported in MPS')
 def 		__snake_case						(			self					:  Dict):







        '''simple docstring'''



        super().test_model_from_pretrained()




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







        '''simple docstring'''



        super().test_output()




 def 		__snake_case						(			self					:  Optional[int]):







        '''simple docstring'''



        lowerCAmelCase__     	=					{
            '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__     	=					self.dummy_input
        return init_dict, inputs_dict




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







        '''simple docstring'''



        lowerCAmelCase__     	=					UNetaDModel.from_pretrained(
            'bglick13/hopper-medium-v2-value-function-hor32'			,    output_loading_info=lowercase__			,    subfolder='unet')
        self.assertIsNotNone(lowercase__)
        self.assertEqual(len(loading_info['missing_keys'])			,    0)

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







        '''simple docstring'''



        lowerCAmelCase__     	=					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__     	=					model.config.in_channels
        lowerCAmelCase__     	=					16
        lowerCAmelCase__     	=					torch.randn((1, seq_len, num_features)).permute(
            0			,    2			,    1)  # match original, we can update values and remove
        lowerCAmelCase__     	=					torch.full((num_features,)			,    0)

        with torch.no_grad():
               lowerCAmelCase__     	=					model(lowercase__			,    lowercase__).sample.permute(0			,    2			,    1)

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




 def 		__snake_case						(			self					:  Union[str, Any]):







        '''simple docstring'''



        pass




 @slow
 def 		__snake_case						(			self					:  str):







        '''simple docstring'''



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

        lowerCAmelCase__     	=					65_536
        lowerCAmelCase__     	=					torch.sin(torch.arange(lowercase__)[None, None, :].repeat(1			,    2			,    1)).to(lowercase__)
        lowerCAmelCase__     	=					torch.tensor([1]).to(lowercase__)

        with torch.no_grad():
               lowerCAmelCase__     	=					model(lowercase__			,    lowercase__).sample

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

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






class 				a_ (						UpperCAmelCase_       ,     UpperCAmelCase_       ,     unittest.TestCase						):






 '''simple docstring'''






 UpperCAmelCase_	 =					UNetaDModel
 UpperCAmelCase_	 =					"""sample"""




 @property
 def 		__snake_case						(			self					:  Dict):







        '''simple docstring'''



        lowerCAmelCase__     	=					4
        lowerCAmelCase__     	=					14
        lowerCAmelCase__     	=					16

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

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




 @property
 def 		__snake_case						(			self					:  str):







        '''simple docstring'''



        return (4, 14, 16)




 @property
 def 		__snake_case						(			self					:  Tuple):







        '''simple docstring'''



        return (4, 14, 1)




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







        '''simple docstring'''



        super().test_determinism()




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







        '''simple docstring'''



        super().test_outputs_equivalence()




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







        '''simple docstring'''



        super().test_from_save_pretrained()




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







        '''simple docstring'''



        super().test_from_save_pretrained_variant()




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







        '''simple docstring'''



        super().test_model_from_pretrained()




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







        '''simple docstring'''



        lowerCAmelCase__     	=					self.prepare_init_args_and_inputs_for_common()
        lowerCAmelCase__     	=					self.model_class(**lowercase__)
        model.to(lowercase__)
        model.eval()

        with torch.no_grad():
               lowerCAmelCase__     	=					model(**lowercase__)

               if isinstance(lowercase__			,    lowercase__):
                      lowerCAmelCase__     	=					output.sample

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




 def 		__snake_case						(			self					:  str):







        '''simple docstring'''



        pass




 def 		__snake_case						(			self					:  List[Any]):







        '''simple docstring'''



        pass




 def 		__snake_case						(			self					:  int):







        '''simple docstring'''



        lowerCAmelCase__     	=					{
            '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__     	=					self.dummy_input
        return init_dict, inputs_dict




 @unittest.skipIf(torch_device == 'mps'			,    'mish op not supported in MPS')
 def 		__snake_case						(			self					:  str):







        '''simple docstring'''



        lowerCAmelCase__     	=					UNetaDModel.from_pretrained(
            'bglick13/hopper-medium-v2-value-function-hor32'			,    output_loading_info=lowercase__			,    subfolder='value_function')
        self.assertIsNotNone(lowercase__)
        self.assertEqual(len(vf_loading_info['missing_keys'])			,    0)

        value_function.to(lowercase__)
        lowerCAmelCase__     	=					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 		__snake_case						(			self					:  Dict):







        '''simple docstring'''



        lowerCAmelCase__     	=					UNetaDModel.from_pretrained(
            'bglick13/hopper-medium-v2-value-function-hor32'			,    output_loading_info=lowercase__			,    subfolder='value_function')
        torch.manual_seed(0)
        if torch.cuda.is_available():
               torch.cuda.manual_seed_all(0)

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

        with torch.no_grad():
               lowerCAmelCase__     	=					value_function(lowercase__			,    lowercase__).sample

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




 def 		__snake_case						(			self					:  List[str]):







        '''simple docstring'''



        pass







import warnings
from typing import List, Optional, Union

from ...processing_utils import ProcessorMixin
from ...tokenization_utils_base import BatchEncoding, PaddingStrategy, PreTokenizedInput, TextInput, TruncationStrategy
from ...utils import TensorType



class 				_lowerCAmelCase			(	UpperCAmelCase_							):


  '''simple docstring'''
  a_							:				List[str]					    =["""image_processor""", """tokenizer"""]
  a_							:				Optional[Any]					    ="""LayoutLMv2ImageProcessor"""
  a_							:				Any					    =("""LayoutLMv2Tokenizer""", """LayoutLMv2TokenizerFast""")







  def __init__(					self					:    Union[str, Any]					, UpperCamelCase					:    str=None					, UpperCamelCase					:    List[Any]=None					, **UpperCamelCase					:    List[str]							):


    '''simple docstring'''





    _snake_case      : Dict														=		None
    if "feature_extractor" in kwargs:
      warnings.warn(
          'The `feature_extractor` argument is deprecated and will be removed in v5, use `image_processor`'
          ' instead.'					, UpperCamelCase					, )
      _snake_case      : Dict														=		kwargs.pop('feature_extractor'							)

    _snake_case      : Optional[Any]														=		image_processor if image_processor is not None else feature_extractor
    if image_processor is None:
      raise ValueError('You need to specify an `image_processor`.'							)
    if tokenizer is None:
      raise ValueError('You need to specify a `tokenizer`.'							)

    super().__init__(UpperCamelCase					, UpperCamelCase							)







  def __call__(					self					:    int					, UpperCamelCase					:    Tuple					, UpperCamelCase					:    Union[TextInput, PreTokenizedInput, List[TextInput], List[PreTokenizedInput]] = None					, UpperCamelCase					:    Optional[Union[PreTokenizedInput, List[PreTokenizedInput]]] = None					, UpperCamelCase					:    Union[List[List[int]], List[List[List[int]]]] = None					, UpperCamelCase					:    Optional[Union[List[int], List[List[int]]]] = None					, UpperCamelCase					:    bool = True					, UpperCamelCase					:    Union[bool, str, PaddingStrategy] = False					, UpperCamelCase					:    Union[bool, str, TruncationStrategy] = False					, UpperCamelCase					:    Optional[int] = None					, UpperCamelCase					:    int = 0					, UpperCamelCase					:    Optional[int] = None					, UpperCamelCase					:    Optional[bool] = None					, UpperCamelCase					:    Optional[bool] = None					, UpperCamelCase					:    bool = False					, UpperCamelCase					:    bool = False					, UpperCamelCase					:    bool = False					, UpperCamelCase					:    bool = False					, UpperCamelCase					:    bool = True					, UpperCamelCase					:    Optional[Union[str, TensorType]] = None					, **UpperCamelCase					:    Optional[Any]					, ):


    '''simple docstring'''





    if self.image_processor.apply_ocr and (boxes is not None):
      raise ValueError(
          'You cannot provide bounding boxes if you initialized the image processor with apply_ocr set to True.'							)

    if self.image_processor.apply_ocr and (word_labels is not None):
      raise ValueError(
          'You cannot provide word labels if you initialized the image processor with apply_ocr set to True.'							)

    if return_overflowing_tokens is True and return_offsets_mapping is False:
      raise ValueError('You cannot return overflowing tokens without returning the offsets mapping.'							)

    # first, apply the image processor
    _snake_case      : Optional[Any]														=		self.image_processor(images=UpperCamelCase					, return_tensors=UpperCamelCase							)

    # second, apply the tokenizer
    if text is not None and self.image_processor.apply_ocr and text_pair is None:
      if isinstance(UpperCamelCase					, UpperCamelCase							):
        _snake_case      : List[Any]														=		[text]  # add batch dimension (as the image processor always adds a batch dimension)
      _snake_case      : Dict														=		features['words']

    _snake_case      : Optional[int]														=		self.tokenizer(
        text=text if text is not None else features['words']					, text_pair=text_pair if text_pair is not None else None					, boxes=boxes if boxes is not None else features['boxes']					, word_labels=UpperCamelCase					, add_special_tokens=UpperCamelCase					, padding=UpperCamelCase					, truncation=UpperCamelCase					, max_length=UpperCamelCase					, stride=UpperCamelCase					, pad_to_multiple_of=UpperCamelCase					, return_token_type_ids=UpperCamelCase					, return_attention_mask=UpperCamelCase					, return_overflowing_tokens=UpperCamelCase					, return_special_tokens_mask=UpperCamelCase					, return_offsets_mapping=UpperCamelCase					, return_length=UpperCamelCase					, verbose=UpperCamelCase					, return_tensors=UpperCamelCase					, **UpperCamelCase					, )

    # add pixel values
    _snake_case      : Optional[int]														=		features.pop('pixel_values'							)
    if return_overflowing_tokens is True:
      _snake_case      : Any														=		self.get_overflowing_images(UpperCamelCase					, encoded_inputs['overflow_to_sample_mapping']							)
    _snake_case      : int														=		images

    return encoded_inputs







  def UpperCamelCase_       (					self					:    List[Any]					, UpperCamelCase					:    List[str]					, UpperCamelCase					:    List[str]							):


    '''simple docstring'''





    _snake_case      : Optional[Any]														=		[]
    for sample_idx in overflow_to_sample_mapping:
      images_with_overflow.append(images[sample_idx]							)

    if len(UpperCamelCase							) != len(UpperCamelCase							):
      raise ValueError(
          'Expected length of images to be the same as the length of `overflow_to_sample_mapping`, but got'
          f""" {len(UpperCamelCase							)} and {len(UpperCamelCase							)}"""							)

    return images_with_overflow







  def UpperCamelCase_       (					self					:    List[Any]					, *UpperCamelCase					:    Any					, **UpperCamelCase					:    int							):


    '''simple docstring'''





    return self.tokenizer.batch_decode(*UpperCamelCase					, **UpperCamelCase							)







  def UpperCamelCase_       (					self					:    Any					, *UpperCamelCase					:    Optional[int]					, **UpperCamelCase					:    Any							):


    '''simple docstring'''





    return self.tokenizer.decode(*UpperCamelCase					, **UpperCamelCase							)







  @property
  def UpperCamelCase_       (					self					:    str							):


    '''simple docstring'''





    return ["input_ids", "bbox", "token_type_ids", "attention_mask", "image"]







  @property
  def UpperCamelCase_       (					self					:    Dict							):


    '''simple docstring'''





    warnings.warn(
        '`feature_extractor_class` is deprecated and will be removed in v5. Use `image_processor_class` instead.'					, UpperCamelCase					, )
    return self.image_processor_class







  @property
  def UpperCamelCase_       (					self					:    Tuple							):


    '''simple docstring'''





    warnings.warn(
        '`feature_extractor` is deprecated and will be removed in v5. Use `image_processor` instead.'					, UpperCamelCase					, )
    return self.image_processor




