





"""simple docstring"""





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

import torch
from transformers import CLIPTextModel, CLIPTokenizer

from ...configuration_utils import ConfigMixin, register_to_config
from ...models import ModelMixin, TransformeraDModel, VQModel
from ...schedulers import VQDiffusionScheduler
from ...utils import logging
from ..pipeline_utils import DiffusionPipeline, ImagePipelineOutput


_UpperCamelCase							=      logging.get_logger(__name__)  # pylint: disable=invalid-name
class 		SCREAMING_SNAKE_CASE_							(    snake_case__       ,  snake_case__   ):
				"""simple docstring"""




				@register_to_config
				def __init__(     self  :Union[str, Any]						,    __lowercase  :bool						,    __lowercase  :Optional[int] = None						,    __lowercase  :Optional[int] = None			):
											super().__init__()

											__lowerCamelCase      :      int         =learnable

											if self.learnable:
																		assert hidden_size is not None, "learnable=True requires `hidden_size` to be set"
																		assert length is not None, "learnable=True requires `length` to be set"

																		__lowerCamelCase      :      Dict         =torch.zeros(__lowercase						,    __lowercase			)
											else:
																		__lowerCamelCase      :      str         =None

											__lowerCamelCase      :      List[Any]         =torch.nn.Parameter(__lowercase			)

