






import math
from typing import Any, Callable, List, Optional, Tuple, Union

import numpy as np
import torch

from ...models import TaFilmDecoder
from ...schedulers import DDPMScheduler
from ...utils import is_onnx_available, logging, randn_tensor


if is_onnx_available():
 from ..onnx_utils import OnnxRuntimeModel

from ..pipeline_utils import AudioPipelineOutput, DiffusionPipeline
from .continous_encoder import SpectrogramContEncoder
from .notes_encoder import SpectrogramNotesEncoder


_a    :   Dict        =			logging.get_logger(__name__)  # pylint: disable=invalid-name

_a    :   List[str]        =			256







class 	_UpperCAmelCase  (  _A		):




      """simple docstring"""



      A          =    ["""melgan"""]
      def __init__(   self    ,	_lowerCAmelCase    ,	_lowerCAmelCase    ,	_lowerCAmelCase    ,	_lowerCAmelCase    ,	_lowerCAmelCase    ,	):



          '''simple docstring'''
          super().__init__()

          # From MELGAN
          lowerCAmelCase__   :Any      							=				math.log(1e-5		)  # Matches MelGAN training.
          lowerCAmelCase__   :Union[str, Any]      							=				4.0  # Largest value for most examples
          lowerCAmelCase__   :str      							=				128

          self.register_modules(
              notes_encoder=_lowerCAmelCase    ,	continuous_encoder=_lowerCAmelCase    ,	decoder=_lowerCAmelCase    ,	scheduler=_lowerCAmelCase    ,	melgan=_lowerCAmelCase    ,	)
      def 			snake_case_      (   self    ,	_lowerCAmelCase    ,	_lowerCAmelCase=(-1.0, 1.0)    ,	_lowerCAmelCase=False		):



          '''simple docstring'''
          lowerCAmelCase__							,lowerCAmelCase__   :List[Any]      							=				output_range
          if clip:
              lowerCAmelCase__   :Any      							=				torch.clip(_lowerCAmelCase    ,	self.min_value    ,	self.max_value		)
          # Scale to [0, 1].
          lowerCAmelCase__   :Union[str, Any]      							=				(features - self.min_value) / (self.max_value - self.min_value)
          # Scale to [min_out, max_out].
          return zero_one * (max_out - min_out) + min_out
      def 			snake_case_      (   self    ,	_lowerCAmelCase    ,	_lowerCAmelCase=(-1.0, 1.0)    ,	_lowerCAmelCase=False		):



          '''simple docstring'''
          lowerCAmelCase__							,lowerCAmelCase__   :Any      							=				input_range
          lowerCAmelCase__   :Dict      							=				torch.clip(_lowerCAmelCase    ,	_lowerCAmelCase    ,	_lowerCAmelCase		) if clip else outputs
          # Scale to [0, 1].
          lowerCAmelCase__   :Optional[int]      							=				(outputs - min_out) / (max_out - min_out)
          # Scale to [self.min_value, self.max_value].
          return zero_one * (self.max_value - self.min_value) + self.min_value
      def 			snake_case_      (   self    ,	_lowerCAmelCase    ,	_lowerCAmelCase    ,	_lowerCAmelCase		):



          '''simple docstring'''
          lowerCAmelCase__   :List[Any]      							=				input_tokens > 0
          lowerCAmelCase__							,lowerCAmelCase__   :str      							=				self.notes_encoder(
              encoder_input_tokens=_lowerCAmelCase    ,	encoder_inputs_mask=_lowerCAmelCase		)

          lowerCAmelCase__							,lowerCAmelCase__   :Any      							=				self.continuous_encoder(
              encoder_inputs=_lowerCAmelCase    ,	encoder_inputs_mask=_lowerCAmelCase		)

          return [(tokens_encoded, tokens_mask), (continuous_encoded, continuous_mask)]
      def 			snake_case_      (   self    ,	_lowerCAmelCase    ,	_lowerCAmelCase    ,	_lowerCAmelCase		):



          '''simple docstring'''
          lowerCAmelCase__   :Dict      							=				noise_time
          if not torch.is_tensor(_lowerCAmelCase		):
              lowerCAmelCase__   :Dict      							=				torch.tensor([timesteps]    ,	dtype=torch.long    ,	device=input_tokens.device		)
          elif torch.is_tensor(_lowerCAmelCase		) and len(timesteps.shape		) == 0:
              lowerCAmelCase__   :Union[str, Any]      							=				timesteps[None].to(input_tokens.device		)

          # broadcast to batch dimension in a way that's compatible with ONNX/Core ML
          lowerCAmelCase__   :List[str]      							=				timesteps * torch.ones(input_tokens.shape[0]    ,	dtype=timesteps.dtype    ,	device=timesteps.device		)

          lowerCAmelCase__   :Dict      							=				self.decoder(
              encodings_and_masks=_lowerCAmelCase    ,	decoder_input_tokens=_lowerCAmelCase    ,	decoder_noise_time=_lowerCAmelCase		)
          return logits



      @torch.no_grad()
      def __call__(   self    ,	_lowerCAmelCase    ,	_lowerCAmelCase = None    ,	_lowerCAmelCase = 100    ,	_lowerCAmelCase = True    ,	_lowerCAmelCase = "numpy"    ,	_lowerCAmelCase = None    ,	_lowerCAmelCase = 1    ,	):



          '''simple docstring'''

          if (callback_steps is None) or (
              callback_steps is not None and (not isinstance(_lowerCAmelCase    ,	_lowerCAmelCase		) or callback_steps <= 0)
          ):
              raise ValueError(
                  f'''`callback_steps` has to be a positive integer but is {callback_steps} of type'''
                  f''' {type(_lowerCAmelCase		)}.'''		)

          lowerCAmelCase__   :Any      							=				np.zeros([1, TARGET_FEATURE_LENGTH, self.n_dims]    ,	dtype=np.floataa		)
          lowerCAmelCase__   :Any      							=				np.zeros([1, 0, self.n_dims]    ,	np.floataa		)
          lowerCAmelCase__   :Any      							=				torch.ones((1, TARGET_FEATURE_LENGTH)    ,	dtype=_lowerCAmelCase    ,	device=self.device		)

          for i, encoder_input_tokens in enumerate(_lowerCAmelCase		):
              if i == 0:
                  lowerCAmelCase__   :str      							=				torch.from_numpy(pred_mel[:1].copy()		).to(
                      device=self.device    ,	dtype=self.decoder.dtype		)
                  # The first chunk has no previous context.
                  lowerCAmelCase__   :Union[str, Any]      							=				torch.zeros((1, TARGET_FEATURE_LENGTH)    ,	dtype=_lowerCAmelCase    ,	device=self.device		)
              else:
                  # The full song pipeline does not feed in a context feature, so the mask
                  # will be all 0s after the feature converter. Because we know we're
                  # feeding in a full context chunk from the previous prediction, set it
                  # to all 1s.
                  lowerCAmelCase__   :Optional[int]      							=				ones

              lowerCAmelCase__   :Optional[Any]      							=				self.scale_features(
                  _lowerCAmelCase    ,	output_range=[-1.0, 1.0]    ,	clip=_lowerCAmelCase		)

              lowerCAmelCase__   :Tuple      							=				self.encode(
                  input_tokens=torch.IntTensor([encoder_input_tokens]		).to(device=self.device		)    ,	continuous_inputs=_lowerCAmelCase    ,	continuous_mask=_lowerCAmelCase    ,	)

              # Sample encoder_continuous_inputs shaped gaussian noise to begin loop
              lowerCAmelCase__   :Any      							=				randn_tensor(
                  shape=encoder_continuous_inputs.shape    ,	generator=_lowerCAmelCase    ,	device=self.device    ,	dtype=self.decoder.dtype    ,	)

              # set step values
              self.scheduler.set_timesteps(_lowerCAmelCase		)

              # Denoising diffusion loop
              for j, t in enumerate(self.progress_bar(self.scheduler.timesteps		)		):
                  lowerCAmelCase__   :Optional[int]      							=				self.decode(
                      encodings_and_masks=_lowerCAmelCase    ,	input_tokens=_lowerCAmelCase    ,	noise_time=t / self.scheduler.config.num_train_timesteps    ,	)

                  # Compute previous output: x_t -> x_t-1
                  lowerCAmelCase__   :Any      							=				self.scheduler.step(_lowerCAmelCase    ,	_lowerCAmelCase    ,	_lowerCAmelCase    ,	generator=_lowerCAmelCase		).prev_sample

              lowerCAmelCase__   :int      							=				self.scale_to_features(_lowerCAmelCase    ,	input_range=[-1.0, 1.0]		)
              lowerCAmelCase__   :List[Any]      							=				mel[:1]
              lowerCAmelCase__   :Optional[Any]      							=				mel.cpu().float().numpy()

              lowerCAmelCase__   :Any      							=				np.concatenate([full_pred_mel, pred_mel[:1]]    ,	axis=1		)

              # call the callback, if provided
              if callback is not None and i % callback_steps == 0:
                  callback(_lowerCAmelCase    ,	_lowerCAmelCase		)

              logger.info("Generated segment"    ,	_lowerCAmelCase		)

          if output_type == "numpy" and not is_onnx_available():
              raise ValueError(
                  "Cannot return output in \'np\' format if ONNX is not available. Make sure to have ONNX installed or set \'output_type\' to \'mel\'."		)
          elif output_type == "numpy" and self.melgan is None:
              raise ValueError(
                  "Cannot return output in \'np\' format if melgan component is not defined. Make sure to define `self.melgan` or set \'output_type\' to \'mel\'."		)

          if output_type == "numpy":
              lowerCAmelCase__   :Any      							=				self.melgan(input_features=full_pred_mel.astype(np.floataa		)		)
          else:
              lowerCAmelCase__   :int      							=				full_pred_mel

          if not return_dict:
              return (output,)

          return AudioPipelineOutput(audios=_lowerCAmelCase		)













