





"""simple docstring"""




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 			__UpperCamelCase							(		a__		):
							lowerCamelCase  :   Union[str, Any]    =["""vqvae"""]






							def __init__(  self ,   lowerCAmelCase__ ,   lowerCAmelCase__ ,   lowerCAmelCase__ ,   lowerCAmelCase__ ,   )      ->					Optional[int]:
								super().__init__()
								self.register_modules(unet=lowerCAmelCase__ ,   scheduler=lowerCAmelCase__ ,   mel=lowerCAmelCase__ ,   vqvae=lowerCAmelCase__ )






							def        __a  (  self )      ->					int:
								return 50 if isinstance(self.scheduler ,   lowerCAmelCase__ ) else 1000






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

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

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

									if start_step > 0:
										a :						Any      =							self.scheduler.add_noise(lowerCAmelCase__ ,   lowerCAmelCase__ ,   self.scheduler.timesteps[start_step - 1] )

									a :						List[str]      =							(
									    self.unet.config.sample_size[1] * self.mel.get_sample_rate() / self.mel.x_res / self.mel.hop_length
									)
									a :						Dict      =							int(mask_start_secs * pixels_per_second )
									a :						int      =							int(mask_end_secs * pixels_per_second )
									a :						List[Any]      =							self.scheduler.add_noise(lowerCAmelCase__ ,   lowerCAmelCase__ ,   torch.tensor(self.scheduler.timesteps[start_step:] ) )

								for step, t in enumerate(self.progress_bar(self.scheduler.timesteps[start_step:] ) ):
									if isinstance(self.unet ,   lowerCAmelCase__ ):
										a :						int      =							self.unet(lowerCAmelCase__ ,   lowerCAmelCase__ ,   lowerCAmelCase__ )["sample"]
									else:
										a :						List[str]      =							self.unet(lowerCAmelCase__ ,   lowerCAmelCase__ )["sample"]

									if isinstance(self.scheduler ,   lowerCAmelCase__ ):
										a :						Union[str, Any]      =							self.scheduler.step(
										    model_output=lowerCAmelCase__ ,   timestep=lowerCAmelCase__ ,   sample=lowerCAmelCase__ ,   eta=lowerCAmelCase__ ,   generator=lowerCAmelCase__ ,   )["prev_sample"]
									else:
										a :						Any      =							self.scheduler.step(
										    model_output=lowerCAmelCase__ ,   timestep=lowerCAmelCase__ ,   sample=lowerCAmelCase__ ,   generator=lowerCAmelCase__ ,   )["prev_sample"]

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

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

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

								a :						Union[str, Any]      =							[self.mel.image_to_audio(lowerCAmelCase__ ) for _ in images]
								if not return_dict:
									return images, (self.mel.get_sample_rate(), audios)

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






							@torch.no_grad()
							def        __a  (  self ,   lowerCAmelCase__ ,   lowerCAmelCase__ = 50 )      ->					np.ndarray:
								assert isinstance(self.scheduler ,   lowerCAmelCase__ )
								self.scheduler.set_timesteps(lowerCAmelCase__ )
								a :						Dict      =							np.array(
								    [np.frombuffer(image.tobytes() ,   dtype="uint8" ).reshape((1, image.height, image.width) ) for image in images] )
								a :						Optional[Any]      =							(sample / 255) * 2 - 1
								a :						Dict      =							torch.Tensor(lowerCAmelCase__ ).to(self.device )

								for t in self.progress_bar(torch.flip(self.scheduler.timesteps ,   (0,) ) ):
									a :						Optional[int]      =							t - self.scheduler.config.num_train_timesteps // self.scheduler.num_inference_steps
									a :						str      =							self.scheduler.alphas_cumprod[t]
									a :						List[str]      =							(
									    self.scheduler.alphas_cumprod[prev_timestep]
									    if prev_timestep >= 0
									    else self.scheduler.final_alpha_cumprod
									)
									a :						List[Any]      =							1 - alpha_prod_t
									a :						List[str]      =							self.unet(lowerCAmelCase__ ,   lowerCAmelCase__ )["sample"]
									a :						Dict      =							(1 - alpha_prod_t_prev) ** 0.5 * model_output
									a :						Tuple      =							(sample - pred_sample_direction) * alpha_prod_t_prev ** (-0.5)
									a :						Tuple      =							sample * alpha_prod_t ** 0.5 + beta_prod_t ** 0.5 * model_output

								return sample






							@staticmethod
							def        __a  (  lowerCAmelCase__ ,   lowerCAmelCase__ ,   lowerCAmelCase__ )      ->					torch.Tensor:
								a :						Union[str, Any]      =							acos(torch.dot(torch.flatten(lowerCAmelCase__ ) ,   torch.flatten(lowerCAmelCase__ ) ) / torch.norm(lowerCAmelCase__ ) / torch.norm(lowerCAmelCase__ ) )
								return sin((1 - alpha) * theta ) * xa / sin(lowerCAmelCase__ ) + sin(alpha * theta ) * xa / sin(lowerCAmelCase__ )









