



"""simple docstring"""
from typing import List, Optional, Union

import torch

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

UpperCAmelCase						:		Union[str, Any]			= "\n    Examples:\n        ```py\n        >>> from diffusers import KandinskyV22Pipeline, KandinskyV22PriorPipeline\n        >>> import torch\n\n        >>> pipe_prior = KandinskyV22PriorPipeline.from_pretrained(\"kandinsky-community/kandinsky-2-2-prior\")\n        >>> pipe_prior.to(\"cuda\")\n        >>> prompt = \"red cat, 4k photo\"\n        >>> out = pipe_prior(prompt)\n        >>> image_emb = out.image_embeds\n        >>> zero_image_emb = out.negative_image_embeds\n        >>> pipe = KandinskyV22Pipeline.from_pretrained(\"kandinsky-community/kandinsky-2-2-decoder\")\n        >>> pipe.to(\"cuda\")\n        >>> image = pipe(\n        ...     image_embeds=image_emb,\n        ...     negative_image_embeds=zero_image_emb,\n        ...     height=768,\n        ...     width=768,\n        ...     num_inference_steps=50,\n        ... ).images\n        >>> image[0].save(\"cat.png\")\n        ```\n"


def 			_SCREAMING_SNAKE_CASE						(__lowerCAmelCase ,  __lowerCAmelCase ,  __lowerCAmelCase=8							) ->						Tuple:







  '''simple docstring'''


  lowercase_								= height // scale_factor**2
  if height % scale_factor**2 != 0:
    new_height += 1
  lowercase_								= width // scale_factor**2
  if width % scale_factor**2 != 0:
    new_width += 1
  return new_height * scale_factor, new_width * scale_factor


