



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






	"""simple docstring"""


	def __init__(     self     :						int						,  __UpperCAmelCase     :						str						,  __UpperCAmelCase     :						List[Any]=13						,  __UpperCAmelCase     :						Optional[int]=32						,  __UpperCAmelCase     :						Any=3						,  __UpperCAmelCase     :						int=4						,  __UpperCAmelCase     :						Optional[int]=[10, 20, 30, 40]						,  __UpperCAmelCase     :						List[Any]=[2, 2, 3, 2]						,  __UpperCAmelCase     :						List[Any]=True						,  __UpperCAmelCase     :						Optional[Any]=True						,  __UpperCAmelCase     :						int=37						,  __UpperCAmelCase     :						Union[str, Any]="gelu"						,  __UpperCAmelCase     :						Tuple=10						,  __UpperCAmelCase     :						int=0.02						,  __UpperCAmelCase     :						Any=["stage2", "stage3", "stage4"]						,  __UpperCAmelCase     :						Optional[Any]=3						,  __UpperCAmelCase     :						Tuple=None						,  )      ->			Optional[int]:
							UpperCAmelCase_=       parent
							UpperCAmelCase_=       batch_size
							UpperCAmelCase_=       image_size
							UpperCAmelCase_=       num_channels
							UpperCAmelCase_=       num_stages
							UpperCAmelCase_=       hidden_sizes
							UpperCAmelCase_=       depths
							UpperCAmelCase_=       is_training
							UpperCAmelCase_=       use_labels
							UpperCAmelCase_=       intermediate_size
							UpperCAmelCase_=       hidden_act
							UpperCAmelCase_=       type_sequence_label_size
							UpperCAmelCase_=       initializer_range
							UpperCAmelCase_=       out_features
							UpperCAmelCase_=       num_labels
							UpperCAmelCase_=       scope
							UpperCAmelCase_=       num_stages




	def   _SCREAMING_SNAKE_CASE		(     self     :						str					)      ->			List[Any]:
							UpperCAmelCase_=       floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]					)

							UpperCAmelCase_=       None
							if self.use_labels:
													UpperCAmelCase_=       ids_tensor([self.batch_size]						,  self.type_sequence_label_size					)

							UpperCAmelCase_=       self.get_config()

							return config, pixel_values, labels




	def   _SCREAMING_SNAKE_CASE		(     self     :						str					)      ->			Union[str, Any]:
							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   _SCREAMING_SNAKE_CASE		(     self     :						List[Any]					)      ->			int:
							return UperNetConfig(
							    backbone_config=self.get_backbone_config()						,  hidden_size=512						,  pool_scales=[1, 2, 3, 6]						,  use_auxiliary_head=lowercase__						,  auxiliary_loss_weight=0.4						,  auxiliary_in_channels=40						,  auxiliary_channels=256						,  auxiliary_num_convs=1						,  auxiliary_concat_input=lowercase__						,  loss_ignore_index=255						,  num_labels=self.num_labels						,  )




	def   _SCREAMING_SNAKE_CASE		(     self     :						Any						,  __UpperCAmelCase     :						List[str]						,  __UpperCAmelCase     :						Any						,  __UpperCAmelCase     :						Any					)      ->			Any:
							UpperCAmelCase_=       UperNetForSemanticSegmentation(config=lowercase__					)
							model.to(lowercase__					)
							model.eval()
							UpperCAmelCase_=       model(lowercase__					)
							self.parent.assertEqual(
							    result.logits.shape						,  (self.batch_size, self.num_labels, self.image_size, self.image_size)					)




	def   _SCREAMING_SNAKE_CASE		(     self     :						Union[str, Any]					)      ->			Tuple:
							UpperCAmelCase_=       self.prepare_config_and_inputs()
							(
							    (
							    UpperCAmelCase_     
							),       (
							    UpperCAmelCase_     
							),       (
							    UpperCAmelCase_     
							),            
							)=       config_and_inputs
							UpperCAmelCase_=       {"""pixel_values""": pixel_values}
							return config, inputs_dict