"""simple docstring"""




import gc
import unittest

from transformers import CTRLConfig, is_torch_available
from transformers.testing_utils import require_torch, slow, torch_device

from ...generation.test_utils import GenerationTesterMixin
from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, ids_tensor, random_attention_mask
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
		import torch

		from transformers import (
		    CTRL_PRETRAINED_MODEL_ARCHIVE_LIST,
		    CTRLForSequenceClassification,
		    CTRLLMHeadModel,
		    CTRLModel,
		)


class 			__UpperCamelCase							:






							def __init__(  self ,   lowerCAmelCase__ ,   lowerCAmelCase__=14 ,   lowerCAmelCase__=7 ,   lowerCAmelCase__=True ,   lowerCAmelCase__=True ,   lowerCAmelCase__=True ,   lowerCAmelCase__=True ,   lowerCAmelCase__=True ,   lowerCAmelCase__=99 ,   lowerCAmelCase__=32 ,   lowerCAmelCase__=5 ,   lowerCAmelCase__=4 ,   lowerCAmelCase__=37 ,   lowerCAmelCase__="gelu" ,   lowerCAmelCase__=0.1 ,   lowerCAmelCase__=0.1 ,   lowerCAmelCase__=512 ,   lowerCAmelCase__=16 ,   lowerCAmelCase__=2 ,   lowerCAmelCase__=0.02 ,   lowerCAmelCase__=3 ,   lowerCAmelCase__=4 ,   lowerCAmelCase__=None ,   )      ->					Optional[int]:
								a :						Optional[int]      =							parent
								a :						Optional[int]      =							batch_size
								a :						str      =							seq_length
								a :						int      =							is_training
								a :						Union[str, Any]      =							use_token_type_ids
								a :						Tuple      =							use_input_mask
								a :						Optional[Any]      =							use_labels
								a :						str      =							use_mc_token_ids
								a :						str      =							vocab_size
								a :						str      =							hidden_size
								a :						List[Any]      =							num_hidden_layers
								a :						str      =							num_attention_heads
								a :						Tuple      =							intermediate_size
								a :						List[Any]      =							hidden_act
								a :						List[Any]      =							hidden_dropout_prob
								a :						List[str]      =							attention_probs_dropout_prob
								a :						Dict      =							max_position_embeddings
								a :						int      =							type_vocab_size
								a :						Any      =							type_sequence_label_size
								a :						List[Any]      =							initializer_range
								a :						List[str]      =							num_labels
								a :						Optional[int]      =							num_choices
								a :						int      =							scope
								a :						Tuple      =							self.vocab_size - 1






							def        __a  (  self )      ->					int:
								a :						Optional[Any]      =							ids_tensor([self.batch_size, self.seq_length] ,   self.vocab_size )

								a :						str      =							None
								if self.use_input_mask:
									a :						Dict      =							random_attention_mask([self.batch_size, self.seq_length] )

								a :						Tuple      =							None
								if self.use_token_type_ids:
									a :						int      =							ids_tensor([self.batch_size, self.seq_length] ,   self.type_vocab_size )

								a :						int      =							None
								if self.use_mc_token_ids:
									a :						Dict      =							ids_tensor([self.batch_size, self.num_choices] ,   self.seq_length )

								a :						List[Any]      =							None
								a :						Optional[int]      =							None
								a :						str      =							None
								if self.use_labels:
									a :						List[str]      =							ids_tensor([self.batch_size] ,   self.type_sequence_label_size )
									a :						List[str]      =							ids_tensor([self.batch_size, self.seq_length] ,   self.num_labels )
									a :						Any      =							ids_tensor([self.batch_size] ,   self.num_choices )

								a :						Tuple      =							self.get_config()

								a :						int      =							ids_tensor([self.num_hidden_layers, self.num_attention_heads] ,   2 )

								return (
								    config,
								    input_ids,
								    input_mask,
								    head_mask,
								    token_type_ids,
								    mc_token_ids,
								    sequence_labels,
								    token_labels,
								    choice_labels,
								)






							def        __a  (  self )      ->					Tuple:
								return CTRLConfig(
								    vocab_size=self.vocab_size ,   n_embd=self.hidden_size ,   n_layer=self.num_hidden_layers ,   n_head=self.num_attention_heads ,   n_positions=self.max_position_embeddings ,   pad_token_id=self.pad_token_id ,   )






							def        __a  (  self ,   lowerCAmelCase__ ,   lowerCAmelCase__ ,   lowerCAmelCase__ ,   lowerCAmelCase__ ,   lowerCAmelCase__ ,   *lowerCAmelCase__ )      ->					str:
								a :						Union[str, Any]      =							CTRLModel(config=lowerCAmelCase__ )
								model.to(lowerCAmelCase__ )
								model.eval()

								model(lowerCAmelCase__ ,   token_type_ids=lowerCAmelCase__ ,   head_mask=lowerCAmelCase__ )
								model(lowerCAmelCase__ ,   token_type_ids=lowerCAmelCase__ )
								a :						List[Any]      =							model(lowerCAmelCase__ )
								self.parent.assertEqual(result.last_hidden_state.shape ,   (self.batch_size, self.seq_length, self.hidden_size) )
								self.parent.assertEqual(len(result.past_key_values ) ,   config.n_layer )






							def        __a  (  self ,   lowerCAmelCase__ ,   lowerCAmelCase__ ,   lowerCAmelCase__ ,   lowerCAmelCase__ ,   lowerCAmelCase__ ,   *lowerCAmelCase__ )      ->					Optional[int]:
								a :						Tuple      =							CTRLLMHeadModel(lowerCAmelCase__ )
								model.to(lowerCAmelCase__ )
								model.eval()

								a :						List[str]      =							model(lowerCAmelCase__ ,   token_type_ids=lowerCAmelCase__ ,   labels=lowerCAmelCase__ )
								self.parent.assertEqual(result.loss.shape ,   () )
								self.parent.assertEqual(result.logits.shape ,   (self.batch_size, self.seq_length, self.vocab_size) )






							def        __a  (  self )      ->					Optional[int]:
								a :						Tuple      =							self.prepare_config_and_inputs()

								(
								    (
								    a
								),			(
								    a
								),			(
								    a
								),			(
								    a
								),			(
								    a
								),			(
								    a
								),			(
								    a
								),			(
								    a
								),			(
								    a
								),			
								) :						Optional[Any]      =							config_and_inputs

								a :						List[Any]      =							{"input_ids": input_ids, "token_type_ids": token_type_ids, "head_mask": head_mask}

								return config, inputs_dict






							def        __a  (  self ,   lowerCAmelCase__ ,   lowerCAmelCase__ ,   lowerCAmelCase__ ,   lowerCAmelCase__ ,   *lowerCAmelCase__ )      ->					Optional[Any]:
								a :						List[str]      =							self.num_labels
								a :						Union[str, Any]      =							CTRLForSequenceClassification(lowerCAmelCase__ )
								model.to(lowerCAmelCase__ )
								model.eval()
								a :						Optional[int]      =							ids_tensor([self.batch_size] ,   self.type_sequence_label_size )
								a :						Dict      =							model(lowerCAmelCase__ ,   token_type_ids=lowerCAmelCase__ ,   labels=lowerCAmelCase__ )
								self.parent.assertEqual(result.logits.shape ,   (self.batch_size, self.num_labels) )


