



'''simple docstring'''






from typing import List, Optional, Union

import torch
from transformers import (
    XLMRobertaTokenizer,
)

from ...models import UNetaDConditionModel, VQModel
from ...pipelines import DiffusionPipeline
from ...pipelines.pipeline_utils import ImagePipelineOutput
from ...schedulers import DDIMScheduler, DDPMScheduler
from ...utils import (
    is_accelerate_available,
    is_accelerate_version,
    logging,
    randn_tensor,
    replace_example_docstring,
)
from .text_encoder import MultilingualCLIP


UpperCamelCase__:     Union[str, Any]										=    logging.get_logger(__name__)  # pylint: disable=invalid-name

UpperCamelCase__:     Dict										=    "\n    Examples:\n        ```py\n        >>> from diffusers import KandinskyPipeline, KandinskyPriorPipeline\n        >>> import torch\n\n        >>> pipe_prior = KandinskyPriorPipeline.from_pretrained(\"kandinsky-community/Kandinsky-2-1-prior\")\n        >>> pipe_prior.to(\"cuda\")\n\n        >>> prompt = \"red cat, 4k photo\"\n        >>> out = pipe_prior(prompt)\n        >>> image_emb = out.image_embeds\n        >>> negative_image_emb = out.negative_image_embeds\n\n        >>> pipe = KandinskyPipeline.from_pretrained(\"kandinsky-community/kandinsky-2-1\")\n        >>> pipe.to(\"cuda\")\n\n        >>> image = pipe(\n        ...     prompt,\n        ...     image_embeds=image_emb,\n        ...     negative_image_embeds=negative_image_emb,\n        ...     height=768,\n        ...     width=768,\n        ...     num_inference_steps=100,\n        ... ).images\n\n        >>> image[0].save(\"cat.png\")\n        ```\n"


def 				snake_case_  (      _lowerCAmelCase						:		Dict     ,   _lowerCAmelCase						:		Dict     ,   _lowerCAmelCase						:		Any=8						)   ->     List[Any]:
       UpperCAmelCase				: Optional[Any]											=	h // scale_factor**2
       if h % scale_factor**2 != 0:
              new_h += 1
       UpperCAmelCase				: Tuple											=	w // scale_factor**2
       if w % scale_factor**2 != 0:
              new_w += 1
       return new_h * scale_factor, new_w * scale_factor