@require_torch
class 		lowercase						(			lowerCamelCase__			,   lowerCamelCase__			,   unittest.TestCase):






	"""simple docstring"""


	a__				: Tuple		       =							(UperNetForSemanticSegmentation,) if is_torch_available() else ()
	a__				: Tuple		       =							{'image-segmentation': UperNetForSemanticSegmentation} if is_torch_available() else {}
	a__				: Tuple		       =							False
	a__				: Dict		       =							False
	a__				: Any		       =							False
	a__				: Optional[Any]		       =							False
	a__				: List[Any]		       =							False
	a__				: List[Any]		       =							False




	def   _SCREAMING_SNAKE_CASE		(     self     :						Union[str, Any]					)      ->			Optional[Any]:
							UpperCAmelCase_=       UperNetModelTester(self					)
							UpperCAmelCase_=       ConfigTester(self						,  config_class=lowercase__						,  has_text_modality=lowercase__						,  hidden_size=37					)




	def   _SCREAMING_SNAKE_CASE		(     self     :						Union[str, Any]					)      ->			str:
							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   _SCREAMING_SNAKE_CASE		(     self     :						Optional[int]					)      ->			List[Any]:
							return




	def   _SCREAMING_SNAKE_CASE		(     self     :						Optional[Any]					)      ->			Union[str, Any]:
							UpperCAmelCase_,       UpperCAmelCase_=       self.model_tester.prepare_config_and_inputs_for_common()

							for model_class in self.all_model_classes:
													UpperCAmelCase_=       model_class(lowercase__					)
													UpperCAmelCase_=       inspect.signature(model.forward					)
													# signature.parameters is an OrderedDict => so arg_names order is deterministic
													UpperCAmelCase_=       [*signature.parameters.keys()]

													UpperCAmelCase_=       ["""pixel_values"""]
													self.assertListEqual(arg_names[:1]						,  lowercase__					)




	def   _SCREAMING_SNAKE_CASE		(     self     :						Optional[Any]					)      ->			Dict:
							UpperCAmelCase_=       self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_for_semantic_segmentation(*lowercase__					)




	@unittest.skip(reason="""UperNet does not use inputs_embeds"""					)
	def   _SCREAMING_SNAKE_CASE		(     self     :						int					)      ->			List[Any]:
							pass




	@unittest.skip(reason="""UperNet does not support input and output embeddings"""					)
	def   _SCREAMING_SNAKE_CASE		(     self     :						str					)      ->			Optional[int]:
							pass




	@unittest.skip(reason="""UperNet does not have a base model"""					)
	def   _SCREAMING_SNAKE_CASE		(     self     :						Optional[Any]					)      ->			List[str]:
							pass




	@unittest.skip(reason="""UperNet does not have a base model"""					)
	def   _SCREAMING_SNAKE_CASE		(     self     :						Union[str, Any]					)      ->			Tuple:
							pass




	@require_torch_multi_gpu
	@unittest.skip(reason="""UperNet has some layers using `add_module` which doesn\'t work well with `nn.DataParallel`"""					)
	def   _SCREAMING_SNAKE_CASE		(     self     :						Any					)      ->			Any:
							pass




	@unittest.skip("""Will be fixed soon by reducing the size of the model used for common tests."""					)
	def   _SCREAMING_SNAKE_CASE		(     self     :						List[Any]					)      ->			Union[str, Any]:
							pass




	def   _SCREAMING_SNAKE_CASE		(     self     :						Any					)      ->			List[Any]:
							def check_hidden_states_output(__UpperCAmelCase     :						List[str]						,  __UpperCAmelCase     :						int						,  __UpperCAmelCase     :						List[str]					):
													UpperCAmelCase_=       model_class(lowercase__					)
													model.to(lowercase__					)
													model.eval()

													with torch.no_grad():
																			UpperCAmelCase_=       model(**self._prepare_for_class(lowercase__						,  lowercase__					)					)

													UpperCAmelCase_=       outputs.encoder_hidden_states if config.is_encoder_decoder else outputs.hidden_states

													UpperCAmelCase_=       self.model_tester.num_stages
													self.assertEqual(len(lowercase__					)						,  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]						,  )

							UpperCAmelCase_,       UpperCAmelCase_=       self.model_tester.prepare_config_and_inputs_for_common()

							for model_class in self.all_model_classes:
													UpperCAmelCase_=       True
													check_hidden_states_output(lowercase__						,  lowercase__						,  lowercase__					)

													# check that output_hidden_states also work using config
													del inputs_dict["output_hidden_states"]
													UpperCAmelCase_=       True

													check_hidden_states_output(lowercase__						,  lowercase__						,  lowercase__					)




	def   _SCREAMING_SNAKE_CASE		(     self     :						Dict					)      ->			Optional[int]:
							UpperCAmelCase_,       UpperCAmelCase_=       self.model_tester.prepare_config_and_inputs_for_common()

							UpperCAmelCase_=       _config_zero_init(lowercase__					)
							UpperCAmelCase_=       _config_zero_init(configs_no_init.backbone_config					)
							for model_class in self.all_model_classes:
													UpperCAmelCase_=       model_class(config=lowercase__					)
													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   _SCREAMING_SNAKE_CASE		(     self     :						Optional[Any]					)      ->			Any:
							pass




	@slow
	def   _SCREAMING_SNAKE_CASE		(     self     :						Optional[Any]					)      ->			Any:
							for model_name in UPERNET_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
													UpperCAmelCase_=       UperNetForSemanticSegmentation.from_pretrained(lowercase__					)
													self.assertIsNotNone(lowercase__					)


