

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,
)


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

SCREAMING_SNAKE_CASE_            = '\n    Examples:\n        ```py\n        >>> import torch\n        >>> import numpy as np\n\n        >>> from diffusers import KandinskyV22PriorPipeline, KandinskyV22ControlnetPipeline\n        >>> from transformers import pipeline\n        >>> from diffusers.utils import load_image\n\n\n        >>> def make_hint(image, depth_estimator):\n        ...     image = depth_estimator(image)["depth"]\n        ...     image = np.array(image)\n        ...     image = image[:, :, None]\n        ...     image = np.concatenate([image, image, image], axis=2)\n        ...     detected_map = torch.from_numpy(image).float() / 255.0\n        ...     hint = detected_map.permute(2, 0, 1)\n        ...     return hint\n\n\n        >>> depth_estimator = pipeline("depth-estimation")\n\n        >>> pipe_prior = KandinskyV22PriorPipeline.from_pretrained(\n        ...     "kandinsky-community/kandinsky-2-2-prior", torch_dtype=torch.float16\n        ... )\n        >>> pipe_prior = pipe_prior.to("cuda")\n\n        >>> pipe = KandinskyV22ControlnetPipeline.from_pretrained(\n        ...     "kandinsky-community/kandinsky-2-2-controlnet-depth", torch_dtype=torch.float16\n        ... )\n        >>> pipe = pipe.to("cuda")\n\n\n        >>> img = load_image(\n        ...     "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main"\n        ...     "/kandinsky/cat.png"\n        ... ).resize((768, 768))\n\n        >>> hint = make_hint(img, depth_estimator).unsqueeze(0).half().to("cuda")\n\n        >>> prompt = "A robot, 4k photo"\n        >>> negative_prior_prompt = "lowres, text, error, cropped, worst quality, low quality, jpeg artifacts, ugly, duplicate, morbid, mutilated, out of frame, extra fingers, mutated hands, poorly drawn hands, poorly drawn face, mutation, deformed, blurry, dehydrated, bad anatomy, bad proportions, extra limbs, cloned face, disfigured, gross proportions, malformed limbs, missing arms, missing legs, extra arms, extra legs, fused fingers, too many fingers, long neck, username, watermark, signature"\n\n        >>> generator = torch.Generator(device="cuda").manual_seed(43)\n\n        >>> image_emb, zero_image_emb = pipe_prior(\n        ...     prompt=prompt, negative_prompt=negative_prior_prompt, generator=generator\n        ... ).to_tuple()\n\n        >>> images = pipe(\n        ...     image_embeds=image_emb,\n        ...     negative_image_embeds=zero_image_emb,\n        ...     hint=hint,\n        ...     num_inference_steps=50,\n        ...     generator=generator,\n        ...     height=768,\n        ...     width=768,\n        ... ).images\n\n        >>> images[0].save("robot_cat.png")\n        ```\n'

def 							__SCREAMING_SNAKE_CASE							(			lowerCAmelCase:					Any					, lowerCAmelCase:					List[str]					, lowerCAmelCase:					int=8	)				->    int:
    _UpperCAmelCase    :						int							=      height // scale_factor**2
    if height % scale_factor**2 != 0:
        new_height += 1
    _UpperCAmelCase    :						List[str]							=      width // scale_factor**2
    if width % scale_factor**2 != 0:
        new_width += 1
    return new_height * scale_factor, new_width * scale_factor