class        SCREAMING_SNAKE_CASE(							A__    ):







 """simple docstring"""






 def __init__(		self    :  int	,    __snake_case    :  MultilingualCLIP	,    __snake_case    :  XLMRobertaTokenizer	,    __snake_case    :  UNetaDConditionModel	,    __snake_case    :  Union[DDIMScheduler, DDPMScheduler]	,    __snake_case    :  VQModel	,    )    ->	Optional[int]:
        super().__init__()

        self.register_modules(
            text_encoder=UpperCAmelCase__	,    tokenizer=UpperCAmelCase__	,    unet=UpperCAmelCase__	,    scheduler=UpperCAmelCase__	,    movq=UpperCAmelCase__	,    )
        UpperCAmelCase				: Union[str, Any]											=	2 ** (len(self.movq.config.block_out_channels	) - 1)







 def   A			(		self    :  Union[str, Any]	,    __snake_case    :  List[str]	,    __snake_case    :  List[str]	,    __snake_case    :  str	,    __snake_case    :  List[str]	,    __snake_case    :  List[Any]	,    __snake_case    :  Tuple	)    ->	Dict:
        if latents is None:
               UpperCAmelCase				: Any											=	randn_tensor(UpperCAmelCase__	,    generator=UpperCAmelCase__	,    device=UpperCAmelCase__	,    dtype=UpperCAmelCase__	)
        else:
               if latents.shape != shape:
                      raise ValueError(F"""Unexpected latents shape, got {latents.shape}, expected {shape}"""	)
               UpperCAmelCase				: int											=	latents.to(UpperCAmelCase__	)

        UpperCAmelCase				: int											=	latents * scheduler.init_noise_sigma
        return latents







 def   A			(		self    :  Tuple	,    __snake_case    :  Optional[int]	,    __snake_case    :  Any	,    __snake_case    :  Optional[Any]	,    __snake_case    :  str	,    __snake_case    :  Any=None	,    )    ->	List[Any]:
        UpperCAmelCase				: Dict											=	len(UpperCAmelCase__	) if isinstance(UpperCAmelCase__	,    UpperCAmelCase__	) else 1
        # get prompt text embeddings
        UpperCAmelCase				: Any											=	self.tokenizer(
            UpperCAmelCase__	,    padding='''max_length'''	,    truncation=UpperCAmelCase__	,    max_length=77	,    return_attention_mask=UpperCAmelCase__	,    add_special_tokens=UpperCAmelCase__	,    return_tensors='''pt'''	,    )

        UpperCAmelCase				: List[str]											=	text_inputs.input_ids
        UpperCAmelCase				: List[Any]											=	self.tokenizer(UpperCAmelCase__	,    padding='''longest'''	,    return_tensors='''pt'''	).input_ids

        if untruncated_ids.shape[-1] >= text_input_ids.shape[-1] and not torch.equal(UpperCAmelCase__	,    UpperCAmelCase__	):
               UpperCAmelCase				: Optional[Any]											=	self.tokenizer.batch_decode(untruncated_ids[:, self.tokenizer.model_max_length - 1 : -1]	)
               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}"""	)

        UpperCAmelCase				: Dict											=	text_input_ids.to(UpperCAmelCase__	)
        UpperCAmelCase				: List[str]											=	text_inputs.attention_mask.to(UpperCAmelCase__	)

        UpperCAmelCase							,	UpperCAmelCase				: Tuple											=	self.text_encoder(
            input_ids=UpperCAmelCase__	,    attention_mask=UpperCAmelCase__	)

        UpperCAmelCase				: List[Any]											=	prompt_embeds.repeat_interleave(UpperCAmelCase__	,    dim=0	)
        UpperCAmelCase				: List[Any]											=	text_encoder_hidden_states.repeat_interleave(UpperCAmelCase__	,    dim=0	)
        UpperCAmelCase				: List[str]											=	text_mask.repeat_interleave(UpperCAmelCase__	,    dim=0	)

        if do_classifier_free_guidance:
               UpperCAmelCase				: Optional[int]											=	42
               if negative_prompt is None:
                      UpperCAmelCase				: Union[str, Any]											=	[''''''] * batch_size
               elif type(UpperCAmelCase__	) is not type(UpperCAmelCase__	):
                      raise TypeError(
                          F"""`negative_prompt` should be the same type to `prompt`, but got {type(UpperCAmelCase__	)} !="""
                          F""" {type(UpperCAmelCase__	)}."""	)
               elif isinstance(UpperCAmelCase__	,    UpperCAmelCase__	):
                      UpperCAmelCase				: Optional[Any]											=	[negative_prompt]
               elif batch_size != len(UpperCAmelCase__	):
                      raise ValueError(
                          F"""`negative_prompt`: {negative_prompt} has batch size {len(UpperCAmelCase__	)}, but `prompt`:"""
                          F""" {prompt} has batch size {batch_size}. Please make sure that passed `negative_prompt` matches"""
                          ''' the batch size of `prompt`.'''	)
               else:
                      UpperCAmelCase				: List[Any]											=	negative_prompt

               UpperCAmelCase				: Optional[int]											=	self.tokenizer(
                   UpperCAmelCase__	,    padding='''max_length'''	,    max_length=77	,    truncation=UpperCAmelCase__	,    return_attention_mask=UpperCAmelCase__	,    add_special_tokens=UpperCAmelCase__	,    return_tensors='''pt'''	,    )
               UpperCAmelCase				: List[str]											=	uncond_input.input_ids.to(UpperCAmelCase__	)
               UpperCAmelCase				: str											=	uncond_input.attention_mask.to(UpperCAmelCase__	)

               UpperCAmelCase							,	UpperCAmelCase				: Dict											=	self.text_encoder(
                   input_ids=UpperCAmelCase__	,    attention_mask=UpperCAmelCase__	)

               # duplicate unconditional embeddings for each generation per prompt, using mps friendly method

               UpperCAmelCase				: Optional[Any]											=	negative_prompt_embeds.shape[1]
               UpperCAmelCase				: List[Any]											=	negative_prompt_embeds.repeat(1	,    UpperCAmelCase__	)
               UpperCAmelCase				: int											=	negative_prompt_embeds.view(batch_size * num_images_per_prompt	,    UpperCAmelCase__	)

               UpperCAmelCase				: Any											=	uncond_text_encoder_hidden_states.shape[1]
               UpperCAmelCase				: int											=	uncond_text_encoder_hidden_states.repeat(1	,    UpperCAmelCase__	,    1	)
               UpperCAmelCase				: Union[str, Any]											=	uncond_text_encoder_hidden_states.view(
                   batch_size * num_images_per_prompt	,    UpperCAmelCase__	,    -1	)
               UpperCAmelCase				: str											=	uncond_text_mask.repeat_interleave(UpperCAmelCase__	,    dim=0	)

               # done duplicates

               # 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
               UpperCAmelCase				: Tuple											=	torch.cat([negative_prompt_embeds, prompt_embeds]	)
               UpperCAmelCase				: Tuple											=	torch.cat([uncond_text_encoder_hidden_states, text_encoder_hidden_states]	)

               UpperCAmelCase				: str											=	torch.cat([uncond_text_mask, text_mask]	)

        return prompt_embeds, text_encoder_hidden_states, text_mask







 def   A			(		self    :  int	,    __snake_case    :  Optional[Any]=0	)    ->	List[str]:
        if is_accelerate_available():
               from accelerate import cpu_offload
        else:
               raise ImportError('''Please install accelerate via `pip install accelerate`'''	)

        UpperCAmelCase				: Optional[Any]											=	torch.device(F"""cuda:{gpu_id}"""	)

        UpperCAmelCase				: Any											=	[
            self.unet,
            self.text_encoder,
            self.movq,
        ]
        for cpu_offloaded_model in models:
               if cpu_offloaded_model is not None:
                      cpu_offload(UpperCAmelCase__	,    UpperCAmelCase__	)







 def   A			(		self    :  int	,    __snake_case    :  Dict=0	)    ->	List[str]:
        if is_accelerate_available() and is_accelerate_version('''>='''	,    '''0.17.0.dev0'''	):
               from accelerate import cpu_offload_with_hook
        else:
               raise ImportError('''`enable_model_cpu_offload` requires `accelerate v0.17.0` or higher.'''	)

        UpperCAmelCase				: Any											=	torch.device(F"""cuda:{gpu_id}"""	)

        if self.device.type != "cpu":
               self.to('''cpu'''	,    silence_dtype_warnings=UpperCAmelCase__	)
               torch.cuda.empty_cache()  # otherwise we don't see the memory savings (but they probably exist)

        UpperCAmelCase				: Any											=	None
        for cpu_offloaded_model in [self.text_encoder, self.unet, self.movq]:
               UpperCAmelCase							,	UpperCAmelCase				: str											=	cpu_offload_with_hook(UpperCAmelCase__	,    UpperCAmelCase__	,    prev_module_hook=UpperCAmelCase__	)

        if self.safety_checker is not None:
               UpperCAmelCase							,	UpperCAmelCase				: Dict											=	cpu_offload_with_hook(self.safety_checker	,    UpperCAmelCase__	,    prev_module_hook=UpperCAmelCase__	)

        # We'll offload the last model manually.
        UpperCAmelCase				: Optional[int]											=	hook







 @property
 # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline._execution_device
 def   A			(		self    :  List[Any]	)    ->	int:
        if not hasattr(self.unet	,    '''_hf_hook'''	):
               return self.device
        for module in self.unet.modules():
               if (
                   hasattr(UpperCAmelCase__	,    '''_hf_hook'''	)
                   and hasattr(module._hf_hook	,    '''execution_device'''	)
                   and module._hf_hook.execution_device is not None
               ):
                      return torch.device(module._hf_hook.execution_device	)
        return self.device







 @torch.no_grad()
 @replace_example_docstring(UpperCAmelCase__	)
 def __call__(		self    :  int	,    __snake_case    :  Union[str, List[str]]	,    __snake_case    :  Union[torch.FloatTensor, List[torch.FloatTensor]]	,    __snake_case    :  Union[torch.FloatTensor, List[torch.FloatTensor]]	,    __snake_case    :  Optional[Union[str, List[str]]] = None	,    __snake_case    :  int = 512	,    __snake_case    :  int = 512	,    __snake_case    :  int = 100	,    __snake_case    :  float = 4.0	,    __snake_case    :  int = 1	,    __snake_case    :  Optional[Union[torch.Generator, List[torch.Generator]]] = None	,    __snake_case    :  Optional[torch.FloatTensor] = None	,    __snake_case    :  Optional[str] = "pil"	,    __snake_case    :  bool = True	,    )    ->	str:

        if isinstance(UpperCAmelCase__	,    UpperCAmelCase__	):
               UpperCAmelCase				: Tuple											=	1
        elif isinstance(UpperCAmelCase__	,    UpperCAmelCase__	):
               UpperCAmelCase				: Any											=	len(UpperCAmelCase__	)
        else:
               raise ValueError(F"""`prompt` has to be of type `str` or `list` but is {type(UpperCAmelCase__	)}"""	)

        UpperCAmelCase				: List[str]											=	self._execution_device

        UpperCAmelCase				: Dict											=	batch_size * num_images_per_prompt
        UpperCAmelCase				: List[str]											=	guidance_scale > 1.0

        UpperCAmelCase							,	UpperCAmelCase							,	UpperCAmelCase				: Any											=	self._encode_prompt(
            UpperCAmelCase__	,    UpperCAmelCase__	,    UpperCAmelCase__	,    UpperCAmelCase__	,    UpperCAmelCase__	)

        if isinstance(UpperCAmelCase__	,    UpperCAmelCase__	):
               UpperCAmelCase				: List[Any]											=	torch.cat(UpperCAmelCase__	,    dim=0	)
        if isinstance(UpperCAmelCase__	,    UpperCAmelCase__	):
               UpperCAmelCase				: str											=	torch.cat(UpperCAmelCase__	,    dim=0	)

        if do_classifier_free_guidance:
               UpperCAmelCase				: Optional[Any]											=	image_embeds.repeat_interleave(UpperCAmelCase__	,    dim=0	)
               UpperCAmelCase				: Union[str, Any]											=	negative_image_embeds.repeat_interleave(UpperCAmelCase__	,    dim=0	)

        UpperCAmelCase				: Optional[Any]											=	torch.cat([negative_image_embeds, image_embeds]	,    dim=0	).to(
            dtype=prompt_embeds.dtype	,    device=UpperCAmelCase__	)

        self.scheduler.set_timesteps(UpperCAmelCase__	,    device=UpperCAmelCase__	)
        UpperCAmelCase				: int											=	self.scheduler.timesteps

        UpperCAmelCase				: List[Any]											=	self.unet.config.in_channels

        UpperCAmelCase							,	UpperCAmelCase				: Union[str, Any]											=	get_new_h_w(UpperCAmelCase__	,    UpperCAmelCase__	,    self.movq_scale_factor	)

        # create initial latent
        UpperCAmelCase				: Dict											=	self.prepare_latents(
            (batch_size, num_channels_latents, height, width)	,    text_encoder_hidden_states.dtype	,    UpperCAmelCase__	,    UpperCAmelCase__	,    UpperCAmelCase__	,    self.scheduler	,    )

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

               UpperCAmelCase				: int											=	{'''text_embeds''': prompt_embeds, '''image_embeds''': image_embeds}
               UpperCAmelCase				: Tuple											=	self.unet(
                   sample=UpperCAmelCase__	,    timestep=UpperCAmelCase__	,    encoder_hidden_states=UpperCAmelCase__	,    added_cond_kwargs=UpperCAmelCase__	,    return_dict=UpperCAmelCase__	,    )[0]

               if do_classifier_free_guidance:
                      UpperCAmelCase							,	UpperCAmelCase				: Any											=	noise_pred.split(latents.shape[1]	,    dim=1	)
                      UpperCAmelCase							,	UpperCAmelCase				: Optional[int]											=	noise_pred.chunk(2	)
                      UpperCAmelCase							,	UpperCAmelCase				: Any											=	variance_pred.chunk(2	)
                      UpperCAmelCase				: int											=	noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond)
                      UpperCAmelCase				: Dict											=	torch.cat([noise_pred, variance_pred_text]	,    dim=1	)

               if not (
                   hasattr(self.scheduler.config	,    '''variance_type'''	)
                   and self.scheduler.config.variance_type in ["learned", "learned_range"]
               ):
                      UpperCAmelCase							,	UpperCAmelCase				: int											=	noise_pred.split(latents.shape[1]	,    dim=1	)

               # compute the previous noisy sample x_t -> x_t-1
               UpperCAmelCase				: Any											=	self.scheduler.step(
                   UpperCAmelCase__	,    UpperCAmelCase__	,    UpperCAmelCase__	,    generator=UpperCAmelCase__	,    ).prev_sample
        # post-processing
        UpperCAmelCase				: str											=	self.movq.decode(UpperCAmelCase__	,    force_not_quantize=UpperCAmelCase__	)['''sample''']

        if output_type not in ["pt", "np", "pil"]:
               raise ValueError(F"""Only the output types `pt`, `pil` and `np` are supported not output_type={output_type}"""	)

        if output_type in ["np", "pil"]:
               UpperCAmelCase				: Optional[Any]											=	image * 0.5 + 0.5
               UpperCAmelCase				: Tuple											=	image.clamp(0	,    1	)
               UpperCAmelCase				: List[Any]											=	image.cpu().permute(0	,    2	,    3	,    1	).float().numpy()

        if output_type == "pil":
               UpperCAmelCase				: Optional[int]											=	self.numpy_to_pil(UpperCAmelCase__	)

        if not return_dict:
               return (image,)

        return ImagePipelineOutput(images=UpperCAmelCase__	)










'''simple docstring'''






from typing import List, Optional, Union

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


UpperCamelCase__:     Optional[Any]										=    logging.get_logger(__name__)

UpperCamelCase__:     Tuple										=    {
    "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(							A__    ):







	"""simple docstring"""






	lowerCamelCase__    				= """time_series_transformer"""
	lowerCamelCase__    				= {
	    """hidden_size""": """d_model""",
	    """num_attention_heads""": """encoder_attention_heads""",
	    """num_hidden_layers""": """encoder_layers""",
	}







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

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

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

								UpperCAmelCase				: Tuple											=	activation_function
								UpperCAmelCase				: Dict											=	init_std

								UpperCAmelCase				: Optional[int]											=	use_cache

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







	@property
	def   A			(		self    :  Union[str, Any]	)    ->	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
								)