import inspect
import unittest

from huggingface_hub import hf_hub_download

from transformers import ConvNextConfig, UperNetConfig
from transformers.testing_utils import require_torch, require_torch_multi_gpu, require_vision, slow, torch_device
from transformers.utils import is_torch_available, is_vision_available

from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, _config_zero_init, floats_tensor, ids_tensor
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
 import torch

 from transformers import UperNetForSemanticSegmentation
 from transformers.models.upernet.modeling_upernet import UPERNET_PRETRAINED_MODEL_ARCHIVE_LIST


if is_vision_available():
 from PIL import Image

 from transformers import AutoImageProcessor







class 	_UpperCAmelCase  :




      """simple docstring"""



      def __init__(   self    ,	_lowerCAmelCase    ,	_lowerCAmelCase=13    ,	_lowerCAmelCase=32    ,	_lowerCAmelCase=3    ,	_lowerCAmelCase=4    ,	_lowerCAmelCase=[10, 20, 30, 40]    ,	_lowerCAmelCase=[2, 2, 3, 2]    ,	_lowerCAmelCase=True    ,	_lowerCAmelCase=True    ,	_lowerCAmelCase=37    ,	_lowerCAmelCase="gelu"    ,	_lowerCAmelCase=10    ,	_lowerCAmelCase=0.02    ,	_lowerCAmelCase=["stage2", "stage3", "stage4"]    ,	_lowerCAmelCase=3    ,	_lowerCAmelCase=None    ,	):



          '''simple docstring'''
          lowerCAmelCase__   :str      							=				parent
          lowerCAmelCase__   :Dict      							=				batch_size
          lowerCAmelCase__   :str      							=				image_size
          lowerCAmelCase__   :int      							=				num_channels
          lowerCAmelCase__   :Tuple      							=				num_stages
          lowerCAmelCase__   :Optional[int]      							=				hidden_sizes
          lowerCAmelCase__   :Tuple      							=				depths
          lowerCAmelCase__   :Tuple      							=				is_training
          lowerCAmelCase__   :str      							=				use_labels
          lowerCAmelCase__   :int      							=				intermediate_size
          lowerCAmelCase__   :List[str]      							=				hidden_act
          lowerCAmelCase__   :Dict      							=				type_sequence_label_size
          lowerCAmelCase__   :Any      							=				initializer_range
          lowerCAmelCase__   :List[str]      							=				out_features
          lowerCAmelCase__   :List[Any]      							=				num_labels
          lowerCAmelCase__   :str      							=				scope
          lowerCAmelCase__   :Dict      							=				num_stages
      def 			snake_case_      (   self		):



          '''simple docstring'''
          lowerCAmelCase__   :int      							=				floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]		)

          lowerCAmelCase__   :List[str]      							=				None
          if self.use_labels:
              lowerCAmelCase__   :str      							=				ids_tensor([self.batch_size]    ,	self.type_sequence_label_size		)

          lowerCAmelCase__   :Optional[Any]      							=				self.get_config()

          return config, pixel_values, labels
      def 			snake_case_      (   self		):



          '''simple docstring'''
          return ConvNextConfig(
              num_channels=self.num_channels    ,	num_stages=self.num_stages    ,	hidden_sizes=self.hidden_sizes    ,	depths=self.depths    ,	is_training=self.is_training    ,	intermediate_size=self.intermediate_size    ,	hidden_act=self.hidden_act    ,	out_features=self.out_features    ,	)
      def 			snake_case_      (   self		):



          '''simple docstring'''
          return UperNetConfig(
              backbone_config=self.get_backbone_config()    ,	hidden_size=512    ,	pool_scales=[1, 2, 3, 6]    ,	use_auxiliary_head=_lowerCAmelCase    ,	auxiliary_loss_weight=0.4    ,	auxiliary_in_channels=40    ,	auxiliary_channels=256    ,	auxiliary_num_convs=1    ,	auxiliary_concat_input=_lowerCAmelCase    ,	loss_ignore_index=255    ,	num_labels=self.num_labels    ,	)
      def 			snake_case_      (   self    ,	_lowerCAmelCase    ,	_lowerCAmelCase    ,	_lowerCAmelCase		):



          '''simple docstring'''
          lowerCAmelCase__   :Dict      							=				UperNetForSemanticSegmentation(config=_lowerCAmelCase		)
          model.to(_lowerCAmelCase		)
          model.eval()
          lowerCAmelCase__   :List[Any]      							=				model(_lowerCAmelCase		)
          self.parent.assertEqual(
              result.logits.shape    ,	(self.batch_size, self.num_labels, self.image_size, self.image_size)		)



      def 			snake_case_      (   self		):



          '''simple docstring'''
          lowerCAmelCase__   :Optional[int]      							=				self.prepare_config_and_inputs()
          (
              (
              lowerCAmelCase__
          )							,(
              lowerCAmelCase__
          )							,(
              lowerCAmelCase__
          )							,
          )   :Dict      							=				config_and_inputs
          lowerCAmelCase__   :Any      							=				{"pixel_values": pixel_values}
          return config, inputs_dict







