

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 			=		logging.get_logger(__name__)  # pylint: disable=invalid-name

_UpperCamelCase 			=		"\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 		_lowercase (     lowercase__   ,	lowercase__   ,	lowercase__=8						):
      __lowerCAmelCase					:						Tuple					=					h // scale_factor**2
      if h % scale_factor**2 != 0:
            new_h += 1
      __lowerCAmelCase					:						List[str]					=					w // scale_factor**2
      if w % scale_factor**2 != 0:
            new_w += 1
      return new_h * scale_factor, new_w * scale_factor







class   __lowercase      (_UpperCAmelCase ):
 def __init__(							self      ,							A_      ,							A_      ,							A_      ,							A_      ,							A_      ,							)					->str:



       '''simple docstring'''





       super().__init__()

       self.register_modules(
           text_encoder=A_      ,							tokenizer=A_      ,							unet=A_      ,							scheduler=A_      ,							movq=A_      ,							)
       __lowerCAmelCase					:						str					=					2 ** (len(self.movq.config.block_out_channels					) - 1)
 def 	UpperCamelCase__  (							self      ,							A_      ,							A_      ,							A_      ,							A_      ,							A_      ,							A_					)					->List[Any]:



       '''simple docstring'''





       if latents is None:
             __lowerCAmelCase					:						int					=					randn_tensor(A_      ,							generator=A_      ,							device=A_      ,							dtype=A_					)
       else:
             if latents.shape != shape:
                   raise ValueError(f"""Unexpected latents shape, got {latents.shape}, expected {shape}"""					)
             __lowerCAmelCase					:						List[Any]					=					latents.to(A_					)

       __lowerCAmelCase					:						Any					=					latents * scheduler.init_noise_sigma
       return latents
 def 	UpperCamelCase__  (							self      ,							A_      ,							A_      ,							A_      ,							A_      ,							A_=None      ,							)					->Union[str, Any]:



       '''simple docstring'''





       __lowerCAmelCase					:						List[str]					=					len(A_					) if isinstance(A_      ,							A_					) else 1
       # get prompt text embeddings
       __lowerCAmelCase					:						str					=					self.tokenizer(
           A_      ,							padding='''max_length'''      ,							truncation=A_      ,							max_length=77      ,							return_attention_mask=A_      ,							add_special_tokens=A_      ,							return_tensors='''pt'''      ,							)

       __lowerCAmelCase					:						Optional[Any]					=					text_inputs.input_ids
       __lowerCAmelCase					:						Dict					=					self.tokenizer(A_      ,							padding='''longest'''      ,							return_tensors='''pt'''					).input_ids

       if untruncated_ids.shape[-1] >= text_input_ids.shape[-1] and not torch.equal(A_      ,							A_					):
             __lowerCAmelCase					:						str					=					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}"""					)

       __lowerCAmelCase					:						List[Any]					=					text_input_ids.to(A_					)
       __lowerCAmelCase					:						Tuple					=					text_inputs.attention_mask.to(A_					)

       __lowerCAmelCase,							__lowerCAmelCase					:						int					=					self.text_encoder(
           input_ids=A_      ,							attention_mask=A_					)

       __lowerCAmelCase					:						Dict					=					prompt_embeds.repeat_interleave(A_      ,							dim=0					)
       __lowerCAmelCase					:						str					=					text_encoder_hidden_states.repeat_interleave(A_      ,							dim=0					)
       __lowerCAmelCase					:						Dict					=					text_mask.repeat_interleave(A_      ,							dim=0					)

       if do_classifier_free_guidance:
             __lowerCAmelCase					:						List[str]
             if negative_prompt is None:
                   __lowerCAmelCase					:						int					=					[''''''] * batch_size
             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					:						Tuple					=					[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					:						str					=					negative_prompt

             __lowerCAmelCase					:						int					=					self.tokenizer(
                 A_      ,							padding='''max_length'''      ,							max_length=77      ,							truncation=A_      ,							return_attention_mask=A_      ,							add_special_tokens=A_      ,							return_tensors='''pt'''      ,							)
             __lowerCAmelCase					:						int					=					uncond_input.input_ids.to(A_					)
             __lowerCAmelCase					:						Union[str, Any]					=					uncond_input.attention_mask.to(A_					)

             __lowerCAmelCase,							__lowerCAmelCase					:						Union[str, Any]					=					self.text_encoder(
                 input_ids=A_      ,							attention_mask=A_					)

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

             __lowerCAmelCase					:						List[Any]					=					negative_prompt_embeds.shape[1]
             __lowerCAmelCase					:						str					=					negative_prompt_embeds.repeat(1      ,							A_					)
             __lowerCAmelCase					:						Union[str, Any]					=					negative_prompt_embeds.view(batch_size * num_images_per_prompt      ,							A_					)

             __lowerCAmelCase					:						Optional[Any]					=					uncond_text_encoder_hidden_states.shape[1]
             __lowerCAmelCase					:						Union[str, Any]					=					uncond_text_encoder_hidden_states.repeat(1      ,							A_      ,							1					)
             __lowerCAmelCase					:						Optional[int]					=					uncond_text_encoder_hidden_states.view(
                 batch_size * num_images_per_prompt      ,							A_      ,							-1					)
             __lowerCAmelCase					:						Union[str, Any]					=					uncond_text_mask.repeat_interleave(A_      ,							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
             __lowerCAmelCase					:						Union[str, Any]					=					torch.cat([negative_prompt_embeds, prompt_embeds]					)
             __lowerCAmelCase					:						str					=					torch.cat([uncond_text_encoder_hidden_states, text_encoder_hidden_states]					)

             __lowerCAmelCase					:						Optional[int]					=					torch.cat([uncond_text_mask, text_mask]					)

       return prompt_embeds, text_encoder_hidden_states, text_mask
 def 	UpperCamelCase__  (							self      ,							A_=0					)					->int:



       '''simple docstring'''





       if is_accelerate_available():
             from accelerate import cpu_offload
       else:
             raise ImportError('''Please install accelerate via `pip install accelerate`'''					)

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

       __lowerCAmelCase					:						Dict					=					[
           self.unet,
           self.text_encoder,
           self.movq,
       ]
       for cpu_offloaded_model in models:
             if cpu_offloaded_model is not None:
                   cpu_offload(A_      ,							A_					)
 def 	UpperCamelCase__  (							self      ,							A_=0					)					->Optional[int]:



       '''simple docstring'''





       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.'''					)

       __lowerCAmelCase					:						Tuple					=					torch.device(f"""cuda:{gpu_id}"""					)

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

       __lowerCAmelCase					:						Union[str, Any]					=					None
       for cpu_offloaded_model in [self.text_encoder, self.unet, self.movq]:
             __lowerCAmelCase,							__lowerCAmelCase					:						Tuple					=					cpu_offload_with_hook(A_      ,							A_      ,							prev_module_hook=A_					)

       if self.safety_checker is not None:
             __lowerCAmelCase,							__lowerCAmelCase					:						List[Any]					=					cpu_offload_with_hook(self.safety_checker      ,							A_      ,							prev_module_hook=A_					)

       # We'll offload the last model manually.
       __lowerCAmelCase					:						Optional[int]					=					hook
 @property
 # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline._execution_device
 def 	UpperCamelCase__  (							self					)					->str:



       '''simple docstring'''





       if not hasattr(self.unet      ,							'''_hf_hook'''					):
             return self.device
       for module in self.unet.modules():
             if (
                 hasattr(A_      ,							'''_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(A_					)
 def __call__(							self      ,							A_      ,							A_      ,							A_      ,							A_ = None      ,							A_ = 512      ,							A_ = 512      ,							A_ = 100      ,							A_ = 4.0      ,							A_ = 1      ,							A_ = None      ,							A_ = None      ,							A_ = "pil"      ,							A_ = True      ,							)					->Dict:



       '''simple docstring'''





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

       __lowerCAmelCase					:						Union[str, Any]					=					self._execution_device

       __lowerCAmelCase					:						Any					=					batch_size * num_images_per_prompt
       __lowerCAmelCase					:						Tuple					=					guidance_scale > 1.0

       __lowerCAmelCase,							__lowerCAmelCase,							__lowerCAmelCase					:						List[Any]					=					self._encode_prompt(
           A_      ,							A_      ,							A_      ,							A_      ,							A_					)

       if isinstance(A_      ,							A_					):
             __lowerCAmelCase					:						Dict					=					torch.cat(A_      ,							dim=0					)
       if isinstance(A_      ,							A_					):
             __lowerCAmelCase					:						Dict					=					torch.cat(A_      ,							dim=0					)

       if do_classifier_free_guidance:
             __lowerCAmelCase					:						Dict					=					image_embeds.repeat_interleave(A_      ,							dim=0					)
             __lowerCAmelCase					:						Optional[Any]					=					negative_image_embeds.repeat_interleave(A_      ,							dim=0					)

       __lowerCAmelCase					:						Optional[int]					=					torch.cat([negative_image_embeds, image_embeds]      ,							dim=0					).to(
           dtype=prompt_embeds.dtype      ,							device=A_					)

       self.scheduler.set_timesteps(A_      ,							device=A_					)
       __lowerCAmelCase					:						Any					=					self.scheduler.timesteps

       __lowerCAmelCase					:						Union[str, Any]					=					self.unet.config.in_channels

       __lowerCAmelCase,							__lowerCAmelCase					:						Optional[int]					=					get_new_h_w(A_      ,							A_      ,							self.movq_scale_factor					)

       # create initial latent
       __lowerCAmelCase					:						Optional[Any]					=					self.prepare_latents(
           (batch_size, num_channels_latents, height, width)      ,							text_encoder_hidden_states.dtype      ,							A_      ,							A_      ,							A_      ,							self.scheduler      ,							)

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

             __lowerCAmelCase					:						List[str]					=					{'''text_embeds''': prompt_embeds, '''image_embeds''': image_embeds}
             __lowerCAmelCase					:						Any					=					self.unet(
                 sample=A_      ,							timestep=A_      ,							encoder_hidden_states=A_      ,							added_cond_kwargs=A_      ,							return_dict=A_      ,							)[0]

             if do_classifier_free_guidance:
                   __lowerCAmelCase,							__lowerCAmelCase					:						Dict					=					noise_pred.split(latents.shape[1]      ,							dim=1					)
                   __lowerCAmelCase,							__lowerCAmelCase					:						Optional[Any]					=					noise_pred.chunk(2					)
                   __lowerCAmelCase,							__lowerCAmelCase					:						str					=					variance_pred.chunk(2					)
                   __lowerCAmelCase					:						Tuple					=					noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond)
                   __lowerCAmelCase					:						Optional[Any]					=					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"]
             ):
                   __lowerCAmelCase,							__lowerCAmelCase					:						Union[str, Any]					=					noise_pred.split(latents.shape[1]      ,							dim=1					)

             # compute the previous noisy sample x_t -> x_t-1
             __lowerCAmelCase					:						int					=					self.scheduler.step(
                 A_      ,							A_      ,							A_      ,							generator=A_      ,							).prev_sample
       # post-processing
       __lowerCAmelCase					:						List[str]					=					self.movq.decode(A_      ,							force_not_quantize=A_					)['''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"]:
             __lowerCAmelCase					:						Dict					=					image * 0.5 + 0.5
             __lowerCAmelCase					:						int					=					image.clamp(0      ,							1					)
             __lowerCAmelCase					:						Dict					=					image.cpu().permute(0      ,							2      ,							3      ,							1					).float().numpy()

       if output_type == "pil":
             __lowerCAmelCase					:						Dict					=					self.numpy_to_pil(A_					)

       if not return_dict:
             return (image,)

       return ImagePipelineOutput(images=A_					)





from dataclasses import dataclass
from typing import Optional, Tuple, Union

import numpy as np
import torch

from ..configuration_utils import ConfigMixin, register_to_config
from ..utils import BaseOutput, randn_tensor
from .scheduling_utils import SchedulerMixin







@dataclass
class   __lowercase      (_UpperCAmelCase ):
 _UpperCamelCase			  =  42
 _UpperCamelCase			  =  42
 _UpperCamelCase			  =  None







class   __lowercase      (_UpperCAmelCase		,   _UpperCAmelCase ):
 _UpperCamelCase			  =  2
 @register_to_config
 def __init__(							self      ,							A_ = 0.02      ,							A_ = 100      ,							A_ = 1.007      ,							A_ = 80      ,							A_ = 0.05      ,							A_ = 50      ,							)					->int:



       '''simple docstring'''





       __lowerCAmelCase					:						Optional[int]					=					sigma_max

       # setable values
       __lowerCAmelCase					:						int					=					None
       __lowerCAmelCase					:						np.IntTensor					=					None
       __lowerCAmelCase					:						torch.FloatTensor					=					None  # sigma(t_i)
 def 	UpperCamelCase__  (							self      ,							A_      ,							A_ = None					)					->torch.FloatTensor:



       '''simple docstring'''





       return sample
 def 	UpperCamelCase__  (							self      ,							A_      ,							A_ = None					)					->List[str]:



       '''simple docstring'''





       __lowerCAmelCase					:						str					=					num_inference_steps
       __lowerCAmelCase					:						Dict					=					np.arange(0      ,							self.num_inference_steps					)[::-1].copy()
       __lowerCAmelCase					:						Optional[Any]					=					torch.from_numpy(A_					).to(A_					)
       __lowerCAmelCase					:						Tuple					=					[
           (
               self.config.sigma_max**2
               * (self.config.sigma_min**2 / self.config.sigma_max**2) ** (i / (num_inference_steps - 1))
           )
           for i in self.timesteps
       ]
       __lowerCAmelCase					:						Optional[int]					=					torch.tensor(A_      ,							dtype=torch.floataa      ,							device=A_					)
 def 	UpperCamelCase__  (							self      ,							A_      ,							A_      ,							A_ = None					)					->Tuple[torch.FloatTensor, float]:



       '''simple docstring'''





       if self.config.s_min <= sigma <= self.config.s_max:
             __lowerCAmelCase					:						List[str]					=					min(self.config.s_churn / self.num_inference_steps      ,							2**0.5 - 1					)
       else:
             __lowerCAmelCase					:						List[str]					=					0

       # sample eps ~ N(0, S_noise^2 * I)
       __lowerCAmelCase					:						int					=					self.config.s_noise * randn_tensor(sample.shape      ,							generator=A_					).to(sample.device					)
       __lowerCAmelCase					:						str					=					sigma + gamma * sigma
       __lowerCAmelCase					:						Any					=					sample + ((sigma_hat**2 - sigma**2) ** 0.5 * eps)

       return sample_hat, sigma_hat
 def 	UpperCamelCase__  (							self      ,							A_      ,							A_      ,							A_      ,							A_      ,							A_ = True      ,							)					->Union[KarrasVeOutput, Tuple]:



       '''simple docstring'''





       __lowerCAmelCase					:						Union[str, Any]					=					sample_hat + sigma_hat * model_output
       __lowerCAmelCase					:						int					=					(sample_hat - pred_original_sample) / sigma_hat
       __lowerCAmelCase					:						Tuple					=					sample_hat + (sigma_prev - sigma_hat) * derivative

       if not return_dict:
             return (sample_prev, derivative)

       return KarrasVeOutput(
           prev_sample=A_      ,							derivative=A_      ,							pred_original_sample=A_					)
 def 	UpperCamelCase__  (							self      ,							A_      ,							A_      ,							A_      ,							A_      ,							A_      ,							A_      ,							A_ = True      ,							)					->Union[KarrasVeOutput, Tuple]:



       '''simple docstring'''





       __lowerCAmelCase					:						str					=					sample_prev + sigma_prev * model_output
       __lowerCAmelCase					:						List[Any]					=					(sample_prev - pred_original_sample) / sigma_prev
       __lowerCAmelCase					:						Any					=					sample_hat + (sigma_prev - sigma_hat) * (0.5 * derivative + 0.5 * derivative_corr)

       if not return_dict:
             return (sample_prev, derivative)

       return KarrasVeOutput(
           prev_sample=A_      ,							derivative=A_      ,							pred_original_sample=A_					)





 def 	UpperCamelCase__  (							self      ,							A_      ,							A_      ,							A_					)					->Any:



       '''simple docstring'''





       raise NotImplementedError()