class 		SCREAMING_SNAKE_CASE_							(    snake_case__   ):
				"""simple docstring"""
				__snake_case   :       VQModel
				__snake_case   :       CLIPTextModel
				__snake_case   :       CLIPTokenizer
				__snake_case   :       TransformeraDModel
				__snake_case   :       LearnedClassifierFreeSamplingEmbeddings
				__snake_case   :       VQDiffusionScheduler




				def __init__(     self  :Any						,    __lowercase  :VQModel						,    __lowercase  :CLIPTextModel						,    __lowercase  :CLIPTokenizer						,    __lowercase  :TransformeraDModel						,    __lowercase  :VQDiffusionScheduler						,    __lowercase  :LearnedClassifierFreeSamplingEmbeddings						,    ):
											super().__init__()

											self.register_modules(
											    vqvae=__lowercase						,    transformer=__lowercase						,    text_encoder=__lowercase						,    tokenizer=__lowercase						,    scheduler=__lowercase						,    learned_classifier_free_sampling_embeddings=__lowercase						,    )




				def 					__lowercase  (     self  :Any						,    __lowercase  :Tuple						,    __lowercase  :Optional[int]						,    __lowercase  :List[Any]			):
											__lowerCamelCase      :      List[Any]         =len(__lowercase			) if isinstance(__lowercase						,    __lowercase			) else 1

											# get prompt text embeddings
											__lowerCamelCase      :      Tuple         =self.tokenizer(
											    __lowercase						,    padding='''max_length'''						,    max_length=self.tokenizer.model_max_length						,    return_tensors='''pt'''						,    )
											__lowerCamelCase      :      List[str]         =text_inputs.input_ids

											if text_input_ids.shape[-1] > self.tokenizer.model_max_length:
																		__lowerCamelCase      :      int         =self.tokenizer.batch_decode(text_input_ids[:, self.tokenizer.model_max_length :]			)
																		logger.warning(
																		    '''The following part of your input was truncated because CLIP can only handle sequences up to'''
																		    f' {self.tokenizer.model_max_length} tokens: {removed_text}'			)
																		__lowerCamelCase      :      Tuple         =text_input_ids[:, : self.tokenizer.model_max_length]
											__lowerCamelCase      :      Any         =self.text_encoder(text_input_ids.to(self.device			)			)[0]

											# NOTE: This additional step of normalizing the text embeddings is from VQ-Diffusion.
											# While CLIP does normalize the pooled output of the text transformer when combining
											# the image and text embeddings, CLIP does not directly normalize the last hidden state.
											#
											# CLIP normalizing the pooled output.
											# https://github.com/huggingface/transformers/blob/d92e22d1f28324f513f3080e5c47c071a3916721/src/transformers/models/clip/modeling_clip.py#L1052-L1053
											__lowerCamelCase      :      List[str]         =prompt_embeds / prompt_embeds.norm(dim=-1						,    keepdim=__lowercase			)

											# duplicate text embeddings for each generation per prompt
											__lowerCamelCase      :      Optional[int]         =prompt_embeds.repeat_interleave(__lowercase						,    dim=0			)

											if do_classifier_free_guidance:
																		if self.learned_classifier_free_sampling_embeddings.learnable:
																									__lowerCamelCase      :      Union[str, Any]         =self.learned_classifier_free_sampling_embeddings.embeddings
																									__lowerCamelCase      :      Union[str, Any]         =negative_prompt_embeds.unsqueeze(0			).repeat(__lowercase						,    1						,    1			)
																		else:
																									__lowerCamelCase      :      Optional[int]         =[''''''] * batch_size

																									__lowerCamelCase      :      Tuple         =text_input_ids.shape[-1]
																									__lowerCamelCase      :      Dict         =self.tokenizer(
																									    __lowercase						,    padding='''max_length'''						,    max_length=__lowercase						,    truncation=__lowercase						,    return_tensors='''pt'''						,    )
																									__lowerCamelCase      :      Optional[Any]         =self.text_encoder(uncond_input.input_ids.to(self.device			)			)[0]
																									# See comment for normalizing text embeddings
																									__lowerCamelCase      :      List[str]         =negative_prompt_embeds / negative_prompt_embeds.norm(dim=-1						,    keepdim=__lowercase			)

																		# duplicate unconditional embeddings for each generation per prompt, using mps friendly method
																		__lowerCamelCase      :      Dict         =negative_prompt_embeds.shape[1]
																		__lowerCamelCase      :      int         =negative_prompt_embeds.repeat(1						,    __lowercase						,    1			)
																		__lowerCamelCase      :      Dict         =negative_prompt_embeds.view(batch_size * num_images_per_prompt						,    __lowercase						,    -1			)

																		# For classifier free guidance, we need to do two forward passes.
																		# Here we concatenate the unconditional and text embeddings into a single batch
																		# to avoid doing two forward passes
																		__lowerCamelCase      :      Union[str, Any]         =torch.cat([negative_prompt_embeds, prompt_embeds]			)

											return prompt_embeds




				@torch.no_grad()
				def __call__(     self  :Dict						,    __lowercase  :Union[str, List[str]]						,    __lowercase  :int = 100						,    __lowercase  :float = 5.0						,    __lowercase  :float = 1.0						,    __lowercase  :int = 1						,    __lowercase  :Optional[Union[torch.Generator, List[torch.Generator]]] = None						,    __lowercase  :Optional[torch.FloatTensor] = None						,    __lowercase  :Optional[str] = "pil"						,    __lowercase  :bool = True						,    __lowercase  :Optional[Callable[[int, int, torch.FloatTensor], None]] = None						,    __lowercase  :int = 1						,    ):
											if isinstance(__lowercase						,    __lowercase			):
																		__lowerCamelCase      :      Tuple         =1
											elif isinstance(__lowercase						,    __lowercase			):
																		__lowerCamelCase      :      List[str]         =len(__lowercase			)
											else:
																		raise ValueError(f'`prompt` has to be of type `str` or `list` but is {type(__lowercase			)}'			)

											__lowerCamelCase      :      int         =batch_size * num_images_per_prompt

											__lowerCamelCase      :      Optional[int]         =guidance_scale > 1.0

											__lowerCamelCase      :      str         =self._encode_prompt(__lowercase						,    __lowercase						,    __lowercase			)

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

											# get the initial completely masked latents unless the user supplied it

											__lowerCamelCase      :      Optional[int]         =(batch_size, self.transformer.num_latent_pixels)
											if latents is None:
																		__lowerCamelCase      :      Tuple         =self.transformer.num_vector_embeds - 1
																		__lowerCamelCase      :      Any         =torch.full(__lowercase						,    __lowercase			).to(self.device			)
											else:
																		if latents.shape != latents_shape:
																									raise ValueError(f'Unexpected latents shape, got {latents.shape}, expected {latents_shape}'			)
																		if (latents < 0).any() or (latents >= self.transformer.num_vector_embeds).any():
																									raise ValueError(
																									    '''Unexpected latents value(s). All latents be valid embedding indices i.e. in the range 0,'''
																									    f' {self.transformer.num_vector_embeds - 1} (inclusive).'			)
																		__lowerCamelCase      :      Optional[int]         =latents.to(self.device			)

											# set timesteps
											self.scheduler.set_timesteps(__lowercase						,    device=self.device			)

											__lowerCamelCase      :      Optional[int]         =self.scheduler.timesteps.to(self.device			)

											__lowerCamelCase      :      Optional[Any]         =latents

											for i, t in enumerate(self.progress_bar(__lowercase			)			):
																		# expand the sample if we are doing classifier free guidance
																		__lowerCamelCase      :      Union[str, Any]         =torch.cat([sample] * 2			) if do_classifier_free_guidance else sample

																		# predict the un-noised image
																		# model_output == `log_p_x_0`
																		__lowerCamelCase      :      Union[str, Any]         =self.transformer(__lowercase						,    encoder_hidden_states=__lowercase						,    timestep=__lowercase			).sample

																		if do_classifier_free_guidance:
																									__lowerCamelCase      :      Dict         =model_output.chunk(2			)
																									__lowerCamelCase      :      Optional[Any]         =model_output_uncond + guidance_scale * (model_output_text - model_output_uncond)
																									model_output -= torch.logsumexp(__lowercase						,    dim=1						,    keepdim=__lowercase			)

																		__lowerCamelCase      :      List[Any]         =self.truncate(__lowercase						,    __lowercase			)

																		# remove `log(0)`'s (`-inf`s)
																		__lowerCamelCase      :      List[Any]         =model_output.clamp(-70			)

																		# compute the previous noisy sample x_t -> x_t-1
																		__lowerCamelCase      :      Union[str, Any]         =self.scheduler.step(__lowercase						,    timestep=__lowercase						,    sample=__lowercase						,    generator=__lowercase			).prev_sample

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

											__lowerCamelCase      :      Any         =self.vqvae.config.vq_embed_dim
											__lowerCamelCase      :      List[str]         =(batch_size, self.transformer.height, self.transformer.width, embedding_channels)
											__lowerCamelCase      :      Optional[int]         =self.vqvae.quantize.get_codebook_entry(__lowercase						,    shape=__lowercase			)
											__lowerCamelCase      :      List[str]         =self.vqvae.decode(__lowercase						,    force_not_quantize=__lowercase			).sample

											__lowerCamelCase      :      str         =(image / 2 + 0.5).clamp(0						,    1			)
											__lowerCamelCase      :      Dict         =image.cpu().permute(0						,    2						,    3						,    1			).numpy()

											if output_type == "pil":
																		__lowerCamelCase      :      str         =self.numpy_to_pil(__lowercase			)

											if not return_dict:
																		return (image,)

											return ImagePipelineOutput(images=__lowercase			)




				def 					__lowercase  (     self  :Any						,    __lowercase  :torch.FloatTensor						,    __lowercase  :float			):
											__lowerCamelCase      :      Optional[Any]         =torch.sort(__lowercase						,    1						,    descending=__lowercase			)
											__lowerCamelCase      :      List[str]         =torch.exp(__lowercase			)
											__lowerCamelCase      :      List[str]         =sorted_p_x_0.cumsum(dim=1			) < truncation_rate

											# Ensure that at least the largest probability is not zeroed out
											__lowerCamelCase      :      Dict         =torch.full_like(keep_mask[:, 0:1, :]						,    __lowercase			)
											__lowerCamelCase      :      Tuple         =torch.cat((all_true, keep_mask)						,    dim=1			)
											__lowerCamelCase      :      Dict         =keep_mask[:, :-1, :]

											__lowerCamelCase      :      int         =keep_mask.gather(1						,    indices.argsort(1			)			)

											__lowerCamelCase      :      Union[str, Any]         =log_p_x_0.clone()

											__lowerCamelCase      :      Any         =-torch.inf  # -inf = log(0)

											return rv











