


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






    '''simple docstring'''





    A_	:  Tuple								=			UNetaDModel
    A_	:  Optional[int]								=			"sample"





    @property
    def        _lowerCamelCase			(						self					: Optional[Any]						):





       '''simple docstring'''
       UpperCAmelCase_				:   Union[str, Any]							= 4
       UpperCAmelCase_				:   List[Any]							= 14
       UpperCAmelCase_				:   Optional[int]							= 16

       UpperCAmelCase_				:   Union[str, Any]							= floats_tensor((batch_size, num_features, seq_len)						).to(_lowerCamelCase						)
       UpperCAmelCase_				:   Union[str, Any]							= torch.tensor([10] * batch_size						).to(_lowerCamelCase						)

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





    @property
    def        _lowerCamelCase			(						self					: Tuple						):





       '''simple docstring'''
       return (4, 14, 16)





    @property
    def        _lowerCamelCase			(						self					: Any						):





       '''simple docstring'''
       return (4, 14, 16)





    def        _lowerCamelCase			(						self					: str						):





       '''simple docstring'''
       pass





    def        _lowerCamelCase			(						self					: Any						):





       '''simple docstring'''
       pass





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





       '''simple docstring'''
       super().test_determinism()





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





       '''simple docstring'''
       super().test_outputs_equivalence()





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





       '''simple docstring'''
       super().test_from_save_pretrained()





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





       '''simple docstring'''
       super().test_from_save_pretrained_variant()





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





       '''simple docstring'''
       super().test_model_from_pretrained()





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





       '''simple docstring'''
       super().test_output()





    def        _lowerCamelCase			(						self					: int						):





       '''simple docstring'''
       UpperCAmelCase_				:   Optional[Any]							= {
           '''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_				:   Optional[Any]							= self.dummy_input
       return init_dict, inputs_dict





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





       '''simple docstring'''
       UpperCAmelCase_		,   UpperCAmelCase_				:   Optional[int]							= UNetaDModel.from_pretrained(
           '''bglick13/hopper-medium-v2-value-function-hor32'''							, output_loading_info=_lowerCamelCase							, subfolder='''unet'''						)
       self.assertIsNotNone(_lowerCamelCase						)
       self.assertEqual(len(loading_info['''missing_keys''']						)							, 0						)

       model.to(_lowerCamelCase						)
       UpperCAmelCase_				:   Optional[Any]							= 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        _lowerCamelCase			(						self					: Any						):





       '''simple docstring'''
       UpperCAmelCase_				:   int							= 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_				:   Dict							= model.config.in_channels
       UpperCAmelCase_				:   Any							= 16
       UpperCAmelCase_				:   str							= torch.randn((1, seq_len, num_features)						).permute(
           0							, 2							, 1						)  # match original, we can update values and remove
       UpperCAmelCase_				:   List[Any]							= torch.full((num_features,)							, 0						)

       with torch.no_grad():
          UpperCAmelCase_				:   List[Any]							= model(_lowerCamelCase							, _lowerCamelCase						).sample.permute(0							, 2							, 1						)

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





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





       '''simple docstring'''
       # Not implemented yet for this UNet
       pass





    @slow
    def        _lowerCamelCase			(						self					: List[Any]						):





       '''simple docstring'''
       UpperCAmelCase_				:   Optional[Any]							= '''harmonai/maestro-150k'''
       UpperCAmelCase_				:   List[Any]							= UNetaDModel.from_pretrained(_lowerCamelCase							, subfolder='''unet'''						)
       model.to(_lowerCamelCase						)

       UpperCAmelCase_				:   List[Any]							= 65_536
       UpperCAmelCase_				:   List[Any]							= torch.sin(torch.arange(_lowerCamelCase						)[None, None, :].repeat(1							, 2							, 1						)						).to(_lowerCamelCase						)
       UpperCAmelCase_				:   Tuple							= torch.tensor([1]						).to(_lowerCamelCase						)

       with torch.no_grad():
          UpperCAmelCase_				:   int							= model(_lowerCamelCase							, _lowerCamelCase						).sample

       UpperCAmelCase_				:   int							= output.abs().sum()
       UpperCAmelCase_				:   Tuple							= output.abs().max()

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





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






    '''simple docstring'''





    A_	:  Dict								=			UNetaDModel
    A_	:  str								=			"sample"





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





       '''simple docstring'''
       UpperCAmelCase_				:   Tuple							= 4
       UpperCAmelCase_				:   Any							= 14
       UpperCAmelCase_				:   Union[str, Any]							= 16

       UpperCAmelCase_				:   Union[str, Any]							= floats_tensor((batch_size, num_features, seq_len)						).to(_lowerCamelCase						)
       UpperCAmelCase_				:   List[Any]							= torch.tensor([10] * batch_size						).to(_lowerCamelCase						)

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





    @property
    def        _lowerCamelCase			(						self					: Any						):





       '''simple docstring'''
       return (4, 14, 16)





    @property
    def        _lowerCamelCase			(						self					: Any						):





       '''simple docstring'''
       return (4, 14, 1)





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





       '''simple docstring'''
       super().test_determinism()





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





       '''simple docstring'''
       super().test_outputs_equivalence()





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





       '''simple docstring'''
       super().test_from_save_pretrained()





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





       '''simple docstring'''
       super().test_from_save_pretrained_variant()





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





       '''simple docstring'''
       super().test_model_from_pretrained()





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





       '''simple docstring'''
       # UNetRL is a value-function is different output shape
       UpperCAmelCase_		,   UpperCAmelCase_				:   Optional[int]							= self.prepare_init_args_and_inputs_for_common()
       UpperCAmelCase_				:   Optional[int]							= self.model_class(**_lowerCamelCase						)
       model.to(_lowerCamelCase						)
       model.eval()

       with torch.no_grad():
          UpperCAmelCase_				:   Optional[Any]							= model(**_lowerCamelCase						)

          if isinstance(_lowerCamelCase							, _lowerCamelCase						):
             UpperCAmelCase_				:   List[Any]							= output.sample

       self.assertIsNotNone(_lowerCamelCase						)
       UpperCAmelCase_				:   Union[str, Any]							= torch.Size((inputs_dict['''sample'''].shape[0], 1)						)
       self.assertEqual(output.shape							, _lowerCamelCase							, '''Input and output shapes do not match'''						)





    def        _lowerCamelCase			(						self					: Optional[int]						):





       '''simple docstring'''
       pass





    def        _lowerCamelCase			(						self					: Optional[Any]						):





       '''simple docstring'''
       pass





    def        _lowerCamelCase			(						self					: Optional[int]						):





       '''simple docstring'''
       UpperCAmelCase_				:   str							= {
           '''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_				:   int							= self.dummy_input
       return init_dict, inputs_dict





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





       '''simple docstring'''
       UpperCAmelCase_		,   UpperCAmelCase_				:   Optional[int]							= UNetaDModel.from_pretrained(
           '''bglick13/hopper-medium-v2-value-function-hor32'''							, output_loading_info=_lowerCamelCase							, subfolder='''value_function'''						)
       self.assertIsNotNone(_lowerCamelCase						)
       self.assertEqual(len(vf_loading_info['''missing_keys''']						)							, 0						)

       value_function.to(_lowerCamelCase						)
       UpperCAmelCase_				:   int							= 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        _lowerCamelCase			(						self					: Optional[int]						):





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

       UpperCAmelCase_				:   Any							= value_function.config.in_channels
       UpperCAmelCase_				:   List[Any]							= 14
       UpperCAmelCase_				:   Tuple							= torch.randn((1, seq_len, num_features)						).permute(
           0							, 2							, 1						)  # match original, we can update values and remove
       UpperCAmelCase_				:   Optional[Any]							= torch.full((num_features,)							, 0						)

       with torch.no_grad():
          UpperCAmelCase_				:   List[Any]							= value_function(_lowerCamelCase							, _lowerCamelCase						).sample

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





    def        _lowerCamelCase			(						self					: Tuple						):





       '''simple docstring'''
       # Not implemented yet for this UNet
       pass







"""simple docstring"""



import json
import os
import random
import shutil
import tempfile
import unittest

import numpy as np
import pytest

from transformers import BertTokenizer, BertTokenizerFast
from transformers.models.bert.tokenization_bert import VOCAB_FILES_NAMES
from transformers.testing_utils import require_vision
from transformers.utils import IMAGE_PROCESSOR_NAME, is_vision_available


if is_vision_available():
    from PIL import Image

    from transformers import FlavaImageProcessor, FlavaProcessor
    from transformers.models.flava.image_processing_flava import (
        FLAVA_CODEBOOK_MEAN,
        FLAVA_CODEBOOK_STD,
        FLAVA_IMAGE_MEAN,
        FLAVA_IMAGE_STD,
    )

@require_vision
class 			a	(  unittest.TestCase      ):







    def      UpperCamelCase_ (     self			):
      lowercase				       =			tempfile.mkdtemp()

      # fmt: off
      lowercase				       =			['[UNK]', '[CLS]', '[SEP]', '[PAD]', '[MASK]', 'want', '##want', '##ed', 'wa', 'un', 'runn', '##ing', ',', 'low', 'lowest']
      # fmt: on
      lowercase				       =			os.path.join(self.tmpdirname				,						VOCAB_FILES_NAMES['vocab_file']			)

      with open(self.vocab_file				,						'w'				,						encoding='utf-8'			) as fp:
        fp.write(''.join([x + '\n' for x in vocab_tokens]			)			)

      lowercase				       =			{
          'image_mean': FLAVA_IMAGE_MEAN,
          'image_std': FLAVA_IMAGE_STD,
          'do_normalize': True,
          'do_resize': True,
          'size': 2_2_4,
          'do_center_crop': True,
          'crop_size': 2_2_4,
          'input_size_patches': 1_4,
          'total_mask_patches': 7_5,
          'mask_group_max_patches': None,
          'mask_group_min_patches': 1_6,
          'mask_group_min_aspect_ratio': 0.3,
          'mask_group_max_aspect_ratio': None,
          'codebook_do_resize': True,
          'codebook_size': 1_1_2,
          'codebook_do_center_crop': True,
          'codebook_crop_size': 1_1_2,
          'codebook_do_map_pixels': True,
          'codebook_do_normalize': True,
          'codebook_image_mean': FLAVA_CODEBOOK_MEAN,
          'codebook_image_std': FLAVA_CODEBOOK_STD,
      }

      lowercase				       =			os.path.join(self.tmpdirname				,						_lowerCamelCase			)
      with open(self.image_processor_file				,						'w'				,						encoding='utf-8'			) as fp:
        json.dump(_lowerCamelCase				,						_lowerCamelCase			)







    def      UpperCamelCase_ (     self				,						**_lowerCamelCase			):
      return BertTokenizer.from_pretrained(self.tmpdirname				,						**_lowerCamelCase			)







    def      UpperCamelCase_ (     self				,						**_lowerCamelCase			):
      return BertTokenizerFast.from_pretrained(self.tmpdirname				,						**_lowerCamelCase			)







    def      UpperCamelCase_ (     self				,						**_lowerCamelCase			):
      return FlavaImageProcessor.from_pretrained(self.tmpdirname				,						**_lowerCamelCase			)







    def      UpperCamelCase_ (     self			):
      shutil.rmtree(self.tmpdirname			)







    def      UpperCamelCase_ (     self			):
      lowercase				       =			[np.random.randint(2_5_5				,						size=(3, 3_0, 4_0_0)				,						dtype=np.uinta			)]

      lowercase				       =			[Image.fromarray(np.moveaxis(_lowerCamelCase				,						0				,						-1			)			) for x in image_inputs]

      return image_inputs







    def      UpperCamelCase_ (     self			):
      lowercase				       =			self.get_tokenizer()
      lowercase				       =			self.get_rust_tokenizer()
      lowercase				       =			self.get_image_processor()

      lowercase				       =			FlavaProcessor(tokenizer=_lowerCamelCase				,						image_processor=_lowerCamelCase			)
      processor_slow.save_pretrained(self.tmpdirname			)
      lowercase				       =			FlavaProcessor.from_pretrained(self.tmpdirname				,						use_fast=_lowerCamelCase			)

      lowercase				       =			FlavaProcessor(tokenizer=_lowerCamelCase				,						image_processor=_lowerCamelCase			)
      processor_fast.save_pretrained(self.tmpdirname			)
      lowercase				       =			FlavaProcessor.from_pretrained(self.tmpdirname			)

      self.assertEqual(processor_slow.tokenizer.get_vocab()				,						tokenizer_slow.get_vocab()			)
      self.assertEqual(processor_fast.tokenizer.get_vocab()				,						tokenizer_fast.get_vocab()			)
      self.assertEqual(tokenizer_slow.get_vocab()				,						tokenizer_fast.get_vocab()			)
      self.assertIsInstance(processor_slow.tokenizer				,						_lowerCamelCase			)
      self.assertIsInstance(processor_fast.tokenizer				,						_lowerCamelCase			)

      self.assertEqual(processor_slow.image_processor.to_json_string()				,						image_processor.to_json_string()			)
      self.assertEqual(processor_fast.image_processor.to_json_string()				,						image_processor.to_json_string()			)
      self.assertIsInstance(processor_slow.image_processor				,						_lowerCamelCase			)
      self.assertIsInstance(processor_fast.image_processor				,						_lowerCamelCase			)







    def      UpperCamelCase_ (     self			):
      lowercase				       =			FlavaProcessor(tokenizer=self.get_tokenizer()				,						image_processor=self.get_image_processor()			)
      processor.save_pretrained(self.tmpdirname			)

      lowercase				       =			self.get_tokenizer(bos_token='(BOS)'				,						eos_token='(EOS)'			)
      lowercase				       =			self.get_image_processor(do_normalize=_lowerCamelCase				,						padding_value=1.0			)

      lowercase				       =			FlavaProcessor.from_pretrained(
          self.tmpdirname				,						bos_token='(BOS)'				,						eos_token='(EOS)'				,						do_normalize=_lowerCamelCase				,						padding_value=1.0			)

      self.assertEqual(processor.tokenizer.get_vocab()				,						tokenizer_add_kwargs.get_vocab()			)
      self.assertIsInstance(processor.tokenizer				,						_lowerCamelCase			)

      self.assertEqual(processor.image_processor.to_json_string()				,						image_processor_add_kwargs.to_json_string()			)
      self.assertIsInstance(processor.image_processor				,						_lowerCamelCase			)







    def      UpperCamelCase_ (     self			):
      lowercase				       =			self.get_image_processor()
      lowercase				       =			self.get_tokenizer()

      lowercase				       =			FlavaProcessor(tokenizer=_lowerCamelCase				,						image_processor=_lowerCamelCase			)

      lowercase				       =			self.prepare_image_inputs()

      lowercase				       =			image_processor(_lowerCamelCase				,						return_tensors='np'			)
      lowercase				       =			processor(images=_lowerCamelCase				,						return_tensors='np'			)

      for key in input_feat_extract.keys():
        self.assertAlmostEqual(input_feat_extract[key].sum()				,						input_processor[key].sum()				,						delta=1e-2			)

      # With rest of the args
      random.seed(1_2_3_4			)
      lowercase				       =			image_processor(
          _lowerCamelCase				,						return_image_mask=_lowerCamelCase				,						return_codebook_pixels=_lowerCamelCase				,						return_tensors='np'			)
      random.seed(1_2_3_4			)
      lowercase				       =			processor(
          images=_lowerCamelCase				,						return_image_mask=_lowerCamelCase				,						return_codebook_pixels=_lowerCamelCase				,						return_tensors='np'			)

      for key in input_feat_extract.keys():
        self.assertAlmostEqual(input_feat_extract[key].sum()				,						input_processor[key].sum()				,						delta=1e-2			)







    def      UpperCamelCase_ (     self			):
      lowercase				       =			self.get_image_processor()
      lowercase				       =			self.get_tokenizer()

      lowercase				       =			FlavaProcessor(tokenizer=_lowerCamelCase				,						image_processor=_lowerCamelCase			)

      lowercase				       =			'lower newer'

      lowercase				       =			processor(text=_lowerCamelCase			)

      lowercase				       =			tokenizer(_lowerCamelCase			)

      for key in encoded_tok.keys():
        self.assertListEqual(encoded_tok[key]				,						encoded_processor[key]			)







    def      UpperCamelCase_ (     self			):
      lowercase				       =			self.get_image_processor()
      lowercase				       =			self.get_tokenizer()

      lowercase				       =			FlavaProcessor(tokenizer=_lowerCamelCase				,						image_processor=_lowerCamelCase			)

      lowercase				       =			'lower newer'
      lowercase				       =			self.prepare_image_inputs()

      lowercase				       =			processor(text=_lowerCamelCase				,						images=_lowerCamelCase			)

      self.assertListEqual(list(inputs.keys()			)				,						['input_ids', 'token_type_ids', 'attention_mask', 'pixel_values']			)

      # add extra args
      lowercase				       =			processor(text=_lowerCamelCase				,						images=_lowerCamelCase				,						return_codebook_pixels=_lowerCamelCase				,						return_image_mask=_lowerCamelCase			)

      self.assertListEqual(
          list(inputs.keys()			)				,						[
              'input_ids',
              'token_type_ids',
              'attention_mask',
              'pixel_values',
              'codebook_pixel_values',
              'bool_masked_pos',
          ]				,						)

      # test if it raises when no input is passed
      with pytest.raises(_lowerCamelCase			):
        processor()







    def      UpperCamelCase_ (     self			):
      lowercase				       =			self.get_image_processor()
      lowercase				       =			self.get_tokenizer()

      lowercase				       =			FlavaProcessor(tokenizer=_lowerCamelCase				,						image_processor=_lowerCamelCase			)

      lowercase				       =			[[1, 4, 5, 8, 1, 0, 8], [3, 4, 3, 1, 1, 8, 9]]

      lowercase				       =			processor.batch_decode(_lowerCamelCase			)
      lowercase				       =			tokenizer.batch_decode(_lowerCamelCase			)

      self.assertListEqual(_lowerCamelCase				,						_lowerCamelCase			)







    def      UpperCamelCase_ (     self			):
      lowercase				       =			self.get_image_processor()
      lowercase				       =			self.get_tokenizer()

      lowercase				       =			FlavaProcessor(tokenizer=_lowerCamelCase				,						image_processor=_lowerCamelCase			)

      lowercase				       =			'lower newer'
      lowercase				       =			self.prepare_image_inputs()

      lowercase				       =			processor(text=_lowerCamelCase				,						images=_lowerCamelCase			)

      self.assertListEqual(list(inputs.keys()			)				,						processor.model_input_names			)


