

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


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

SCREAMING_SNAKE_CASE							:				Optional[int]     	=			"""
    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 							_UpperCamelCase				(   lowerCAmelCase__:					Optional[int]		,lowerCAmelCase__:					Optional[Any]		,lowerCAmelCase__:					Dict=8   )      -> str:
							SCREAMING_SNAKE_CASE_				=				h // scale_factor**2
							if h % scale_factor**2 != 0:
														new_h += 1
							SCREAMING_SNAKE_CASE_				=				w // scale_factor**2
							if w % scale_factor**2 != 0:
														new_w += 1
							return new_h * scale_factor, new_w * scale_factor






class    snake_case     (      SCREAMING_SNAKE_CASE_		):






				"""simple docstring"""



				def __init__(    self,     _lowercase,     _lowercase,     _lowercase,     _lowercase,     _lowercase,     )     ->							List[Any]:
											super().__init__()

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



				def   a__					(    self,     _lowercase,     _lowercase,     _lowercase,     _lowercase,     _lowercase,     _lowercase       )     ->							List[str]:

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

											SCREAMING_SNAKE_CASE_				=				latents * scheduler.init_noise_sigma
											return latents



				def   a__					(    self,     _lowercase,     _lowercase,     _lowercase,     _lowercase,     _lowercase=None,     )     ->							Optional[Any]:
											SCREAMING_SNAKE_CASE_				=				len(snake_case__       ) if isinstance(snake_case__,     snake_case__       ) else 1
											# get prompt text embeddings
											SCREAMING_SNAKE_CASE_				=				self.tokenizer(
											    snake_case__,     padding='max_length',     truncation=snake_case__,     max_length=77,     return_attention_mask=snake_case__,     add_special_tokens=snake_case__,     return_tensors='pt',     )

											SCREAMING_SNAKE_CASE_				=				text_inputs.input_ids
											SCREAMING_SNAKE_CASE_				=				self.tokenizer(snake_case__,     padding='longest',     return_tensors='pt'       ).input_ids

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

											SCREAMING_SNAKE_CASE_				=				text_input_ids.to(snake_case__       )
											SCREAMING_SNAKE_CASE_				=				text_inputs.attention_mask.to(snake_case__       )

											SCREAMING_SNAKE_CASE_				=				self.text_encoder(
											    input_ids=snake_case__,     attention_mask=snake_case__       )

											SCREAMING_SNAKE_CASE_				=				prompt_embeds.repeat_interleave(snake_case__,     dim=0       )
											SCREAMING_SNAKE_CASE_				=				text_encoder_hidden_states.repeat_interleave(snake_case__,     dim=0       )
											SCREAMING_SNAKE_CASE_				=				text_mask.repeat_interleave(snake_case__,     dim=0       )

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

																		SCREAMING_SNAKE_CASE_				=				self.tokenizer(
																		    snake_case__,     padding='max_length',     max_length=77,     truncation=snake_case__,     return_attention_mask=snake_case__,     add_special_tokens=snake_case__,     return_tensors='pt',     )
																		SCREAMING_SNAKE_CASE_				=				uncond_input.input_ids.to(snake_case__       )
																		SCREAMING_SNAKE_CASE_				=				uncond_input.attention_mask.to(snake_case__       )

																		SCREAMING_SNAKE_CASE_				=				self.text_encoder(
																		    input_ids=snake_case__,     attention_mask=snake_case__       )

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

																		SCREAMING_SNAKE_CASE_				=				negative_prompt_embeds.shape[1]
																		SCREAMING_SNAKE_CASE_				=				negative_prompt_embeds.repeat(1,     snake_case__       )
																		SCREAMING_SNAKE_CASE_				=				negative_prompt_embeds.view(batch_size * num_images_per_prompt,     snake_case__       )

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

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

											return prompt_embeds, text_encoder_hidden_states, text_mask



				def   a__					(    self,     _lowercase=0       )     ->							Any:

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

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

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



				def   a__					(    self,     _lowercase=0       )     ->							Any:

											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_				=				torch.device(f"""cuda:{gpu_id}"""       )

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

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

											if self.safety_checker is not None:
																		SCREAMING_SNAKE_CASE_				=				cpu_offload_with_hook(self.safety_checker,     snake_case__,     prev_module_hook=snake_case__       )

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



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

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

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

											SCREAMING_SNAKE_CASE_				=				self._execution_device

											SCREAMING_SNAKE_CASE_				=				batch_size * num_images_per_prompt
											SCREAMING_SNAKE_CASE_				=				guidance_scale > 1.0

											SCREAMING_SNAKE_CASE_				=				self._encode_prompt(
											    snake_case__,     snake_case__,     snake_case__,     snake_case__,     snake_case__       )

											if isinstance(snake_case__,     snake_case__       ):
																		SCREAMING_SNAKE_CASE_				=				torch.cat(snake_case__,     dim=0       )
											if isinstance(snake_case__,     snake_case__       ):
																		SCREAMING_SNAKE_CASE_				=				torch.cat(snake_case__,     dim=0       )

											if do_classifier_free_guidance:
																		SCREAMING_SNAKE_CASE_				=				image_embeds.repeat_interleave(snake_case__,     dim=0       )
																		SCREAMING_SNAKE_CASE_				=				negative_image_embeds.repeat_interleave(snake_case__,     dim=0       )

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

											self.scheduler.set_timesteps(snake_case__,     device=snake_case__       )
											SCREAMING_SNAKE_CASE_				=				self.scheduler.timesteps

											SCREAMING_SNAKE_CASE_				=				self.unet.config.in_channels

											SCREAMING_SNAKE_CASE_				=				get_new_h_w(snake_case__,     snake_case__,     self.movq_scale_factor       )

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

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

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

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

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

											if output_type == "pil":
																		SCREAMING_SNAKE_CASE_				=				self.numpy_to_pil(snake_case__       )

											if not return_dict:
																		return (image,)

											return ImagePipelineOutput(images=snake_case__       )









'''simple docstring'''







from typing import TYPE_CHECKING

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


lowerCAmelCase     :  int											=    {
    """configuration_owlvit""": [
        """OWLVIT_PRETRAINED_CONFIG_ARCHIVE_MAP""",
        """OwlViTConfig""",
        """OwlViTOnnxConfig""",
        """OwlViTTextConfig""",
        """OwlViTVisionConfig""",
    ],
    """processing_owlvit""": ["""OwlViTProcessor"""],
}


try:
       if not is_vision_available():
              raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
       pass
else:
       lowerCAmelCase     :  Dict											=    ["""OwlViTFeatureExtractor"""]
       lowerCAmelCase     :  int											=    ["""OwlViTImageProcessor"""]

try:
       if not is_torch_available():
              raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
       pass
else:
       lowerCAmelCase     :  Optional[Any]											=    [
           """OWLVIT_PRETRAINED_MODEL_ARCHIVE_LIST""",
           """OwlViTModel""",
           """OwlViTPreTrainedModel""",
           """OwlViTTextModel""",
           """OwlViTVisionModel""",
           """OwlViTForObjectDetection""",
       ]

if TYPE_CHECKING:
       from .configuration_owlvit import (
           OWLVIT_PRETRAINED_CONFIG_ARCHIVE_MAP,
           OwlViTConfig,
           OwlViTOnnxConfig,
           OwlViTTextConfig,
           OwlViTVisionConfig,
       )
       from .processing_owlvit import OwlViTProcessor

       try:
              if not is_vision_available():
                     raise OptionalDependencyNotAvailable()
       except OptionalDependencyNotAvailable:
              pass
       else:
              from .feature_extraction_owlvit import OwlViTFeatureExtractor
              from .image_processing_owlvit import OwlViTImageProcessor

       try:
              if not is_torch_available():
                     raise OptionalDependencyNotAvailable()
       except OptionalDependencyNotAvailable:
              pass
       else:
              from .modeling_owlvit import (
                  OWLVIT_PRETRAINED_MODEL_ARCHIVE_LIST,
                  OwlViTForObjectDetection,
                  OwlViTModel,
                  OwlViTPreTrainedModel,
                  OwlViTTextModel,
                  OwlViTVisionModel,
              )

else:
       import sys

       lowerCAmelCase     :  Union[str, Any]											=    _LazyModule(__name__, globals()["""__file__"""], _import_structure, module_spec=__spec__)