@require_torch
class 			__UpperCamelCase							(		a__						,			a__						,			a__						,			unittest.TestCase		):
							lowerCamelCase  :   Tuple    =(CTRLModel, CTRLLMHeadModel, CTRLForSequenceClassification) if is_torch_available() else ()
							lowerCamelCase  :   List[Any]    =(CTRLLMHeadModel,) if is_torch_available() else ()
							lowerCamelCase  :   List[Any]    =(
							    {
							        """feature-extraction""": CTRLModel,
							        """text-classification""": CTRLForSequenceClassification,
							        """text-generation""": CTRLLMHeadModel,
							        """zero-shot""": CTRLForSequenceClassification,
							    }
							    if is_torch_available()
							    else {}
							)
							lowerCamelCase  :   Optional[int]    =True
							lowerCamelCase  :   List[str]    =False
							lowerCamelCase  :   Optional[Any]    =False






							def        __a  (  self ,   lowerCAmelCase__ ,   lowerCAmelCase__ ,   lowerCAmelCase__ ,   lowerCAmelCase__ ,   lowerCAmelCase__ )      ->					Optional[int]:
								if pipeline_test_casse_name == "ZeroShotClassificationPipelineTests":
									# Get `tokenizer does not have a padding token` error for both fast/slow tokenizers.
									# `CTRLConfig` was never used in pipeline tests, either because of a missing checkpoint or because a tiny
									# config could not be created.
									return True

								return False






							def        __a  (  self )      ->					str:
								a :						str      =							CTRLModelTester(self )
								a :						Dict      =							ConfigTester(self ,   config_class=lowerCAmelCase__ ,   n_embd=37 )






							def        __a  (  self )      ->					Tuple:
								super().tearDown()
								# clean-up as much as possible GPU memory occupied by PyTorch
								gc.collect()
								torch.cuda.empty_cache()






							def        __a  (  self )      ->					Optional[Any]:
								self.config_tester.run_common_tests()






							def        __a  (  self )      ->					Union[str, Any]:
								a :						Optional[int]      =							self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_ctrl_model(*lowerCAmelCase__ )






							def        __a  (  self )      ->					int:
								a :						List[str]      =							self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_lm_head_model(*lowerCAmelCase__ )






							@unittest.skip("Will be fixed soon by reducing the size of the model used for common tests." )
							def        __a  (  self )      ->					Tuple:
								pass






							@slow
							def        __a  (  self )      ->					Optional[Any]:
								for model_name in CTRL_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
									a :						Union[str, Any]      =							CTRLModel.from_pretrained(lowerCAmelCase__ )
									self.assertIsNotNone(lowerCAmelCase__ )






							@unittest.skip("The model doesn't support left padding" )  # and it's not used enough to be worth fixing :)
							def        __a  (  self )      ->					Tuple:
								pass


@require_torch
class 			__UpperCamelCase							(		unittest.TestCase		):






							def        __a  (  self )      ->					int:
								super().tearDown()
								# clean-up as much as possible GPU memory occupied by PyTorch
								gc.collect()
								torch.cuda.empty_cache()






							@slow
							def        __a  (  self )      ->					Dict:
								a :						Dict      =							CTRLLMHeadModel.from_pretrained("ctrl" )
								model.to(lowerCAmelCase__ )
								a :						List[str]      =							torch.tensor(
								    [[1_1859, 0, 1611, 8]] ,   dtype=torch.long ,   device=lowerCAmelCase__ )  # Legal the president is
								a :						Any      =							[
								    1_1859,
								    0,
								    1611,
								    8,
								    5,
								    150,
								    2_6449,
								    2,
								    19,
								    348,
								    469,
								    3,
								    2595,
								    48,
								    2_0740,
								    24_6533,
								    24_6533,
								    19,
								    30,
								    5,
								]  # Legal the president is a good guy and I don't want to lose my job. \n \n I have a

								a :						str      =							model.generate(lowerCAmelCase__ ,   do_sample=lowerCAmelCase__ )
								self.assertListEqual(output_ids[0].tolist() ,   lowerCAmelCase__ )