class 			a					(		UpperCAmelCase			):


      def __init__(    self  ,   A_  ,   A_  ,   A_  ,   ):







          '''simple docstring'''


          super().__init__()

          self.register_modules(
              unet=A_  ,   scheduler=A_  ,   movq=A_  ,   )
          _UpperCAmelCase    :						Optional[int]							=      2 ** (len(self.movq.config.block_out_channels		) - 1)


      def     _UpperCAmelCase			(    self  ,   A_  ,   A_  ,   A_  ,   A_  ,   A_  ,   A_		):







          '''simple docstring'''


          if latents is None:
              _UpperCAmelCase    :						Any							=      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}'		)
              _UpperCAmelCase    :						Optional[int]							=      latents.to(A_		)

          _UpperCAmelCase    :						Tuple							=      latents * scheduler.init_noise_sigma
          return latents


      def     _UpperCAmelCase			(    self  ,   A_=0		):







          '''simple docstring'''


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

          _UpperCAmelCase    :						Union[str, Any]							=      torch.device(f'cuda:{gpu_id}'		)

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


      def     _UpperCAmelCase			(    self  ,   A_=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."		)

          _UpperCAmelCase    :						str							=      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)

          _UpperCAmelCase    :						Optional[int]							=      None
          for cpu_offloaded_model in [self.unet, self.movq]:
              _UpperCAmelCase    :						Dict							=      cpu_offload_with_hook(A_  ,   A_  ,   prev_module_hook=A_		)

          # We'll offload the last model manually.
          _UpperCAmelCase    :						Optional[Any]							=      hook


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







          '''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_ = 512  ,   A_ = 512  ,   A_ = 100  ,   A_ = 4.0  ,   A_ = 1  ,   A_ = None  ,   A_ = None  ,   A_ = "pil"  ,   A_ = True  ,   ):







          '''simple docstring'''


          _UpperCAmelCase    :						str							=      self._execution_device

          _UpperCAmelCase    :						Tuple							=      guidance_scale > 1.0

          if isinstance(A_  ,   A_		):
              _UpperCAmelCase    :						Union[str, Any]							=      torch.cat(A_  ,   dim=0		)
          if isinstance(A_  ,   A_		):
              _UpperCAmelCase    :						Dict							=      torch.cat(A_  ,   dim=0		)
          if isinstance(A_  ,   A_		):
              _UpperCAmelCase    :						Any							=      torch.cat(A_  ,   dim=0		)

          _UpperCAmelCase    :						Optional[Any]							=      image_embeds.shape[0] * num_images_per_prompt

          if do_classifier_free_guidance:
              _UpperCAmelCase    :						Optional[int]							=      image_embeds.repeat_interleave(A_  ,   dim=0		)
              _UpperCAmelCase    :						Union[str, Any]							=      negative_image_embeds.repeat_interleave(A_  ,   dim=0		)
              _UpperCAmelCase    :						Tuple							=      hint.repeat_interleave(A_  ,   dim=0		)

          _UpperCAmelCase    :						List[Any]							=      torch.cat([negative_image_embeds, image_embeds]  ,   dim=0		).to(dtype=self.unet.dtype  ,   device=A_		)
          _UpperCAmelCase    :						List[Any]							=      torch.cat([hint, hint]  ,   dim=0		).to(dtype=self.unet.dtype  ,   device=A_		)
          self.scheduler.set_timesteps(A_  ,   device=A_		)
          _UpperCAmelCase    :						Dict							=      self.scheduler.timesteps

          _UpperCAmelCase    :						Union[str, Any]							=      self.movq.config.latent_channels

          _UpperCAmelCase    :						Optional[int]							=      downscale_height_and_width(A_  ,   A_  ,   self.movq_scale_factor		)

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

              _UpperCAmelCase    :						Optional[Any]							=      {"image_embeds": image_embeds, "hint": hint}
              _UpperCAmelCase    :						Optional[int]							=      self.unet(
                  sample=A_  ,   timestep=A_  ,   encoder_hidden_states=A_  ,   added_cond_kwargs=A_  ,   return_dict=A_  ,   )[0]

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

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

          if output_type == "pil":
              _UpperCAmelCase    :						str							=      self.numpy_to_pil(A_		)

          if not return_dict:
              return (image,)

          return ImagePipelineOutput(images=A_		)





import warnings
from typing import List

import numpy as np

from ...processing_utils import ProcessorMixin
from ...tokenization_utils_base import BatchEncoding
from ...utils import is_flax_available, is_tf_available, is_torch_available

class 			a					(		UpperCAmelCase			):
						_lowercase        =     ["image_processor", "tokenizer"]
						_lowercase        =     "OwlViTImageProcessor"
						_lowercase        =     ("CLIPTokenizer", "CLIPTokenizerFast")


						def __init__(    self  ,   A_=None  ,   A_=None  ,   **A_		):







										'''simple docstring'''


										_UpperCAmelCase    :						List[str]							=      None
										if "feature_extractor" in kwargs:
														warnings.warn(
														    "The `feature_extractor` argument is deprecated and will be removed in v5, use `image_processor`"
														    " instead."  ,   A_  ,   )
														_UpperCAmelCase    :						Union[str, Any]							=      kwargs.pop("feature_extractor"		)

										_UpperCAmelCase    :						Any							=      image_processor if image_processor is not None else feature_extractor
										if image_processor is None:
														raise ValueError("You need to specify an `image_processor`."		)
										if tokenizer is None:
														raise ValueError("You need to specify a `tokenizer`."		)

										super().__init__(A_  ,   A_		)


						def __call__(    self  ,   A_=None  ,   A_=None  ,   A_=None  ,   A_="max_length"  ,   A_="np"  ,   **A_		):







										'''simple docstring'''


										if text is None and query_images is None and images is None:
														raise ValueError(
														    "You have to specify at least one text or query image or image. All three cannot be none."		)

										if text is not None:
														if isinstance(A_  ,   A_		) or (isinstance(A_  ,   A_		) and not isinstance(text[0]  ,   A_		)):
																		_UpperCAmelCase    :						Optional[int]							=      [self.tokenizer(A_  ,   padding=A_  ,   return_tensors=A_  ,   **A_		)]

														elif isinstance(A_  ,   A_		) and isinstance(text[0]  ,   A_		):
																		_UpperCAmelCase    :						Optional[int]							=      []

																		# Maximum number of queries across batch
																		_UpperCAmelCase    :						Optional[Any]							=      max([len(A_		) for t in text]		)

																		# Pad all batch samples to max number of text queries
																		for t in text:
																						if len(A_		) != max_num_queries:
																										_UpperCAmelCase    :						Optional[int]							=      t + [" "] * (max_num_queries - len(A_		))

																						_UpperCAmelCase    :						str							=      self.tokenizer(A_  ,   padding=A_  ,   return_tensors=A_  ,   **A_		)
																						encodings.append(A_		)
														else:
																		raise TypeError("Input text should be a string, a list of strings or a nested list of strings"		)

														if return_tensors == "np":
																		_UpperCAmelCase    :						List[str]							=      np.concatenate([encoding["input_ids"] for encoding in encodings]  ,   axis=0		)
																		_UpperCAmelCase    :						Tuple							=      np.concatenate([encoding["attention_mask"] for encoding in encodings]  ,   axis=0		)

														elif return_tensors == "jax" and is_flax_available():
																		import jax.numpy as jnp

																		_UpperCAmelCase    :						Optional[Any]							=      jnp.concatenate([encoding["input_ids"] for encoding in encodings]  ,   axis=0		)
																		_UpperCAmelCase    :						str							=      jnp.concatenate([encoding["attention_mask"] for encoding in encodings]  ,   axis=0		)

														elif return_tensors == "pt" and is_torch_available():
																		import torch

																		_UpperCAmelCase    :						str							=      torch.cat([encoding["input_ids"] for encoding in encodings]  ,   dim=0		)
																		_UpperCAmelCase    :						Dict							=      torch.cat([encoding["attention_mask"] for encoding in encodings]  ,   dim=0		)

														elif return_tensors == "tf" and is_tf_available():
																		import tensorflow as tf

																		_UpperCAmelCase    :						Union[str, Any]							=      tf.stack([encoding["input_ids"] for encoding in encodings]  ,   axis=0		)
																		_UpperCAmelCase    :						Optional[int]							=      tf.stack([encoding["attention_mask"] for encoding in encodings]  ,   axis=0		)

														else:
																		raise ValueError("Target return tensor type could not be returned"		)

														_UpperCAmelCase    :						Optional[int]							=      BatchEncoding()
														_UpperCAmelCase    :						str							=      input_ids
														_UpperCAmelCase    :						Optional[Any]							=      attention_mask

										if query_images is not None:
														_UpperCAmelCase    :						int							=      BatchEncoding()
														_UpperCAmelCase    :						str							=      self.image_processor(
														    A_  ,   return_tensors=A_  ,   **A_		).pixel_values
														_UpperCAmelCase    :						Optional[Any]							=      query_pixel_values

										if images is not None:
														_UpperCAmelCase    :						int							=      self.image_processor(A_  ,   return_tensors=A_  ,   **A_		)

										if text is not None and images is not None:
														_UpperCAmelCase    :						Optional[int]							=      image_features.pixel_values
														return encoding
										elif query_images is not None and images is not None:
														_UpperCAmelCase    :						Any							=      image_features.pixel_values
														return encoding
										elif text is not None or query_images is not None:
														return encoding
										else:
														return BatchEncoding(data=dict(**A_		)  ,   tensor_type=A_		)


						def     _UpperCAmelCase			(    self  ,   *A_  ,   **A_		):







										'''simple docstring'''


										return self.image_processor.post_process(*A_  ,   **A_		)


						def     _UpperCAmelCase			(    self  ,   *A_  ,   **A_		):







										'''simple docstring'''


										return self.image_processor.post_process_object_detection(*A_  ,   **A_		)


						def     _UpperCAmelCase			(    self  ,   *A_  ,   **A_		):







										'''simple docstring'''


										return self.image_processor.post_process_image_guided_detection(*A_  ,   **A_		)


						def     _UpperCAmelCase			(    self  ,   *A_  ,   **A_		):







										'''simple docstring'''


										return self.tokenizer.batch_decode(*A_  ,   **A_		)


						def     _UpperCAmelCase			(    self  ,   *A_  ,   **A_		):







										'''simple docstring'''


										return self.tokenizer.decode(*A_  ,   **A_		)


						@property
						def     _UpperCAmelCase			(    self		):







										'''simple docstring'''


										warnings.warn(
										    "`feature_extractor_class` is deprecated and will be removed in v5. Use `image_processor_class` instead."  ,   A_  ,   )
										return self.image_processor_class





						@property
						def     _UpperCAmelCase			(    self		):







										'''simple docstring'''


										warnings.warn(
										    "`feature_extractor` is deprecated and will be removed in v5. Use `image_processor` instead."  ,   A_  ,   )
										return self.image_processor