def 	__a					(   )				->			int:






						'''simple docstring'''



						UpperCAmelCase_=       hf_hub_download(
						    repo_id="""hf-internal-testing/fixtures_ade20k"""      ,repo_type="""dataset"""      ,filename="""ADE_val_00000001.jpg"""    )
						UpperCAmelCase_=       Image.open(A__    ).convert("""RGB"""    )
						return image






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






	"""simple docstring"""


	def   _SCREAMING_SNAKE_CASE		(     self     :						Optional[Any]					)      ->			Tuple:
							UpperCAmelCase_=       AutoImageProcessor.from_pretrained("""openmmlab/upernet-swin-tiny"""					)
							UpperCAmelCase_=       UperNetForSemanticSegmentation.from_pretrained("""openmmlab/upernet-swin-tiny"""					).to(lowercase__					)

							UpperCAmelCase_=       prepare_img()
							UpperCAmelCase_=       processor(images=lowercase__						,  return_tensors="""pt"""					).to(lowercase__					)

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

							UpperCAmelCase_=       torch.Size((1, model.config.num_labels, 512, 512)					)
							self.assertEqual(outputs.logits.shape						,  lowercase__					)

							UpperCAmelCase_=       torch.tensor(
							    [[-7.5_958, -7.5_958, -7.4_302], [-7.5_958, -7.5_958, -7.4_302], [-7.4_797, -7.4_797, -7.3_068]]					).to(lowercase__					)
							self.assertTrue(torch.allclose(outputs.logits[0, 0, :3, :3]						,  lowercase__						,  atol=1E-4					)					)




	def   _SCREAMING_SNAKE_CASE		(     self     :						Tuple					)      ->			Union[str, Any]:
							UpperCAmelCase_=       AutoImageProcessor.from_pretrained("""openmmlab/upernet-convnext-tiny"""					)
							UpperCAmelCase_=       UperNetForSemanticSegmentation.from_pretrained("""openmmlab/upernet-convnext-tiny"""					).to(lowercase__					)

							UpperCAmelCase_=       prepare_img()
							UpperCAmelCase_=       processor(images=lowercase__						,  return_tensors="""pt"""					).to(lowercase__					)

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

							UpperCAmelCase_=       torch.Size((1, model.config.num_labels, 512, 512)					)
							self.assertEqual(outputs.logits.shape						,  lowercase__					)

							UpperCAmelCase_=       torch.tensor(
							    [[-8.8_110, -8.8_110, -8.6_521], [-8.8_110, -8.8_110, -8.6_521], [-8.7_746, -8.7_746, -8.6_130]]					).to(lowercase__					)
							self.assertTrue(torch.allclose(outputs.logits[0, 0, :3, :3]						,  lowercase__						,  atol=1E-4					)					)









from math import acos, sin
from typing import List, Tuple, Union

import numpy as np
import torch
from PIL import Image

from ...models import AutoencoderKL, UNetaDConditionModel
from ...schedulers import DDIMScheduler, DDPMScheduler
from ...utils import randn_tensor
from ..pipeline_utils import AudioPipelineOutput, BaseOutput, DiffusionPipeline, ImagePipelineOutput
from .mel import Mel