class    SCREAMING_SNAKE_CASE__ (							__UpperCAmelCase   ):

    def __init__(				self			:				Optional[Any] ,  lowerCAmelCase_			:				UNetaDConditionModel ,  lowerCAmelCase_			:				DDPMScheduler ,  lowerCAmelCase_			:				VQModel ,  ):







      """simple docstring"""







      super().__init__()

      self.register_modules(
          unet=lowerCAmelCase_ ,  scheduler=lowerCAmelCase_ ,  movq=lowerCAmelCase_ ,  )
      lowercase_								= 2 ** (len(self.movq.config.block_out_channels) - 1)

    def 					_UpperCAmelCase	(				self			:				Dict ,  lowerCAmelCase_			:				Optional[int] ,  lowerCAmelCase_			:				List[str] ,  lowerCAmelCase_			:				List[str] ,  lowerCAmelCase_			:				Tuple ,  lowerCAmelCase_			:				Tuple ,  lowerCAmelCase_			:				Union[str, Any]):







      """simple docstring"""







      if latents is None:
        lowercase_								= 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}''')
        lowercase_								= latents.to(lowerCAmelCase_)

      lowercase_								= latents * scheduler.init_noise_sigma
      return latents

    def 					_UpperCAmelCase	(				self			:				Union[str, Any] ,  lowerCAmelCase_			:				Optional[Any]=0):







      """simple docstring"""







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

      lowercase_								= torch.device(F'''cuda:{gpu_id}''')

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

    def 					_UpperCAmelCase	(				self			:				Optional[int] ,  lowerCAmelCase_			:				Optional[int]=0):







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

      lowercase_								= 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)

      lowercase_								= None
      for cpu_offloaded_model in [self.unet, self.movq]:
        lowercase_					,    lowercase_								= cpu_offload_with_hook(lowerCAmelCase_ ,  lowerCAmelCase_ ,  prev_module_hook=lowerCAmelCase_)

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

    @property
    # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline._execution_device
    def 					_UpperCAmelCase	(				self			:				List[str]):







      """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			:				Dict ,  lowerCAmelCase_			:				Union[torch.FloatTensor, List[torch.FloatTensor]] ,  lowerCAmelCase_			:				Union[torch.FloatTensor, List[torch.FloatTensor]] ,  lowerCAmelCase_			:				int = 5_1_2 ,  lowerCAmelCase_			:				int = 5_1_2 ,  lowerCAmelCase_			:				int = 1_0_0 ,  lowerCAmelCase_			:				float = 4.0 ,  lowerCAmelCase_			:				int = 1 ,  lowerCAmelCase_			:				Optional[Union[torch.Generator, List[torch.Generator]]] = None ,  lowerCAmelCase_			:				Optional[torch.FloatTensor] = None ,  lowerCAmelCase_			:				Optional[str] = "pil" ,  lowerCAmelCase_			:				bool = True ,  ):







      """simple docstring"""







      lowercase_								= self._execution_device

      lowercase_								= guidance_scale > 1.0

      if isinstance(lowerCAmelCase_ ,  lowerCAmelCase_):
        lowercase_								= torch.cat(lowerCAmelCase_ ,  dim=0)
      lowercase_								= image_embeds.shape[0] * num_images_per_prompt
      if isinstance(lowerCAmelCase_ ,  lowerCAmelCase_):
        lowercase_								= torch.cat(lowerCAmelCase_ ,  dim=0)

      if do_classifier_free_guidance:
        lowercase_								= image_embeds.repeat_interleave(lowerCAmelCase_ ,  dim=0)
        lowercase_								= negative_image_embeds.repeat_interleave(lowerCAmelCase_ ,  dim=0)

      lowercase_								= torch.cat([negative_image_embeds, image_embeds] ,  dim=0).to(dtype=self.unet.dtype ,  device=lowerCAmelCase_)

      self.scheduler.set_timesteps(lowerCAmelCase_ ,  device=lowerCAmelCase_)
      lowercase_								= self.scheduler.timesteps

      lowercase_								= self.unet.config.in_channels

      lowercase_					,    lowercase_								= downscale_height_and_width(lowerCAmelCase_ ,  lowerCAmelCase_ ,  self.movq_scale_factor)

      # create initial latent
      lowercase_								= self.prepare_latents(
          (batch_size, num_channels_latents, height, width) ,  image_embeds.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
        lowercase_								= torch.cat([latents] * 2) if do_classifier_free_guidance else latents

        lowercase_								= {"""image_embeds""": image_embeds}
        lowercase_								= self.unet(
            sample=lowerCAmelCase_ ,  timestep=lowerCAmelCase_ ,  encoder_hidden_states=lowerCAmelCase_ ,  added_cond_kwargs=lowerCAmelCase_ ,  return_dict=lowerCAmelCase_ ,  )[0]

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

        # compute the previous noisy sample x_t -> x_t-1
        lowercase_								= self.scheduler.step(
            lowerCAmelCase_ ,  lowerCAmelCase_ ,  lowerCAmelCase_ ,  generator=lowerCAmelCase_ ,  )[0]
      # post-processing
      lowercase_								= 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"]:
        lowercase_								= image * 0.5 + 0.5
        lowercase_								= image.clamp(0 ,  1)
        lowercase_								= image.cpu().permute(0 ,  2 ,  3 ,  1).float().numpy()

      if output_type == "pil":
        lowercase_								= self.numpy_to_pil(lowerCAmelCase_)

      if not return_dict:
        return (image,)

      return ImagePipelineOutput(images=lowerCAmelCase_)






"""simple docstring"""
from typing import Any, Dict, List, Union

from ..utils import add_end_docstrings, is_torch_available, is_vision_available, logging, requires_backends
from .base import PIPELINE_INIT_ARGS, ChunkPipeline


if is_vision_available():
       from PIL import Image

       from ..image_utils import load_image

if is_torch_available():
       import torch

       from transformers.modeling_outputs import BaseModelOutput

       from ..models.auto.modeling_auto import MODEL_FOR_ZERO_SHOT_OBJECT_DETECTION_MAPPING

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


@add_end_docstrings(__UpperCAmelCase   )
class    SCREAMING_SNAKE_CASE__ (							__UpperCAmelCase   ):

    def __init__(				self			:				Optional[int] ,  **lowerCAmelCase_			:				List[Any]):







      """simple docstring"""







      super().__init__(**lowerCAmelCase_)

      if self.framework == "tf":
        raise ValueError(F'''The {self.__class__} is only available in PyTorch.''')

      requires_backends(self ,  """vision""")
      self.check_model_type(lowerCAmelCase_)

    def __call__(				self			:				Dict ,  lowerCAmelCase_			:				Union[str, "Image.Image", List[Dict[str, Any]]] ,  lowerCAmelCase_			:				Union[str, List[str]] = None ,  **lowerCAmelCase_			:				Dict ,  ):







      """simple docstring"""







      if "text_queries" in kwargs:
        lowercase_								= kwargs.pop("""text_queries""")

      if isinstance(lowerCAmelCase_ ,  (str, Image.Image)):
        lowercase_								= {"""image""": image, """candidate_labels""": candidate_labels}
      else:
        lowercase_								= image
      lowercase_								= super().__call__(lowerCAmelCase_ ,  **lowerCAmelCase_)
      return results

    def 					_UpperCAmelCase	(				self			:				Optional[Any] ,  **lowerCAmelCase_			:				Optional[Any]):







      """simple docstring"""







      lowercase_								= {}
      if "threshold" in kwargs:
        lowercase_								= kwargs["""threshold"""]
      if "top_k" in kwargs:
        lowercase_								= kwargs["""top_k"""]
      return {}, {}, postprocess_params

    def 					_UpperCAmelCase	(				self			:				int ,  lowerCAmelCase_			:				List[Any]):







      """simple docstring"""







      lowercase_								= load_image(inputs["""image"""])
      lowercase_								= inputs["""candidate_labels"""]
      if isinstance(lowerCAmelCase_ ,  lowerCAmelCase_):
        lowercase_								= candidate_labels.split(""",""")

      lowercase_								= torch.tensor([[image.height, image.width]] ,  dtype=torch.intaa)
      for i, candidate_label in enumerate(lowerCAmelCase_):
        lowercase_								= self.tokenizer(lowerCAmelCase_ ,  return_tensors=self.framework)
        lowercase_								= self.image_processor(lowerCAmelCase_ ,  return_tensors=self.framework)
        yield {
            "is_last": i == len(lowerCAmelCase_) - 1,
            "target_size": target_size,
            "candidate_label": candidate_label,
            **text_inputs,
            **image_features,
        }

    def 					_UpperCAmelCase	(				self			:				List[Any] ,  lowerCAmelCase_			:				List[Any]):







      """simple docstring"""







      lowercase_								= model_inputs.pop("""target_size""")
      lowercase_								= model_inputs.pop("""candidate_label""")
      lowercase_								= model_inputs.pop("""is_last""")

      lowercase_								= self.model(**lowerCAmelCase_)

      lowercase_								= {"""target_size""": target_size, """candidate_label""": candidate_label, """is_last""": is_last, **outputs}
      return model_outputs

    def 					_UpperCAmelCase	(				self			:				int ,  lowerCAmelCase_			:				Optional[int] ,  lowerCAmelCase_			:				Tuple=0.1 ,  lowerCAmelCase_			:				List[str]=None):







      """simple docstring"""







      lowercase_								= []
      for model_output in model_outputs:
        lowercase_								= model_output["""candidate_label"""]
        lowercase_								= BaseModelOutput(lowerCAmelCase_)
        lowercase_								= self.image_processor.post_process_object_detection(
            outputs=lowerCAmelCase_ ,  threshold=lowerCAmelCase_ ,  target_sizes=model_output["""target_size"""])[0]

        for index in outputs["scores"].nonzero():
          lowercase_								= outputs["""scores"""][index].item()
          lowercase_								= self._get_bounding_box(outputs["""boxes"""][index][0])

          lowercase_								= {"""score""": score, """label""": label, """box""": box}
          results.append(lowerCAmelCase_)

      lowercase_								= sorted(lowerCAmelCase_ ,  key=lambda lowerCAmelCase_: x["score"] ,  reverse=lowerCAmelCase_)
      if top_k:
        lowercase_								= results[:top_k]

      return results




    def 					_UpperCAmelCase	(				self			:				Union[str, Any] ,  lowerCAmelCase_			:				"torch.Tensor"):







      """simple docstring"""







      if self.framework != "pt":
        raise ValueError("""The ZeroShotObjectDetectionPipeline is only available in PyTorch.""")
      lowercase_					,    lowercase_					,    lowercase_					,    lowercase_								= box.int().tolist()
      lowercase_								= {
          """xmin""": xmin,
          """ymin""": ymin,
          """xmax""": xmax,
          """ymax""": ymax,
      }
      return bbox