"""simple docstring"""





from typing import List, Optional, Union

from ...configuration_utils import PretrainedConfig
from ...utils import logging


_UpperCamelCase							=      logging.get_logger(__name__)

_UpperCamelCase							=      {
    'huggingface/time-series-transformer-tourism-monthly': (
        'https://huggingface.co/huggingface/time-series-transformer-tourism-monthly/resolve/main/config.json'
    ),
    # See all TimeSeriesTransformer models at https://huggingface.co/models?filter=time_series_transformer
}
class 		SCREAMING_SNAKE_CASE_							(    snake_case__   ):
				"""simple docstring"""
				__snake_case   :       Union[str, Any]	 =		"""time_series_transformer"""
				__snake_case   :       Any	 =		{
				    """hidden_size""": """d_model""",
				    """num_attention_heads""": """encoder_attention_heads""",
				    """num_hidden_layers""": """encoder_layers""",
				}




				def __init__(     self  :List[Any]						,    __lowercase  :Optional[int] = None						,    __lowercase  :Optional[int] = None						,    __lowercase  :str = "student_t"						,    __lowercase  :str = "nll"						,    __lowercase  :int = 1						,    __lowercase  :List[int] = [1, 2, 3, 4, 5, 6, 7]						,    __lowercase  :Optional[Union[str, bool]] = "mean"						,    __lowercase  :int = 0						,    __lowercase  :int = 0						,    __lowercase  :int = 0						,    __lowercase  :int = 0						,    __lowercase  :Optional[List[int]] = None						,    __lowercase  :Optional[List[int]] = None						,    __lowercase  :int = 32						,    __lowercase  :int = 32						,    __lowercase  :int = 2						,    __lowercase  :int = 2						,    __lowercase  :int = 2						,    __lowercase  :int = 2						,    __lowercase  :bool = True						,    __lowercase  :str = "gelu"						,    __lowercase  :int = 64						,    __lowercase  :float = 0.1						,    __lowercase  :float = 0.1						,    __lowercase  :float = 0.1						,    __lowercase  :float = 0.1						,    __lowercase  :float = 0.1						,    __lowercase  :int = 100						,    __lowercase  :float = 0.02						,    __lowercase  :int=True						,    **__lowercase  :Optional[Any]						,    ):
											# time series specific configuration
											__lowerCamelCase      :      Tuple         =prediction_length
											__lowerCamelCase      :      List[Any]         =context_length or prediction_length
											__lowerCamelCase      :      Dict         =distribution_output
											__lowerCamelCase      :      str         =loss
											__lowerCamelCase      :      Tuple         =input_size
											__lowerCamelCase      :      int         =num_time_features
											__lowerCamelCase      :      int         =lags_sequence
											__lowerCamelCase      :      Optional[int]         =scaling
											__lowerCamelCase      :      str         =num_dynamic_real_features
											__lowerCamelCase      :      Optional[Any]         =num_static_real_features
											__lowerCamelCase      :      List[Any]         =num_static_categorical_features
											if cardinality and num_static_categorical_features > 0:
																		if len(__lowercase			) != num_static_categorical_features:
																									raise ValueError(
																									    '''The cardinality should be a list of the same length as `num_static_categorical_features`'''			)
																		__lowerCamelCase      :      Optional[int]         =cardinality
											else:
																		__lowerCamelCase      :      str         =[0]
											if embedding_dimension and num_static_categorical_features > 0:
																		if len(__lowercase			) != num_static_categorical_features:
																									raise ValueError(
																									    '''The embedding dimension should be a list of the same length as `num_static_categorical_features`'''			)
																		__lowerCamelCase      :      int         =embedding_dimension
											else:
																		__lowerCamelCase      :      Optional[int]         =[min(50						,    (cat + 1) // 2			) for cat in self.cardinality]
											__lowerCamelCase      :      List[str]         =num_parallel_samples

											# Transformer architecture configuration
											__lowerCamelCase      :      str         =input_size * len(__lowercase			) + self._number_of_features
											__lowerCamelCase      :      Union[str, Any]         =d_model
											__lowerCamelCase      :      int         =encoder_attention_heads
											__lowerCamelCase      :      int         =decoder_attention_heads
											__lowerCamelCase      :      Dict         =encoder_ffn_dim
											__lowerCamelCase      :      int         =decoder_ffn_dim
											__lowerCamelCase      :      List[Any]         =encoder_layers
											__lowerCamelCase      :      int         =decoder_layers

											__lowerCamelCase      :      Union[str, Any]         =dropout
											__lowerCamelCase      :      Optional[Any]         =attention_dropout
											__lowerCamelCase      :      List[str]         =activation_dropout
											__lowerCamelCase      :      List[str]         =encoder_layerdrop
											__lowerCamelCase      :      int         =decoder_layerdrop

											__lowerCamelCase      :      Tuple         =activation_function
											__lowerCamelCase      :      str         =init_std

											__lowerCamelCase      :      Dict         =use_cache

											super().__init__(is_encoder_decoder=__lowercase						,    **__lowercase			)




				@property
				def 					__lowercase  (     self  :int			):
											return (
											    sum(self.embedding_dimension			)
											    + self.num_dynamic_real_features
											    + self.num_time_features
											    + self.num_static_real_features
											    + self.input_size * 2  # the log1p(abs(loc)) and log(scale) features
											)




