


"""simple docstring"""


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

import torch
from transformers import CLIPImageProcessor, CLIPTextModel, CLIPTokenizer

from diffusers import DiffusionPipeline
from diffusers.models import AutoencoderKL, UNetaDConditionModel
from diffusers.pipelines.stable_diffusion import StableDiffusionPipelineOutput
from diffusers.pipelines.stable_diffusion.safety_checker import StableDiffusionSafetyChecker
from diffusers.schedulers import DDIMScheduler, LMSDiscreteScheduler, PNDMScheduler
from diffusers.utils import logging


__lowerCamelCase      :List[str]	    =		logging.get_logger(__name__)  # pylint: disable=invalid-name


class 			A__	(						__lowercase):


     """simple docstring"""

     def __init__(	self:						List[Any]     ,		__a:						AutoencoderKL     ,		__a:						CLIPTextModel     ,		__a:						CLIPTokenizer     ,		__a:						UNetaDConditionModel     ,		__a:						Union[DDIMScheduler, PNDMScheduler, LMSDiscreteScheduler]     ,		__a:						StableDiffusionSafetyChecker     ,		__a:						CLIPImageProcessor     ,		)->    List[str]:
           super().__init__()
           self.register_modules(
               vae=__a     ,		text_encoder=__a     ,		tokenizer=__a     ,		unet=__a     ,		scheduler=__a     ,		safety_checker=__a     ,		feature_extractor=__a     ,		)

     def 		a__						(	self:						Any     ,		__a:						Optional[Union[str, int]] = "auto"	)->    Tuple:
           if slice_size == "auto":
                 # half the attention head size is usually a good trade-off between
                 # speed and memory
                 lowerCamelCase					:					Tuple        =	self.unet.config.attention_head_dim // 2
           self.unet.set_attention_slice(__a	)

     def 		a__						(	self:						Optional[Any]	)->    int:
           self.enable_attention_slicing(__a	)







     @torch.no_grad()
     def __call__(	self:						List[str]     ,		__a:						Union[str, List[str]]     ,		__a:						int = 512     ,		__a:						int = 512     ,		__a:						int = 50     ,		__a:						float = 7.5     ,		__a:						Optional[Union[str, List[str]]] = None     ,		__a:						Optional[int] = 1     ,		__a:						float = 0.0     ,		__a:						Optional[torch.Generator] = None     ,		__a:						Optional[torch.FloatTensor] = None     ,		__a:						Optional[str] = "pil"     ,		__a:						bool = True     ,		__a:						Optional[Callable[[int, int, torch.FloatTensor], None]] = None     ,		__a:						int = 1     ,		__a:						Optional[torch.FloatTensor] = None     ,		**__a:						Dict     ,		)->    List[str]:

           if isinstance(__a     ,		__a	):
                 lowerCamelCase					:					Union[str, Any]        =	1
           elif isinstance(__a     ,		__a	):
                 lowerCamelCase					:					Tuple        =	len(__a	)
           else:
                 raise ValueError(f'`prompt` has to be of type `str` or `list` but is {type(__a	)}'	)

           if height % 8 != 0 or width % 8 != 0:
                 raise ValueError(f'`height` and `width` have to be divisible by 8 but are {height} and {width}.'	)

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

           # get prompt text embeddings
           lowerCamelCase					:					Tuple        =	self.tokenizer(
               __a     ,		padding="""max_length"""     ,		max_length=self.tokenizer.model_max_length     ,		return_tensors="""pt"""     ,		)
           lowerCamelCase					:					Dict        =	text_inputs.input_ids

           if text_input_ids.shape[-1] > self.tokenizer.model_max_length:
                 lowerCamelCase					:					Union[str, Any]        =	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					:					Union[str, Any]        =	text_input_ids[:, : self.tokenizer.model_max_length]

           if text_embeddings is None:
                 lowerCamelCase					:					Any        =	self.text_encoder(text_input_ids.to(self.device	)	)[0]

           # duplicate text embeddings for each generation per prompt, using mps friendly method
           lowerCamelCase     ,			lowerCamelCase     ,			lowerCamelCase					:					Union[str, Any]        =	text_embeddings.shape
           lowerCamelCase					:					List[Any]        =	text_embeddings.repeat(1     ,		__a     ,		1	)
           lowerCamelCase					:					Optional[Any]        =	text_embeddings.view(bs_embed * num_images_per_prompt     ,		__a     ,		-1	)

           # here `guidance_scale` is defined analog to the guidance weight `w` of equation (2)
           # of the Imagen paper: https://arxiv.org/pdf/2205.11487.pdf . `guidance_scale = 1`
           # corresponds to doing no classifier free guidance.
           lowerCamelCase					:					Union[str, Any]        =	guidance_scale > 1.0
           # get unconditional embeddings for classifier free guidance
           if do_classifier_free_guidance:
                 lowerCamelCase					:					List[str]
                 if negative_prompt is None:
                       lowerCamelCase					:					Optional[int]        =	[""""""]
                 elif type(__a	) is not type(__a	):
                       raise TypeError(
                           f'`negative_prompt` should be the same type to `prompt`, but got {type(__a	)} !='
                           f' {type(__a	)}.'	)
                 elif isinstance(__a     ,		__a	):
                       lowerCamelCase					:					Dict        =	[negative_prompt]
                 elif batch_size != len(__a	):
                       raise ValueError(
                           f'`negative_prompt`: {negative_prompt} has batch size {len(__a	)}, but `prompt`:'
                           f' {prompt} has batch size {batch_size}. Please make sure that passed `negative_prompt` matches'
                           """ the batch size of `prompt`."""	)
                 else:
                       lowerCamelCase					:					Union[str, Any]        =	negative_prompt

                 lowerCamelCase					:					Any        =	text_input_ids.shape[-1]
                 lowerCamelCase					:					List[str]        =	self.tokenizer(
                     __a     ,		padding="""max_length"""     ,		max_length=__a     ,		truncation=__a     ,		return_tensors="""pt"""     ,		)
                 lowerCamelCase					:					int        =	self.text_encoder(uncond_input.input_ids.to(self.device	)	)[0]

                 # duplicate unconditional embeddings for each generation per prompt, using mps friendly method
                 lowerCamelCase					:					List[Any]        =	uncond_embeddings.shape[1]
                 lowerCamelCase					:					Any        =	uncond_embeddings.repeat(__a     ,		__a     ,		1	)
                 lowerCamelCase					:					List[str]        =	uncond_embeddings.view(batch_size * num_images_per_prompt     ,		__a     ,		-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					:					Tuple        =	torch.cat([uncond_embeddings, text_embeddings]	)

           # get the initial random noise unless the user supplied it

           # Unlike in other pipelines, latents need to be generated in the target device
           # for 1-to-1 results reproducibility with the CompVis implementation.
           # However this currently doesn't work in `mps`.
           lowerCamelCase					:					int        =	(batch_size * num_images_per_prompt, self.unet.config.in_channels, height // 8, width // 8)
           lowerCamelCase					:					Optional[int]        =	(batch_size * num_images_per_prompt, self.unet.config.in_channels, 64, 64)
           lowerCamelCase					:					str        =	text_embeddings.dtype
           if latents is None:
                 if self.device.type == "mps":
                       # randn does not exist on mps
                       lowerCamelCase					:					Union[str, Any]        =	torch.randn(
                           __a     ,		generator=__a     ,		device="""cpu"""     ,		dtype=__a	).to(self.device	)
                       lowerCamelCase					:					Optional[int]        =	torch.randn(__a     ,		generator=__a     ,		device="""cpu"""     ,		dtype=__a	).to(
                           self.device	)
                 else:
                       lowerCamelCase					:					int        =	torch.randn(
                           __a     ,		generator=__a     ,		device=self.device     ,		dtype=__a	)
                       lowerCamelCase					:					Union[str, Any]        =	torch.randn(__a     ,		generator=__a     ,		device=self.device     ,		dtype=__a	)
           else:
                 if latents_reference.shape != latents_shape:
                       raise ValueError(f'Unexpected latents shape, got {latents.shape}, expected {latents_shape}'	)
                 lowerCamelCase					:					Union[str, Any]        =	latents_reference.to(self.device	)
                 lowerCamelCase					:					Optional[Any]        =	latents.to(self.device	)

           # This is the key part of the pipeline where we
           # try to ensure that the generated images w/ the same seed
           # but different sizes actually result in similar images
           lowerCamelCase					:					Tuple        =	(latents_shape[3] - latents_shape_reference[3]) // 2
           lowerCamelCase					:					int        =	(latents_shape[2] - latents_shape_reference[2]) // 2
           lowerCamelCase					:					int        =	latents_shape_reference[3] if dx >= 0 else latents_shape_reference[3] + 2 * dx
           lowerCamelCase					:					Union[str, Any]        =	latents_shape_reference[2] if dy >= 0 else latents_shape_reference[2] + 2 * dy
           lowerCamelCase					:					Tuple        =	0 if dx < 0 else dx
           lowerCamelCase					:					Any        =	0 if dy < 0 else dy
           lowerCamelCase					:					Optional[Any]        =	max(-dx     ,		0	)
           lowerCamelCase					:					Tuple        =	max(-dy     ,		0	)
           # import pdb
           # pdb.set_trace()
           lowerCamelCase					:					Tuple        =	latents_reference[:, :, dy : dy + h, dx : dx + w]

           # set timesteps
           self.scheduler.set_timesteps(__a	)

           # Some schedulers like PNDM have timesteps as arrays
           # It's more optimized to move all timesteps to correct device beforehand
           lowerCamelCase					:					List[str]        =	self.scheduler.timesteps.to(self.device	)

           # scale the initial noise by the standard deviation required by the scheduler
           lowerCamelCase					:					int        =	latents * self.scheduler.init_noise_sigma

           # prepare extra kwargs for the scheduler step, since not all schedulers have the same signature
           # eta (η) is only used with the DDIMScheduler, it will be ignored for other schedulers.
           # eta corresponds to η in DDIM paper: https://arxiv.org/abs/2010.02502
           # and should be between [0, 1]
           lowerCamelCase					:					int        =	"""eta""" in set(inspect.signature(self.scheduler.step	).parameters.keys()	)
           lowerCamelCase					:					List[str]        =	{}
           if accepts_eta:
                 lowerCamelCase					:					Optional[Any]        =	eta

           for i, t in enumerate(self.progress_bar(__a	)	):
                 # expand the latents if we are doing classifier free guidance
                 lowerCamelCase					:					List[Any]        =	torch.cat([latents] * 2	) if do_classifier_free_guidance else latents
                 lowerCamelCase					:					Union[str, Any]        =	self.scheduler.scale_model_input(__a     ,		__a	)

                 # predict the noise residual
                 lowerCamelCase					:					Union[str, Any]        =	self.unet(__a     ,		__a     ,		encoder_hidden_states=__a	).sample

                 # perform guidance
                 if do_classifier_free_guidance:
                       lowerCamelCase     ,			lowerCamelCase					:					str        =	noise_pred.chunk(2	)
                       lowerCamelCase					:					List[Any]        =	noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond)

                 # compute the previous noisy sample x_t -> x_t-1
                 lowerCamelCase					:					str        =	self.scheduler.step(__a     ,		__a     ,		__a     ,		**__a	).prev_sample

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

           lowerCamelCase					:					str        =	1 / 0.1_82_15 * latents
           lowerCamelCase					:					Optional[Any]        =	self.vae.decode(__a	).sample

           lowerCamelCase					:					Tuple        =	(image / 2 + 0.5).clamp(0     ,		1	)

           # we always cast to float32 as this does not cause significant overhead and is compatible with bfloat16
           lowerCamelCase					:					Dict        =	image.cpu().permute(0     ,		2     ,		3     ,		1	).float().numpy()

           if self.safety_checker is not None:
                 lowerCamelCase					:					int        =	self.feature_extractor(self.numpy_to_pil(__a	)     ,		return_tensors="""pt"""	).to(
                     self.device	)
                 lowerCamelCase     ,			lowerCamelCase					:					Optional[int]        =	self.safety_checker(
                     images=__a     ,		clip_input=safety_checker_input.pixel_values.to(text_embeddings.dtype	)	)
           else:
                 lowerCamelCase					:					Dict        =	None

           if output_type == "pil":
                 lowerCamelCase					:					Tuple        =	self.numpy_to_pil(__a	)

           if not return_dict:
                 return (image, has_nsfw_concept)

           return StableDiffusionPipelineOutput(images=__a     ,		nsfw_content_detected=__a	)








"""simple docstring"""


from typing import List, Optional, Union

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


__lowerCamelCase      :str	    =		logging.get_logger(__name__)

__lowerCamelCase      :Any	    =		{
    '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 			A__	(						__lowercase):


     """simple docstring"""
     snake_case__	:					List[Any]													='''time_series_transformer'''
     snake_case__	:					List[Any]													={
         '''hidden_size''': '''d_model''',
         '''num_attention_heads''': '''encoder_attention_heads''',
         '''num_hidden_layers''': '''encoder_layers''',
     }

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

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

           lowerCamelCase					:					Optional[int]        =	dropout
           lowerCamelCase					:					List[str]        =	attention_dropout
           lowerCamelCase					:					Tuple        =	activation_dropout
           lowerCamelCase					:					Optional[int]        =	encoder_layerdrop
           lowerCamelCase					:					int        =	decoder_layerdrop

           lowerCamelCase					:					Optional[int]        =	activation_function
           lowerCamelCase					:					Optional[Any]        =	init_std

           lowerCamelCase					:					Optional[Any]        =	use_cache

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







     @property
     def 		a__						(	self:						int	)->    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
           )




