



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


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

lowerCamelCase  			=		'\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   a__      ( lowerCAmelCase__      ,     lowerCAmelCase__      ,     lowerCAmelCase__=8 ):
      UpperCAmelCase_     					=				h // scale_factor**2
      if h % scale_factor**2 != 0:
            new_h += 1
      UpperCAmelCase_     					=				w // scale_factor**2
      if w % scale_factor**2 != 0:
            new_w += 1
      return new_h * scale_factor, new_w * scale_factor


class  lowercase__  (       SCREAMING_SNAKE_CASE ):







       '''simple docstring'''


       def __init__(				self     :    Optional[int]						,							_UpperCAmelCase     :    Union[str, Any]						,							_UpperCAmelCase     :    Any						,							_UpperCAmelCase     :    Union[str, Any]						,							_UpperCAmelCase     :    List[Any]						,							_UpperCAmelCase     :    Union[str, Any]						,							)					->		Optional[Any]:


             '''simple docstring'''




             super().__init__()

             self.register_modules(
                 text_encoder=lowerCamelCase_						,							tokenizer=lowerCamelCase_						,							unet=lowerCamelCase_						,							scheduler=lowerCamelCase_						,							movq=lowerCamelCase_						,							)
             UpperCAmelCase_     					=				2 ** (len(self.movq.config.block_out_channels       ) - 1)



       def 	lowercase__							(				self     :    Optional[int]						,							_UpperCAmelCase     :    List[str]						,							_UpperCAmelCase     :    Optional[Any]						,							_UpperCAmelCase     :    int						,							_UpperCAmelCase     :    Dict						,							_UpperCAmelCase     :    Dict						,							_UpperCAmelCase     :    str       )					->		Any:


             '''simple docstring'''




             if latents is None:
                   UpperCAmelCase_     					=				randn_tensor(lowerCamelCase_						,							generator=lowerCamelCase_						,							device=lowerCamelCase_						,							dtype=lowerCamelCase_       )
             else:
                   if latents.shape != shape:
                         raise ValueError(F"""Unexpected latents shape, got {latents.shape}, expected {shape}"""       )
                   UpperCAmelCase_     					=				latents.to(lowerCamelCase_       )

             UpperCAmelCase_     					=				latents * scheduler.init_noise_sigma
             return latents



       def 	lowercase__							(				self     :    Any						,							_UpperCAmelCase     :    List[Any]						,							_UpperCAmelCase     :    Optional[int]						,							_UpperCAmelCase     :    Tuple						,							_UpperCAmelCase     :    Dict						,							_UpperCAmelCase     :    List[str]=None						,							)					->		List[Any]:


             '''simple docstring'''




             UpperCAmelCase_     					=				len(lowerCamelCase_       ) if isinstance(lowerCamelCase_						,							lowerCamelCase_       ) else 1
             # get prompt text embeddings
             UpperCAmelCase_     					=				self.tokenizer(
                 lowerCamelCase_						,							padding="max_length"						,							truncation=lowerCamelCase_						,							max_length=77						,							return_attention_mask=lowerCamelCase_						,							add_special_tokens=lowerCamelCase_						,							return_tensors="pt"						,							)

             UpperCAmelCase_     					=				text_inputs.input_ids
             UpperCAmelCase_     					=				self.tokenizer(lowerCamelCase_						,							padding="longest"						,							return_tensors="pt"       ).input_ids

             if untruncated_ids.shape[-1] >= text_input_ids.shape[-1] and not torch.equal(lowerCamelCase_						,							lowerCamelCase_       ):
                   UpperCAmelCase_     					=				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_     					=				text_input_ids.to(lowerCamelCase_       )
             UpperCAmelCase_     					=				text_inputs.attention_mask.to(lowerCamelCase_       )

             UpperCAmelCase_     					=				self.text_encoder(
                 input_ids=lowerCamelCase_						,							attention_mask=lowerCamelCase_       )

             UpperCAmelCase_     					=				prompt_embeds.repeat_interleave(lowerCamelCase_						,							dim=0       )
             UpperCAmelCase_     					=				text_encoder_hidden_states.repeat_interleave(lowerCamelCase_						,							dim=0       )
             UpperCAmelCase_     					=				text_mask.repeat_interleave(lowerCamelCase_						,							dim=0       )

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

                   UpperCAmelCase_     					=				self.tokenizer(
                       lowerCamelCase_						,							padding="max_length"						,							max_length=77						,							truncation=lowerCamelCase_						,							return_attention_mask=lowerCamelCase_						,							add_special_tokens=lowerCamelCase_						,							return_tensors="pt"						,							)
                   UpperCAmelCase_     					=				uncond_input.input_ids.to(lowerCamelCase_       )
                   UpperCAmelCase_     					=				uncond_input.attention_mask.to(lowerCamelCase_       )

                   UpperCAmelCase_     					=				self.text_encoder(
                       input_ids=lowerCamelCase_						,							attention_mask=lowerCamelCase_       )

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

                   UpperCAmelCase_     					=				negative_prompt_embeds.shape[1]
                   UpperCAmelCase_     					=				negative_prompt_embeds.repeat(1						,							lowerCamelCase_       )
                   UpperCAmelCase_     					=				negative_prompt_embeds.view(batch_size * num_images_per_prompt						,							lowerCamelCase_       )

                   UpperCAmelCase_     					=				uncond_text_encoder_hidden_states.shape[1]
                   UpperCAmelCase_     					=				uncond_text_encoder_hidden_states.repeat(1						,							lowerCamelCase_						,							1       )
                   UpperCAmelCase_     					=				uncond_text_encoder_hidden_states.view(
                       batch_size * num_images_per_prompt						,							lowerCamelCase_						,							-1       )
                   UpperCAmelCase_     					=				uncond_text_mask.repeat_interleave(lowerCamelCase_						,							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_     					=				torch.cat([negative_prompt_embeds, prompt_embeds]       )
                   UpperCAmelCase_     					=				torch.cat([uncond_text_encoder_hidden_states, text_encoder_hidden_states]       )

                   UpperCAmelCase_     					=				torch.cat([uncond_text_mask, text_mask]       )

             return prompt_embeds, text_encoder_hidden_states, text_mask



       def 	lowercase__							(				self     :    Tuple						,							_UpperCAmelCase     :    Optional[int]=0       )					->		List[str]:


             '''simple docstring'''




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

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

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



       def 	lowercase__							(				self     :    Dict						,							_UpperCAmelCase     :    Any=0       )					->		Union[str, Any]:


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

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

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

             UpperCAmelCase_     					=				None
             for cpu_offloaded_model in [self.text_encoder, self.unet, self.movq]:
                   UpperCAmelCase_     					=				cpu_offload_with_hook(lowerCamelCase_						,							lowerCamelCase_						,							prev_module_hook=lowerCamelCase_       )

             if self.safety_checker is not None:
                   UpperCAmelCase_     					=				cpu_offload_with_hook(self.safety_checker						,							lowerCamelCase_						,							prev_module_hook=lowerCamelCase_       )

             # We'll offload the last model manually.
             UpperCAmelCase_     					=				hook



       @property
       # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline._execution_device
       def 	lowercase__							(				self     :    int       )					->		Optional[Any]:


             '''simple docstring'''




             if not hasattr(self.unet						,							"_hf_hook"       ):
                   return self.device
             for module in self.unet.modules():
                   if (
                       hasattr(lowerCamelCase_						,							"_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(lowerCamelCase_       )
       def __call__(				self     :    Optional[int]						,							_UpperCAmelCase     :    List[Any]						,							_UpperCAmelCase     :    Any						,							_UpperCAmelCase     :    List[str]						,							_UpperCAmelCase     :    Union[str, Any] = None						,							_UpperCAmelCase     :    List[str] = 512						,							_UpperCAmelCase     :    List[str] = 512						,							_UpperCAmelCase     :    Optional[int] = 100						,							_UpperCAmelCase     :    Optional[int] = 4.0						,							_UpperCAmelCase     :    Optional[int] = 1						,							_UpperCAmelCase     :    Any = None						,							_UpperCAmelCase     :    Union[str, Any] = None						,							_UpperCAmelCase     :    Dict = "pil"						,							_UpperCAmelCase     :    Optional[int] = True						,							)					->		Optional[Any]:


             '''simple docstring'''




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

             UpperCAmelCase_     					=				self._execution_device

             UpperCAmelCase_     					=				batch_size * num_images_per_prompt
             UpperCAmelCase_     					=				guidance_scale > 1.0

             UpperCAmelCase_     					=				self._encode_prompt(
                 lowerCamelCase_						,							lowerCamelCase_						,							lowerCamelCase_						,							lowerCamelCase_						,							lowerCamelCase_       )

             if isinstance(lowerCamelCase_						,							lowerCamelCase_       ):
                   UpperCAmelCase_     					=				torch.cat(lowerCamelCase_						,							dim=0       )
             if isinstance(lowerCamelCase_						,							lowerCamelCase_       ):
                   UpperCAmelCase_     					=				torch.cat(lowerCamelCase_						,							dim=0       )

             if do_classifier_free_guidance:
                   UpperCAmelCase_     					=				image_embeds.repeat_interleave(lowerCamelCase_						,							dim=0       )
                   UpperCAmelCase_     					=				negative_image_embeds.repeat_interleave(lowerCamelCase_						,							dim=0       )

             UpperCAmelCase_     					=				torch.cat([negative_image_embeds, image_embeds]						,							dim=0       ).to(
                 dtype=prompt_embeds.dtype						,							device=lowerCamelCase_       )

             self.scheduler.set_timesteps(lowerCamelCase_						,							device=lowerCamelCase_       )
             UpperCAmelCase_     					=				self.scheduler.timesteps

             UpperCAmelCase_     					=				self.unet.config.in_channels

             UpperCAmelCase_     					=				get_new_h_w(lowerCamelCase_						,							lowerCamelCase_						,							self.movq_scale_factor       )

             # create initial latent
             UpperCAmelCase_     					=				self.prepare_latents(
                 (batch_size, num_channels_latents, height, width)						,							text_encoder_hidden_states.dtype						,							lowerCamelCase_						,							lowerCamelCase_						,							lowerCamelCase_						,							self.scheduler						,							)

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

                   UpperCAmelCase_     					=				{'''text_embeds''': prompt_embeds, '''image_embeds''': image_embeds}
                   UpperCAmelCase_     					=				self.unet(
                       sample=lowerCamelCase_						,							timestep=lowerCamelCase_						,							encoder_hidden_states=lowerCamelCase_						,							added_cond_kwargs=lowerCamelCase_						,							return_dict=lowerCamelCase_						,							)[0]

                   if do_classifier_free_guidance:
                         UpperCAmelCase_     					=				noise_pred.split(latents.shape[1]						,							dim=1       )
                         UpperCAmelCase_     					=				noise_pred.chunk(2       )
                         UpperCAmelCase_     					=				variance_pred.chunk(2       )
                         UpperCAmelCase_     					=				noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond)
                         UpperCAmelCase_     					=				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_     					=				noise_pred.split(latents.shape[1]						,							dim=1       )

                   # compute the previous noisy sample x_t -> x_t-1
                   UpperCAmelCase_     					=				self.scheduler.step(
                       lowerCamelCase_						,							lowerCamelCase_						,							lowerCamelCase_						,							generator=lowerCamelCase_						,							).prev_sample
             # post-processing
             UpperCAmelCase_     					=				self.movq.decode(lowerCamelCase_						,							force_not_quantize=lowerCamelCase_       )['''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_     					=				image * 0.5 + 0.5
                   UpperCAmelCase_     					=				image.clamp(0						,							1       )
                   UpperCAmelCase_     					=				image.cpu().permute(0						,							2						,							3						,							1       ).float().numpy()

             if output_type == "pil":
                   UpperCAmelCase_     					=				self.numpy_to_pil(lowerCamelCase_       )

             if not return_dict:
                   return (image,)

             return ImagePipelineOutput(images=lowerCamelCase_       )





'''simple docstring'''

from typing import TYPE_CHECKING

from ...utils import OptionalDependencyNotAvailable, _LazyModule, is_torch_available


UpperCamelCase__  :						str					=      {
    'configuration_git': ['GIT_PRETRAINED_CONFIG_ARCHIVE_MAP', 'GitConfig', 'GitVisionConfig'],
    'processing_git': ['GitProcessor'],
}

try:
  if not is_torch_available():
    raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
  pass
else:
  UpperCamelCase__  :						str					=      [
      'GIT_PRETRAINED_MODEL_ARCHIVE_LIST',
      'GitForCausalLM',
      'GitModel',
      'GitPreTrainedModel',
      'GitVisionModel',
  ]

if TYPE_CHECKING:
  from .configuration_git import GIT_PRETRAINED_CONFIG_ARCHIVE_MAP, GitConfig, GitVisionConfig
  from .processing_git import GitProcessor

  try:
    if not is_torch_available():
      raise OptionalDependencyNotAvailable()
  except OptionalDependencyNotAvailable:
    pass
  else:
    from .modeling_git import (
        GIT_PRETRAINED_MODEL_ARCHIVE_LIST,
        GitForCausalLM,
        GitModel,
        GitPreTrainedModel,
        GitVisionModel,
    )

else:
  import sys

  UpperCamelCase__  :						Tuple					=      _LazyModule(__name__, globals()['__file__'], _import_structure, module_spec=__spec__)




