

from typing import List, Optional, Union

import numpy as np
import PIL
import torch
from PIL import Image

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


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

UpperCAmelCase__ :     Union[str, Any]				    =    """
    Examples:
        ```py
        >>> from diffusers import KandinskyV22Img2ImgPipeline, KandinskyV22PriorPipeline
        >>> from diffusers.utils import load_image
        >>> import torch

        >>> pipe_prior = KandinskyV22PriorPipeline.from_pretrained(
        ...     \"kandinsky-community/kandinsky-2-2-prior\", torch_dtype=torch.float16
        ... )
        >>> pipe_prior.to(\"cuda\")

        >>> prompt = \"A red cartoon frog, 4k\"
        >>> image_emb, zero_image_emb = pipe_prior(prompt, return_dict=False)

        >>> pipe = KandinskyV22Img2ImgPipeline.from_pretrained(
        ...     \"kandinsky-community/kandinsky-2-2-decoder\", torch_dtype=torch.float16
        ... )
        >>> pipe.to(\"cuda\")

        >>> init_image = load_image(
        ...     \"https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main\"
        ...     \"/kandinsky/frog.png\"
        ... )

        >>> image = pipe(
        ...     image=init_image,
        ...     image_embeds=image_emb,
        ...     negative_image_embeds=zero_image_emb,
        ...     height=768,
        ...     width=768,
        ...     num_inference_steps=100,
        ...     strength=0.2,
        ... ).images

        >>> image[0].save(\"red_frog.png\")
        ```
"""






def 					__lowercase      (     _A						,    _A						,    _A=8			)     ->						Any:
      SCREAMING_SNAKE_CASE :      List[str]       =   height // scale_factor**2
      if height % scale_factor**2 != 0:
            new_height += 1
      SCREAMING_SNAKE_CASE :      Optional[Any]       =   width // scale_factor**2
      if width % scale_factor**2 != 0:
            new_width += 1
      return new_height * scale_factor, new_width * scale_factor






def 					__lowercase      (     _A						,    _A=512						,    _A=512			)     ->						Tuple:
      SCREAMING_SNAKE_CASE :      Any       =   pil_image.resize((w, h)						,    resample=Image.BICUBIC						,    reducing_gap=1			)
      SCREAMING_SNAKE_CASE :      int       =   np.array(pil_image.convert("""RGB"""			)			)
      SCREAMING_SNAKE_CASE :      List[str]       =   arr.astype(np.floataa			) / 127.5 - 1
      SCREAMING_SNAKE_CASE :      Tuple       =   np.transpose(_A						,    [2, 0, 1]			)
      SCREAMING_SNAKE_CASE :      Optional[Any]       =   torch.from_numpy(_A			).unsqueeze(0			)
      return image