@require_torch
class 	_UpperCAmelCase  (  _A				,							_A				,							unittest.TestCase		):




      """simple docstring"""



      A          =    (UperNetForSemanticSegmentation,) if is_torch_available() else ()
      A          =    {'''image-segmentation''': UperNetForSemanticSegmentation} if is_torch_available() else {}
      A          =    False
      A          =    False
      A          =    False
      A          =    False
      A          =    False
      A          =    False
      def 			snake_case_      (   self		):



          '''simple docstring'''
          lowerCAmelCase__   :Tuple      							=				UperNetModelTester(self		)
          lowerCAmelCase__   :Any      							=				ConfigTester(self    ,	config_class=_lowerCAmelCase    ,	has_text_modality=_lowerCAmelCase    ,	hidden_size=37		)
      def 			snake_case_      (   self		):



          '''simple docstring'''
          self.create_and_test_config_common_properties()
          self.config_tester.create_and_test_config_to_json_string()
          self.config_tester.create_and_test_config_to_json_file()
          self.config_tester.create_and_test_config_from_and_save_pretrained()
          self.config_tester.create_and_test_config_with_num_labels()
          self.config_tester.check_config_can_be_init_without_params()
          self.config_tester.check_config_arguments_init()
      def 			snake_case_      (   self		):



          '''simple docstring'''
          return
      def 			snake_case_      (   self		):



          '''simple docstring'''
          lowerCAmelCase__							,lowerCAmelCase__   :str      							=				self.model_tester.prepare_config_and_inputs_for_common()

          for model_class in self.all_model_classes:
              lowerCAmelCase__   :Optional[int]      							=				model_class(_lowerCAmelCase		)
              lowerCAmelCase__   :Tuple      							=				inspect.signature(model.forward		)
              # signature.parameters is an OrderedDict => so arg_names order is deterministic
              lowerCAmelCase__   :Optional[Any]      							=				[*signature.parameters.keys()]

              lowerCAmelCase__   :Union[str, Any]      							=				["pixel_values"]
              self.assertListEqual(arg_names[:1]    ,	_lowerCAmelCase		)
      def 			snake_case_      (   self		):



          '''simple docstring'''
          lowerCAmelCase__   :int      							=				self.model_tester.prepare_config_and_inputs()
          self.model_tester.create_and_check_for_semantic_segmentation(*_lowerCAmelCase		)
      @unittest.skip(reason="UperNet does not use inputs_embeds"		)
      def 			snake_case_      (   self		):



          '''simple docstring'''
          pass
      @unittest.skip(reason="UperNet does not support input and output embeddings"		)
      def 			snake_case_      (   self		):



          '''simple docstring'''
          pass
      @unittest.skip(reason="UperNet does not have a base model"		)
      def 			snake_case_      (   self		):



          '''simple docstring'''
          pass
      @unittest.skip(reason="UperNet does not have a base model"		)
      def 			snake_case_      (   self		):



          '''simple docstring'''
          pass
      @require_torch_multi_gpu
      @unittest.skip(reason="UperNet has some layers using `add_module` which doesn't work well with `nn.DataParallel`"		)
      def 			snake_case_      (   self		):



          '''simple docstring'''
          pass
      @unittest.skip("Will be fixed soon by reducing the size of the model used for common tests."		)
      def 			snake_case_      (   self		):



          '''simple docstring'''
          pass
      def 			snake_case_      (   self		):



          '''simple docstring'''
          def check_hidden_states_output(_lowerCAmelCase    ,	_lowerCAmelCase    ,	_lowerCAmelCase		):
              lowerCAmelCase__   :int      							=				model_class(_lowerCAmelCase		)
              model.to(_lowerCAmelCase		)
              model.eval()

              with torch.no_grad():
                  lowerCAmelCase__   :Dict      							=				model(**self._prepare_for_class(_lowerCAmelCase    ,	_lowerCAmelCase		)		)

              lowerCAmelCase__   :int      							=				outputs.encoder_hidden_states if config.is_encoder_decoder else outputs.hidden_states

              lowerCAmelCase__   :Tuple      							=				self.model_tester.num_stages
              self.assertEqual(len(_lowerCAmelCase		)    ,	expected_num_stages + 1		)

              # ConvNext's feature maps are of shape (batch_size, num_channels, height, width)
              self.assertListEqual(
                  list(hidden_states[0].shape[-2:]		)    ,	[self.model_tester.image_size // 4, self.model_tester.image_size // 4]    ,	)

          lowerCAmelCase__							,lowerCAmelCase__   :Any      							=				self.model_tester.prepare_config_and_inputs_for_common()

          for model_class in self.all_model_classes:
              lowerCAmelCase__   :str      							=				True
              check_hidden_states_output(_lowerCAmelCase    ,	_lowerCAmelCase    ,	_lowerCAmelCase		)

              # check that output_hidden_states also work using config
              del inputs_dict["output_hidden_states"]
              lowerCAmelCase__   :str      							=				True

              check_hidden_states_output(_lowerCAmelCase    ,	_lowerCAmelCase    ,	_lowerCAmelCase		)
      def 			snake_case_      (   self		):



          '''simple docstring'''
          lowerCAmelCase__							,lowerCAmelCase__   :Dict      							=				self.model_tester.prepare_config_and_inputs_for_common()

          lowerCAmelCase__   :str      							=				_config_zero_init(_lowerCAmelCase		)
          lowerCAmelCase__   :List[Any]      							=				_config_zero_init(configs_no_init.backbone_config		)
          for model_class in self.all_model_classes:
              lowerCAmelCase__   :List[Any]      							=				model_class(config=_lowerCAmelCase		)
              for name, param in model.named_parameters():
                  if param.requires_grad:
                      self.assertIn(
                          ((param.data.mean() * 1e9).round() / 1e9).item()    ,	[0.0, 1.0]    ,	msg=f'''Parameter {name} of model {model_class} seems not properly initialized'''    ,	)
      @unittest.skip(reason="UperNet does not have tied weights"		)
      def 			snake_case_      (   self		):



          '''simple docstring'''
          pass



      @slow
      def 			snake_case_      (   self		):



          '''simple docstring'''
          for model_name in UPERNET_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
              lowerCAmelCase__   :List[str]      							=				UperNetForSemanticSegmentation.from_pretrained(_lowerCAmelCase		)
              self.assertIsNotNone(_lowerCAmelCase		)





def 				snake_case__							(							):
    lowerCAmelCase__   :int      							=				hf_hub_download(
        repo_id="hf-internal-testing/fixtures_ade20k"    ,			repo_type="dataset"    ,			filename="ADE_val_00000001.jpg"       )
    lowerCAmelCase__   :Union[str, Any]      							=				Image.open(UpperCAmelCase       ).convert("RGB"       )
    return image







@require_torch
@require_vision
@slow
class 	_UpperCAmelCase  (  unittest.TestCase		):




      """simple docstring"""



      def 			snake_case_      (   self		):



          '''simple docstring'''
          lowerCAmelCase__   :Optional[Any]      							=				AutoImageProcessor.from_pretrained("openmmlab/upernet-swin-tiny"		)
          lowerCAmelCase__   :Dict      							=				UperNetForSemanticSegmentation.from_pretrained("openmmlab/upernet-swin-tiny"		).to(_lowerCAmelCase		)

          lowerCAmelCase__   :str      							=				prepare_img()
          lowerCAmelCase__   :str      							=				processor(images=_lowerCAmelCase    ,	return_tensors="pt"		).to(_lowerCAmelCase		)

          with torch.no_grad():
              lowerCAmelCase__   :int      							=				model(**_lowerCAmelCase		)

          lowerCAmelCase__   :str      							=				torch.Size((1, model.config.num_labels, 512, 512)		)
          self.assertEqual(outputs.logits.shape    ,	_lowerCAmelCase		)

          lowerCAmelCase__   :Tuple      							=				torch.tensor(
              [[-7.5958, -7.5958, -7.4302], [-7.5958, -7.5958, -7.4302], [-7.4797, -7.4797, -7.3068]]		).to(_lowerCAmelCase		)
          self.assertTrue(torch.allclose(outputs.logits[0, 0, :3, :3]    ,	_lowerCAmelCase    ,	atol=1e-4		)		)



      def 			snake_case_      (   self		):



          '''simple docstring'''
          lowerCAmelCase__   :Any      							=				AutoImageProcessor.from_pretrained("openmmlab/upernet-convnext-tiny"		)
          lowerCAmelCase__   :Optional[int]      							=				UperNetForSemanticSegmentation.from_pretrained("openmmlab/upernet-convnext-tiny"		).to(_lowerCAmelCase		)

          lowerCAmelCase__   :List[Any]      							=				prepare_img()
          lowerCAmelCase__   :Tuple      							=				processor(images=_lowerCAmelCase    ,	return_tensors="pt"		).to(_lowerCAmelCase		)

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

          lowerCAmelCase__   :str      							=				torch.Size((1, model.config.num_labels, 512, 512)		)
          self.assertEqual(outputs.logits.shape    ,	_lowerCAmelCase		)

          lowerCAmelCase__   :str      							=				torch.tensor(
              [[-8.8110, -8.8110, -8.6521], [-8.8110, -8.8110, -8.6521], [-8.7746, -8.7746, -8.6130]]		).to(_lowerCAmelCase		)
          self.assertTrue(torch.allclose(outputs.logits[0, 0, :3, :3]    ,	_lowerCAmelCase    ,	atol=1e-4		)		)



