



"""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


lowercase__					:						Tuple												=			logging.get_logger(__name__)  # pylint: disable=invalid-name

lowercase__					:						Any												=			'''
    Examples:
        ```py
        >>> from diffusers import KandinskyPipeline, KandinskyPriorPipeline
        >>> import torch

        >>> pipe_prior = KandinskyPriorPipeline.from_pretrained("kandinsky-community/Kandinsky-2-1-prior")
        >>> pipe_prior.to("cuda")

        >>> prompt = "red cat, 4k photo"
        >>> out = pipe_prior(prompt)
        >>> image_emb = out.image_embeds
        >>> negative_image_emb = out.negative_image_embeds

        >>> pipe = KandinskyPipeline.from_pretrained("kandinsky-community/kandinsky-2-1")
        >>> pipe.to("cuda")

        >>> image = pipe(
        ...     prompt,
        ...     image_embeds=image_emb,
        ...     negative_image_embeds=negative_image_emb,
        ...     height=768,
        ...     width=768,
        ...     num_inference_steps=100,
        ... ).images

        >>> image[0].save("cat.png")
        ```
'''




def __lowercase    (							_a	,    _a	,    _a=8					):
    snake_case_				:  Tuple       =			h // scale_factor**2
    if h % scale_factor**2 != 0:
        new_h += 1
    snake_case_				:  Optional[Any]       =			w // scale_factor**2
    if w % scale_factor**2 != 0:
        new_w += 1
    return new_h * scale_factor, new_w * scale_factor






class  _UpperCAmelCase (    lowerCAmelCase__):

 def __init__(  self	: Any						,  lowercase_	: MultilingualCLIP						,  lowercase_	: XLMRobertaTokenizer						,  lowercase_	: UNetaDConditionModel						,  lowercase_	: Union[DDIMScheduler, DDPMScheduler]						,  lowercase_	: VQModel						,  ):
     super().__init__()

     self.register_modules(
         text_encoder=lowercase_						,  tokenizer=lowercase_						,  unet=lowercase_						,  scheduler=lowercase_						,  movq=lowercase_						,  )
     snake_case_				:  Optional[Any]       =			2 ** (len(self.movq.config.block_out_channels       ) - 1)

 def 		_snake_case      (  self	: Union[str, Any]						,  lowercase_	: Optional[Any]						,  lowercase_	: int						,  lowercase_	: Optional[Any]						,  lowercase_	: str						,  lowercase_	: List[str]						,  lowercase_	: str       ):
     if latents is None:
         snake_case_				:  str       =			randn_tensor(lowercase_						,  generator=lowercase_						,  device=lowercase_						,  dtype=lowercase_       )
     else:
         if latents.shape != shape:
             raise ValueError(f"Unexpected latents shape, got {latents.shape}, expected {shape}"       )
         snake_case_				:  List[str]       =			latents.to(lowercase_       )

     snake_case_				:  Any       =			latents * scheduler.init_noise_sigma
     return latents

 def 		_snake_case      (  self	: str						,  lowercase_	: Any						,  lowercase_	: Dict						,  lowercase_	: Optional[Any]						,  lowercase_	: int						,  lowercase_	: Union[str, Any]=None						,  ):
     snake_case_				:  List[str]       =			len(lowercase_       ) if isinstance(lowercase_						,  lowercase_       ) else 1
     # get prompt text embeddings
     snake_case_				:  List[Any]       =			self.tokenizer(
         lowercase_						,  padding='''max_length'''						,  truncation=lowercase_						,  max_length=77						,  return_attention_mask=lowercase_						,  add_special_tokens=lowercase_						,  return_tensors='''pt'''						,  )

     snake_case_				:  List[str]       =			text_inputs.input_ids
     snake_case_				:  List[str]       =			self.tokenizer(lowercase_						,  padding='''longest'''						,  return_tensors='''pt'''       ).input_ids

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

     snake_case_				:  List[Any]       =			text_input_ids.to(lowercase_       )
     snake_case_				:  Union[str, Any]       =			text_inputs.attention_mask.to(lowercase_       )

     snake_case_,			snake_case_				:  Dict       =			self.text_encoder(
         input_ids=lowercase_						,  attention_mask=lowercase_       )

     snake_case_				:  Tuple       =			prompt_embeds.repeat_interleave(lowercase_						,  dim=0       )
     snake_case_				:  Dict       =			text_encoder_hidden_states.repeat_interleave(lowercase_						,  dim=0       )
     snake_case_				:  Any       =			text_mask.repeat_interleave(lowercase_						,  dim=0       )

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

         snake_case_				:  List[str]       =			self.tokenizer(
             lowercase_						,  padding='''max_length'''						,  max_length=77						,  truncation=lowercase_						,  return_attention_mask=lowercase_						,  add_special_tokens=lowercase_						,  return_tensors='''pt'''						,  )
         snake_case_				:  str       =			uncond_input.input_ids.to(lowercase_       )
         snake_case_				:  Tuple       =			uncond_input.attention_mask.to(lowercase_       )

         snake_case_,			snake_case_				:  Optional[Any]       =			self.text_encoder(
             input_ids=lowercase_						,  attention_mask=lowercase_       )

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

         snake_case_				:  Optional[Any]       =			negative_prompt_embeds.shape[1]
         snake_case_				:  Any       =			negative_prompt_embeds.repeat(1						,  lowercase_       )
         snake_case_				:  Union[str, Any]       =			negative_prompt_embeds.view(batch_size * num_images_per_prompt						,  lowercase_       )

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

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

     return prompt_embeds, text_encoder_hidden_states, text_mask

 def 		_snake_case      (  self	: List[Any]						,  lowercase_	: List[Any]=0       ):
     if is_accelerate_available():
         from accelerate import cpu_offload
     else:
         raise ImportError('''Please install accelerate via `pip install accelerate`'''       )

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

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

 def 		_snake_case      (  self	: List[str]						,  lowercase_	: Optional[Any]=0       ):
     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.'''       )

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

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

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

     if self.safety_checker is not None:
         snake_case_,			snake_case_				:  Optional[Any]       =			cpu_offload_with_hook(self.safety_checker						,  lowercase_						,  prev_module_hook=lowercase_       )

     # We'll offload the last model manually.
     snake_case_				:  List[str]       =			hook

 @property
 # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline._execution_device
 def 		_snake_case      (  self	: List[Any]       ):
     if not hasattr(self.unet						,  '''_hf_hook'''       ):
         return self.device
     for module in self.unet.modules():
         if (
             hasattr(lowercase_						,  '''_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(lowercase_       )
 def __call__(  self	: List[str]						,  lowercase_	: Union[str, List[str]]						,  lowercase_	: Union[torch.FloatTensor, List[torch.FloatTensor]]						,  lowercase_	: Union[torch.FloatTensor, List[torch.FloatTensor]]						,  lowercase_	: Optional[Union[str, List[str]]] = None						,  lowercase_	: int = 512						,  lowercase_	: int = 512						,  lowercase_	: int = 100						,  lowercase_	: float = 4.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						,  ):

     if isinstance(lowercase_						,  lowercase_       ):
         snake_case_				:  str       =			1
     elif isinstance(lowercase_						,  lowercase_       ):
         snake_case_				:  Dict       =			len(lowercase_       )
     else:
         raise ValueError(f"`prompt` has to be of type `str` or `list` but is {type(lowercase_       )}"       )

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

     snake_case_				:  Union[str, Any]       =			batch_size * num_images_per_prompt
     snake_case_				:  Tuple       =			guidance_scale > 1.0

     snake_case_,			snake_case_,			snake_case_				:  Optional[int]       =			self._encode_prompt(
         lowercase_						,  lowercase_						,  lowercase_						,  lowercase_						,  lowercase_       )

     if isinstance(lowercase_						,  lowercase_       ):
         snake_case_				:  str       =			torch.cat(lowercase_						,  dim=0       )
     if isinstance(lowercase_						,  lowercase_       ):
         snake_case_				:  Tuple       =			torch.cat(lowercase_						,  dim=0       )

     if do_classifier_free_guidance:
         snake_case_				:  Tuple       =			image_embeds.repeat_interleave(lowercase_						,  dim=0       )
         snake_case_				:  Dict       =			negative_image_embeds.repeat_interleave(lowercase_						,  dim=0       )

     snake_case_				:  List[Any]       =			torch.cat([negative_image_embeds, image_embeds]						,  dim=0       ).to(
         dtype=prompt_embeds.dtype						,  device=lowercase_       )

     self.scheduler.set_timesteps(lowercase_						,  device=lowercase_       )
     snake_case_				:  Union[str, Any]       =			self.scheduler.timesteps

     snake_case_				:  Any       =			self.unet.config.in_channels

     snake_case_,			snake_case_				:  int       =			get_new_h_w(lowercase_						,  lowercase_						,  self.movq_scale_factor       )

     # create initial latent
     snake_case_				:  Any       =			self.prepare_latents(
         (batch_size, num_channels_latents, height, width)						,  text_encoder_hidden_states.dtype						,  lowercase_						,  lowercase_						,  lowercase_						,  self.scheduler						,  )

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

         snake_case_				:  str       =			{'''text_embeds''': prompt_embeds, '''image_embeds''': image_embeds}
         snake_case_				:  int       =			self.unet(
             sample=lowercase_						,  timestep=lowercase_						,  encoder_hidden_states=lowercase_						,  added_cond_kwargs=lowercase_						,  return_dict=lowercase_						,  )[0]

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

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

     if output_type == "pil":
         snake_case_				:  List[str]       =			self.numpy_to_pil(lowercase_       )

     if not return_dict:
         return (image,)

     return ImagePipelineOutput(images=lowercase_       )










"""simple docstring"""

lowercase__					:						Union[str, Any]												=			{
    '''Pillow''': '''Pillow''',
    '''accelerate''': '''accelerate>=0.11.0''',
    '''compel''': '''compel==0.1.8''',
    '''black''': '''black~=23.1''',
    '''datasets''': '''datasets''',
    '''filelock''': '''filelock''',
    '''flax''': '''flax>=0.4.1''',
    '''hf-doc-builder''': '''hf-doc-builder>=0.3.0''',
    '''huggingface-hub''': '''huggingface-hub>=0.13.2''',
    '''requests-mock''': '''requests-mock==1.10.0''',
    '''importlib_metadata''': '''importlib_metadata''',
    '''invisible-watermark''': '''invisible-watermark''',
    '''isort''': '''isort>=5.5.4''',
    '''jax''': '''jax>=0.2.8,!=0.3.2''',
    '''jaxlib''': '''jaxlib>=0.1.65''',
    '''Jinja2''': '''Jinja2''',
    '''k-diffusion''': '''k-diffusion>=0.0.12''',
    '''torchsde''': '''torchsde''',
    '''note_seq''': '''note_seq''',
    '''librosa''': '''librosa''',
    '''numpy''': '''numpy''',
    '''omegaconf''': '''omegaconf''',
    '''parameterized''': '''parameterized''',
    '''protobuf''': '''protobuf>=3.20.3,<4''',
    '''pytest''': '''pytest''',
    '''pytest-timeout''': '''pytest-timeout''',
    '''pytest-xdist''': '''pytest-xdist''',
    '''ruff''': '''ruff>=0.0.241''',
    '''safetensors''': '''safetensors''',
    '''sentencepiece''': '''sentencepiece>=0.1.91,!=0.1.92''',
    '''scipy''': '''scipy''',
    '''onnx''': '''onnx''',
    '''regex''': '''regex!=2019.12.17''',
    '''requests''': '''requests''',
    '''tensorboard''': '''tensorboard''',
    '''torch''': '''torch>=1.4''',
    '''torchvision''': '''torchvision''',
    '''transformers''': '''transformers>=4.25.1''',
    '''urllib3''': '''urllib3<=2.0.0''',
}