class 		lowercase						(			snake_case__):






	"""simple docstring"""


	a__				: str		       =							["vqvae"]




	def __init__(     self     :						List[Any]						,  __UpperCAmelCase     :						AutoencoderKL						,  __UpperCAmelCase     :						UNetaDConditionModel						,  __UpperCAmelCase     :						Mel						,  __UpperCAmelCase     :						Union[DDIMScheduler, DDPMScheduler]						,  )      ->			str:
							super().__init__()
							self.register_modules(unet=__UpperCAmelCase						,  scheduler=__UpperCAmelCase						,  mel=__UpperCAmelCase						,  vqvae=__UpperCAmelCase					)




	def   _SCREAMING_SNAKE_CASE		(     self     :						List[Any]					)      ->			int:
							return 50 if isinstance(self.scheduler						,  __UpperCAmelCase					) else 1_000




	@torch.no_grad()
	def __call__(     self     :						List[Any]						,  __UpperCAmelCase     :						int = 1						,  __UpperCAmelCase     :						str = None						,  __UpperCAmelCase     :						np.ndarray = None						,  __UpperCAmelCase     :						int = 0						,  __UpperCAmelCase     :						int = 0						,  __UpperCAmelCase     :						int = None						,  __UpperCAmelCase     :						torch.Generator = None						,  __UpperCAmelCase     :						float = 0						,  __UpperCAmelCase     :						float = 0						,  __UpperCAmelCase     :						torch.Generator = None						,  __UpperCAmelCase     :						float = 0						,  __UpperCAmelCase     :						torch.Tensor = None						,  __UpperCAmelCase     :						torch.Tensor = None						,  __UpperCAmelCase     :						Union[str, Any]=True						,  )      ->			Union[
	    Union[AudioPipelineOutput, ImagePipelineOutput],
	    Tuple[List[Image.Image], Tuple[int, List[np.ndarray]]],
	]:
							UpperCAmelCase_=       steps or self.get_default_steps()
							self.scheduler.set_timesteps(__UpperCAmelCase					)
							UpperCAmelCase_=       step_generator or generator
							# For backwards compatibility
							if type(self.unet.config.sample_size					) == int:
													UpperCAmelCase_=       (self.unet.config.sample_size, self.unet.config.sample_size)
							if noise is None:
													UpperCAmelCase_=       randn_tensor(
													    (
													        batch_size,
													        self.unet.config.in_channels,
													        self.unet.config.sample_size[0],
													        self.unet.config.sample_size[1],
													    )						,  generator=__UpperCAmelCase						,  device=self.device						,  )
							UpperCAmelCase_=       noise
							UpperCAmelCase_=       None

							if audio_file is not None or raw_audio is not None:
													self.mel.load_audio(__UpperCAmelCase						,  __UpperCAmelCase					)
													UpperCAmelCase_=       self.mel.audio_slice_to_image(__UpperCAmelCase					)
													UpperCAmelCase_=       np.frombuffer(input_image.tobytes()						,  dtype="""uint8"""					).reshape(
													    (input_image.height, input_image.width)					)
													UpperCAmelCase_=       (input_image / 255) * 2 - 1
													UpperCAmelCase_=       torch.tensor(input_image[np.newaxis, :, :]						,  dtype=torch.float					).to(self.device					)

													if self.vqvae is not None:
																			UpperCAmelCase_=       self.vqvae.encode(torch.unsqueeze(__UpperCAmelCase						,  0					)					).latent_dist.sample(
																			    generator=__UpperCAmelCase					)[0]
																			UpperCAmelCase_=       self.vqvae.config.scaling_factor * input_images

													if start_step > 0:
																			UpperCAmelCase_=       self.scheduler.add_noise(__UpperCAmelCase						,  __UpperCAmelCase						,  self.scheduler.timesteps[start_step - 1]					)

													UpperCAmelCase_=       (
													    self.unet.config.sample_size[1] * self.mel.get_sample_rate() / self.mel.x_res / self.mel.hop_length
													)
													UpperCAmelCase_=       int(mask_start_secs * pixels_per_second					)
													UpperCAmelCase_=       int(mask_end_secs * pixels_per_second					)
													UpperCAmelCase_=       self.scheduler.add_noise(__UpperCAmelCase						,  __UpperCAmelCase						,  torch.tensor(self.scheduler.timesteps[start_step:]					)					)

							for step, t in enumerate(self.progress_bar(self.scheduler.timesteps[start_step:]					)					):
													if isinstance(self.unet						,  __UpperCAmelCase					):
																			UpperCAmelCase_=       self.unet(__UpperCAmelCase						,  __UpperCAmelCase						,  __UpperCAmelCase					)["""sample"""]
													else:
																			UpperCAmelCase_=       self.unet(__UpperCAmelCase						,  __UpperCAmelCase					)["""sample"""]

													if isinstance(self.scheduler						,  __UpperCAmelCase					):
																			UpperCAmelCase_=       self.scheduler.step(
																			    model_output=__UpperCAmelCase						,  timestep=__UpperCAmelCase						,  sample=__UpperCAmelCase						,  eta=__UpperCAmelCase						,  generator=__UpperCAmelCase						,  )["""prev_sample"""]
													else:
																			UpperCAmelCase_=       self.scheduler.step(
																			    model_output=__UpperCAmelCase						,  timestep=__UpperCAmelCase						,  sample=__UpperCAmelCase						,  generator=__UpperCAmelCase						,  )["""prev_sample"""]

													if mask is not None:
																			if mask_start > 0:
																									UpperCAmelCase_=       mask[:, step, :, :mask_start]
																			if mask_end > 0:
																									UpperCAmelCase_=       mask[:, step, :, -mask_end:]

							if self.vqvae is not None:
													# 0.18215 was scaling factor used in training to ensure unit variance
													UpperCAmelCase_=       1 / self.vqvae.config.scaling_factor * images
													UpperCAmelCase_=       self.vqvae.decode(__UpperCAmelCase					)["""sample"""]

							UpperCAmelCase_=       (images / 2 + 0.5).clamp(0						,  1					)
							UpperCAmelCase_=       images.cpu().permute(0						,  2						,  3						,  1					).numpy()
							UpperCAmelCase_=       (images * 255).round().astype("""uint8"""					)
							UpperCAmelCase_=       list(
							    (Image.fromarray(_[:, :, 0]					) for _ in images)
							    if images.shape[3] == 1
							    else (Image.fromarray(__UpperCAmelCase						,  mode="""RGB"""					).convert("""L"""					) for _ in images)					)

							UpperCAmelCase_=       [self.mel.image_to_audio(__UpperCAmelCase					) for _ in images]
							if not return_dict:
													return images, (self.mel.get_sample_rate(), audios)

							return BaseOutput(**AudioPipelineOutput(np.array(__UpperCAmelCase					)[:, np.newaxis, :]					)						,  **ImagePipelineOutput(__UpperCAmelCase					)					)




	@torch.no_grad()
	def   _SCREAMING_SNAKE_CASE		(     self     :						List[str]						,  __UpperCAmelCase     :						List[Image.Image]						,  __UpperCAmelCase     :						int = 50					)      ->			np.ndarray:
							assert isinstance(self.scheduler						,  __UpperCAmelCase					)
							self.scheduler.set_timesteps(__UpperCAmelCase					)
							UpperCAmelCase_=       np.array(
							    [np.frombuffer(image.tobytes()						,  dtype="""uint8"""					).reshape((1, image.height, image.width)					) for image in images]					)
							UpperCAmelCase_=       (sample / 255) * 2 - 1
							UpperCAmelCase_=       torch.Tensor(__UpperCAmelCase					).to(self.device					)

							for t in self.progress_bar(torch.flip(self.scheduler.timesteps						,  (0,)					)					):
													UpperCAmelCase_=       t - self.scheduler.config.num_train_timesteps // self.scheduler.num_inference_steps
													UpperCAmelCase_=       self.scheduler.alphas_cumprod[t]
													UpperCAmelCase_=       (
													    self.scheduler.alphas_cumprod[prev_timestep]
													    if prev_timestep >= 0
													    else self.scheduler.final_alpha_cumprod
													)
													UpperCAmelCase_=       1 - alpha_prod_t
													UpperCAmelCase_=       self.unet(__UpperCAmelCase						,  __UpperCAmelCase					)["""sample"""]
													UpperCAmelCase_=       (1 - alpha_prod_t_prev) ** 0.5 * model_output
													UpperCAmelCase_=       (sample - pred_sample_direction) * alpha_prod_t_prev ** (-0.5)
													UpperCAmelCase_=       sample * alpha_prod_t ** 0.5 + beta_prod_t ** 0.5 * model_output

							return sample




	@staticmethod
	def   _SCREAMING_SNAKE_CASE		(     __UpperCAmelCase     :						torch.Tensor						,  __UpperCAmelCase     :						torch.Tensor						,  __UpperCAmelCase     :						float					)      ->			torch.Tensor:
							UpperCAmelCase_=       acos(torch.dot(torch.flatten(__UpperCAmelCase					)						,  torch.flatten(__UpperCAmelCase					)					) / torch.norm(__UpperCAmelCase					) / torch.norm(__UpperCAmelCase					)					)
							return sin((1 - alpha) * theta					) * xa / sin(__UpperCAmelCase					) + sin(alpha * theta					) * xa / sin(__UpperCAmelCase					)