class        a__       (   UpperCAmelCase       ):







  """simple docstring"""




  def __init__(			self    : List[Any]				,				UpperCAmelCase__    : UNetaDConditionModel				,				UpperCAmelCase__    : DDPMScheduler				,				UpperCAmelCase__    : VQModel				,				)					->List[str]:
        """simple docstring"""







        super().__init__()

        self.register_modules(
            unet=UpperCAmelCase__				,				scheduler=UpperCAmelCase__				,				movq=UpperCAmelCase__				,				)
        SCREAMING_SNAKE_CASE :      List[str]       =   2 ** (len(self.movq.config.block_out_channels		) - 1)


  def  _lowercase      (			self    : List[Any]				,				UpperCAmelCase__    : List[Any]				,				UpperCAmelCase__    : Optional[int]				,				UpperCAmelCase__    : Any		)					->Union[str, Any]:
        """simple docstring"""







        SCREAMING_SNAKE_CASE :      Any       =   min(int(num_inference_steps * strength		)				,				UpperCAmelCase__		)

        SCREAMING_SNAKE_CASE :      str       =   max(num_inference_steps - init_timestep				,				0		)
        SCREAMING_SNAKE_CASE :      Union[str, Any]       =   self.scheduler.timesteps[t_start:]

        return timesteps, num_inference_steps - t_start


  def  _lowercase      (			self    : Dict				,				UpperCAmelCase__    : Union[str, Any]				,				UpperCAmelCase__    : Union[str, Any]				,				UpperCAmelCase__    : Optional[int]				,				UpperCAmelCase__    : Any				,				UpperCAmelCase__    : Optional[int]				,				UpperCAmelCase__    : Optional[int]				,				UpperCAmelCase__    : Optional[Any]=None		)					->Optional[Any]:
        """simple docstring"""







        if not isinstance(UpperCAmelCase__				,				(torch.Tensor, PIL.Image.Image, list)		):
              raise ValueError(
                  f"`image` has to be of type `torch.Tensor`, `PIL.Image.Image` or list but is {type(UpperCAmelCase__		)}"		)

        SCREAMING_SNAKE_CASE :      int       =   image.to(device=UpperCAmelCase__				,				dtype=UpperCAmelCase__		)

        SCREAMING_SNAKE_CASE :      Union[str, Any]       =   batch_size * num_images_per_prompt

        if image.shape[1] == 4:
              SCREAMING_SNAKE_CASE :      Union[str, Any]       =   image

        else:
              if isinstance(UpperCAmelCase__				,				UpperCAmelCase__		) and len(UpperCAmelCase__		) != batch_size:
                    raise ValueError(
                        f"You have passed a list of generators of length {len(UpperCAmelCase__		)}, but requested an effective batch"
                        f" size of {batch_size}. Make sure the batch size matches the length of the generators."		)

              elif isinstance(UpperCAmelCase__				,				UpperCAmelCase__		):
                    SCREAMING_SNAKE_CASE :      Any       =   [
                        self.movq.encode(image[i : i + 1]		).latent_dist.sample(generator[i]		) for i in range(UpperCAmelCase__		)
                    ]
                    SCREAMING_SNAKE_CASE :      List[str]       =   torch.cat(UpperCAmelCase__				,				dim=0		)
              else:
                    SCREAMING_SNAKE_CASE :      Optional[int]       =   self.movq.encode(UpperCAmelCase__		).latent_dist.sample(UpperCAmelCase__		)

              SCREAMING_SNAKE_CASE :      Union[str, Any]       =   self.movq.config.scaling_factor * init_latents

        SCREAMING_SNAKE_CASE :      Optional[int]       =   torch.cat([init_latents]				,				dim=0		)

        SCREAMING_SNAKE_CASE :      Optional[int]       =   init_latents.shape
        SCREAMING_SNAKE_CASE :      Dict       =   randn_tensor(UpperCAmelCase__				,				generator=UpperCAmelCase__				,				device=UpperCAmelCase__				,				dtype=UpperCAmelCase__		)

        # get latents
        SCREAMING_SNAKE_CASE :      List[str]       =   self.scheduler.add_noise(UpperCAmelCase__				,				UpperCAmelCase__				,				UpperCAmelCase__		)

        SCREAMING_SNAKE_CASE :      Union[str, Any]       =   init_latents

        return latents


  def  _lowercase      (			self    : int				,				UpperCAmelCase__    : List[Any]=0		)					->int:
        """simple docstring"""







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

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

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


  def  _lowercase      (			self    : Any				,				UpperCAmelCase__    : str=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."""		)

        SCREAMING_SNAKE_CASE :      List[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)

        SCREAMING_SNAKE_CASE :      int       =   None
        for cpu_offloaded_model in [self.unet, self.movq]:
              SCREAMING_SNAKE_CASE     ,	SCREAMING_SNAKE_CASE :      List[Any]       =   cpu_offload_with_hook(UpperCAmelCase__				,				UpperCAmelCase__				,				prev_module_hook=UpperCAmelCase__		)

        # We'll offload the last model manually.
        SCREAMING_SNAKE_CASE :      Any       =   hook


  @property
  # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline._execution_device
  def  _lowercase      (			self    : List[Any]		)					->int:
        """simple docstring"""







        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    : List[str]				,				UpperCAmelCase__    : Union[torch.FloatTensor, List[torch.FloatTensor]]				,				UpperCAmelCase__    : Union[torch.FloatTensor, PIL.Image.Image, List[torch.FloatTensor], List[PIL.Image.Image]]				,				UpperCAmelCase__    : Union[torch.FloatTensor, List[torch.FloatTensor]]				,				UpperCAmelCase__    : int = 5_1_2				,				UpperCAmelCase__    : int = 5_1_2				,				UpperCAmelCase__    : int = 1_0_0				,				UpperCAmelCase__    : float = 4.0				,				UpperCAmelCase__    : float = 0.3				,				UpperCAmelCase__    : int = 1				,				UpperCAmelCase__    : Optional[Union[torch.Generator, List[torch.Generator]]] = None				,				UpperCAmelCase__    : Optional[str] = "pil"				,				UpperCAmelCase__    : bool = True				,				)					->Tuple:
        """simple docstring"""







        SCREAMING_SNAKE_CASE :      Any       =   self._execution_device

        SCREAMING_SNAKE_CASE :      Tuple       =   guidance_scale > 1.0

        if isinstance(UpperCAmelCase__				,				UpperCAmelCase__		):
              SCREAMING_SNAKE_CASE :      List[Any]       =   torch.cat(UpperCAmelCase__				,				dim=0		)
        SCREAMING_SNAKE_CASE :      Optional[int]       =   image_embeds.shape[0]
        if isinstance(UpperCAmelCase__				,				UpperCAmelCase__		):
              SCREAMING_SNAKE_CASE :      int       =   torch.cat(UpperCAmelCase__				,				dim=0		)

        if do_classifier_free_guidance:
              SCREAMING_SNAKE_CASE :      List[Any]       =   image_embeds.repeat_interleave(UpperCAmelCase__				,				dim=0		)
              SCREAMING_SNAKE_CASE :      Dict       =   negative_image_embeds.repeat_interleave(UpperCAmelCase__				,				dim=0		)

        SCREAMING_SNAKE_CASE :      str       =   torch.cat([negative_image_embeds, image_embeds]				,				dim=0		).to(dtype=self.unet.dtype				,				device=UpperCAmelCase__		)

        if not isinstance(UpperCAmelCase__				,				UpperCAmelCase__		):
              SCREAMING_SNAKE_CASE :      int       =   [image]
        if not all(isinstance(UpperCAmelCase__				,				(PIL.Image.Image, torch.Tensor)		) for i in image		):
              raise ValueError(
                  f"Input is in incorrect format: {[type(UpperCAmelCase__		) for i in image]}. Currently, we only support  PIL image and pytorch tensor"		)

        SCREAMING_SNAKE_CASE :      Union[str, Any]       =   torch.cat([prepare_image(UpperCAmelCase__				,				UpperCAmelCase__				,				UpperCAmelCase__		) for i in image]				,				dim=0		)
        SCREAMING_SNAKE_CASE :      Union[str, Any]       =   image.to(dtype=image_embeds.dtype				,				device=UpperCAmelCase__		)

        SCREAMING_SNAKE_CASE :      List[Any]       =   self.movq.encode(UpperCAmelCase__		)["""latents"""]
        SCREAMING_SNAKE_CASE :      Optional[int]       =   latents.repeat_interleave(UpperCAmelCase__				,				dim=0		)
        self.scheduler.set_timesteps(UpperCAmelCase__				,				device=UpperCAmelCase__		)
        SCREAMING_SNAKE_CASE     ,	SCREAMING_SNAKE_CASE :      str       =   self.get_timesteps(UpperCAmelCase__				,				UpperCAmelCase__				,				UpperCAmelCase__		)
        SCREAMING_SNAKE_CASE :      Union[str, Any]       =   timesteps[:1].repeat(batch_size * num_images_per_prompt		)
        SCREAMING_SNAKE_CASE     ,	SCREAMING_SNAKE_CASE :      Optional[Any]       =   downscale_height_and_width(UpperCAmelCase__				,				UpperCAmelCase__				,				self.movq_scale_factor		)
        SCREAMING_SNAKE_CASE :      Dict       =   self.prepare_latents(
            UpperCAmelCase__				,				UpperCAmelCase__				,				UpperCAmelCase__				,				UpperCAmelCase__				,				image_embeds.dtype				,				UpperCAmelCase__				,				UpperCAmelCase__		)
        for i, t in enumerate(self.progress_bar(UpperCAmelCase__		)		):
              # expand the latents if we are doing classifier free guidance
              SCREAMING_SNAKE_CASE :      List[Any]       =   torch.cat([latents] * 2		) if do_classifier_free_guidance else latents

              SCREAMING_SNAKE_CASE :      Optional[int]       =   {"""image_embeds""": image_embeds}
              SCREAMING_SNAKE_CASE :      Dict       =   self.unet(
                  sample=UpperCAmelCase__				,				timestep=UpperCAmelCase__				,				encoder_hidden_states=UpperCAmelCase__				,				added_cond_kwargs=UpperCAmelCase__				,				return_dict=UpperCAmelCase__				,				)[0]

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

              # compute the previous noisy sample x_t -> x_t-1
              SCREAMING_SNAKE_CASE :      Tuple       =   self.scheduler.step(
                  UpperCAmelCase__				,				UpperCAmelCase__				,				UpperCAmelCase__				,				generator=UpperCAmelCase__				,				)[0]

        # post-processing
        SCREAMING_SNAKE_CASE :      Any       =   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"]:
              SCREAMING_SNAKE_CASE :      int       =   image * 0.5 + 0.5
              SCREAMING_SNAKE_CASE :      Any       =   image.clamp(0				,				1		)
              SCREAMING_SNAKE_CASE :      int       =   image.cpu().permute(0				,				2				,				3				,				1		).float().numpy()

        if output_type == "pil":
              SCREAMING_SNAKE_CASE :      Tuple       =   self.numpy_to_pil(UpperCAmelCase__		)

        if not return_dict:
              return (image,)

        return ImagePipelineOutput(images=UpperCAmelCase__		)







from typing import TYPE_CHECKING

from ...utils import (
    OptionalDependencyNotAvailable,
    _LazyModule,
    is_flax_available,
    is_tf_available,
    is_torch_available,
    is_vision_available,
)


UpperCAmelCase__ :     List[Any]				    =    {"""configuration_vit""": ["""VIT_PRETRAINED_CONFIG_ARCHIVE_MAP""", """ViTConfig""", """ViTOnnxConfig"""]}

try:
  if not is_vision_available():
    raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
  pass
else:
  UpperCAmelCase__ :     Union[str, Any]				    =    ["""ViTFeatureExtractor"""]
  UpperCAmelCase__ :     List[Any]				    =    ["""ViTImageProcessor"""]

try:
  if not is_torch_available():
    raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
  pass
else:
  UpperCAmelCase__ :     List[Any]				    =    [
      """VIT_PRETRAINED_MODEL_ARCHIVE_LIST""",
      """ViTForImageClassification""",
      """ViTForMaskedImageModeling""",
      """ViTModel""",
      """ViTPreTrainedModel""",
  ]

try:
  if not is_tf_available():
    raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
  pass
else:
  UpperCAmelCase__ :     Tuple				    =    [
      """TFViTForImageClassification""",
      """TFViTModel""",
      """TFViTPreTrainedModel""",
  ]

try:
  if not is_flax_available():
    raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
  pass
else:
  UpperCAmelCase__ :     str				    =    [
      """FlaxViTForImageClassification""",
      """FlaxViTModel""",
      """FlaxViTPreTrainedModel""",
  ]

if TYPE_CHECKING:
  from .configuration_vit import VIT_PRETRAINED_CONFIG_ARCHIVE_MAP, ViTConfig, ViTOnnxConfig

  try:
    if not is_vision_available():
      raise OptionalDependencyNotAvailable()
  except OptionalDependencyNotAvailable:
    pass
  else:
    from .feature_extraction_vit import ViTFeatureExtractor
    from .image_processing_vit import ViTImageProcessor

  try:
    if not is_torch_available():
      raise OptionalDependencyNotAvailable()
  except OptionalDependencyNotAvailable:
    pass
  else:
    from .modeling_vit import (
        VIT_PRETRAINED_MODEL_ARCHIVE_LIST,
        ViTForImageClassification,
        ViTForMaskedImageModeling,
        ViTModel,
        ViTPreTrainedModel,
    )

  try:
    if not is_tf_available():
      raise OptionalDependencyNotAvailable()
  except OptionalDependencyNotAvailable:
    pass
  else:
    from .modeling_tf_vit import TFViTForImageClassification, TFViTModel, TFViTPreTrainedModel

  try:
    if not is_flax_available():
      raise OptionalDependencyNotAvailable()
  except OptionalDependencyNotAvailable:
    pass
  else:
    from .modeling_flax_vit import FlaxViTForImageClassification, FlaxViTModel, FlaxViTPreTrainedModel


else:
  import sys

  UpperCAmelCase__ :     Optional[Any]				    =    _LazyModule(__name__, globals()["""__file__"""], _import_structure, module_spec=__spec__)




