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

_SCREAMING_SNAKE_CASE	       =      """
    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   (					SCREAMING_SNAKE_CASE							,    SCREAMING_SNAKE_CASE							,    SCREAMING_SNAKE_CASE=8 )  ->     int:







							"""simple docstring"""


							__snake_case     			=       h // scale_factor**2
							if h % scale_factor**2 != 0:
														new_h += 1
							__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 							__magic_name__     (		lowercase__   ):





					def __init__(				self   :							Tuple							,						snake_case_   :							MultilingualCLIP							,						snake_case_   :							XLMRobertaTokenizer							,						snake_case_   :							UNetaDConditionModel							,						snake_case_   :							Union[DDIMScheduler, DDPMScheduler]							,						snake_case_   :							VQModel							,						):
												super().__init__()

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





					def 							lowerCAmelCase      (				self   :							int							,						snake_case_   :							Any							,						snake_case_   :							List[Any]							,						snake_case_   :							int							,						snake_case_   :							Any							,						snake_case_   :							int							,						snake_case_   :							Any    ):
												if latents is None:
																			__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}'''    )
																			__snake_case     			=       latents.to(snake_case_    )

												__snake_case     			=       latents * scheduler.init_noise_sigma
												return latents





					def 							lowerCAmelCase      (				self   :							Optional[int]							,						snake_case_   :							List[Any]							,						snake_case_   :							Optional[Any]							,						snake_case_   :							Optional[Any]							,						snake_case_   :							Dict							,						snake_case_   :							int=None							,						):
												__snake_case     			=       len(snake_case_    ) if isinstance(snake_case_							,						snake_case_    ) else 1
												# get prompt text embeddings
												__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"							,						)

												__snake_case     			=       text_inputs.input_ids
												__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_    ):
																			__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}'''    )

												__snake_case     			=       text_input_ids.to(snake_case_    )
												__snake_case     			=       text_inputs.attention_mask.to(snake_case_    )

												__snake_case			,   __snake_case     			=       self.text_encoder(
												    input_ids=snake_case_							,						attention_mask=snake_case_    )

												__snake_case     			=       prompt_embeds.repeat_interleave(snake_case_							,						dim=0    )
												__snake_case     			=       text_encoder_hidden_states.repeat_interleave(snake_case_							,						dim=0    )
												__snake_case     			=       text_mask.repeat_interleave(snake_case_							,						dim=0    )

												if do_classifier_free_guidance:
																			__snake_case     			=       42
																			if negative_prompt is None:
																										__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_    ):
																										__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:
																										__snake_case     			=       negative_prompt

																			__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"							,						)
																			__snake_case     			=       uncond_input.input_ids.to(snake_case_    )
																			__snake_case     			=       uncond_input.attention_mask.to(snake_case_    )

																			__snake_case			,   __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

																			__snake_case     			=       negative_prompt_embeds.shape[1]
																			__snake_case     			=       negative_prompt_embeds.repeat(1							,						snake_case_    )
																			__snake_case     			=       negative_prompt_embeds.view(batch_size * num_images_per_prompt							,						snake_case_    )

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

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

												return prompt_embeds, text_encoder_hidden_states, text_mask





					def 							lowerCAmelCase      (				self   :							Dict							,						snake_case_   :							Union[str, Any]=0    ):
												if is_accelerate_available():
																			from accelerate import cpu_offload
												else:
																			raise ImportError("Please install accelerate via `pip install accelerate`"    )

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

												__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 							lowerCAmelCase      (				self   :							List[str]							,						snake_case_   :							Optional[Any]=0    ):
												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."    )

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

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

												if self.safety_checker is not None:
																			__snake_case			,   __snake_case     			=       cpu_offload_with_hook(self.safety_checker							,						snake_case_							,						prev_module_hook=snake_case_    )

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





					@property
					# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline._execution_device
					def 							lowerCAmelCase      (				self   :							str    ):
												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   :							int							,						snake_case_   :							Union[str, List[str]]							,						snake_case_   :							Union[torch.FloatTensor, List[torch.FloatTensor]]							,						snake_case_   :							Union[torch.FloatTensor, List[torch.FloatTensor]]							,						snake_case_   :							Optional[Union[str, List[str]]] = None							,						snake_case_   :							int = 512							,						snake_case_   :							int = 512							,						snake_case_   :							int = 100							,						snake_case_   :							float = 4.0							,						snake_case_   :							int = 1							,						snake_case_   :							Optional[Union[torch.Generator, List[torch.Generator]]] = None							,						snake_case_   :							Optional[torch.FloatTensor] = None							,						snake_case_   :							Optional[str] = "pil"							,						snake_case_   :							bool = True							,						):

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

												__snake_case     			=       self._execution_device

												__snake_case     			=       batch_size * num_images_per_prompt
												__snake_case     			=       guidance_scale > 1.0

												__snake_case			,   __snake_case			,   __snake_case     			=       self._encode_prompt(
												    snake_case_							,						snake_case_							,						snake_case_							,						snake_case_							,						snake_case_    )

												if isinstance(snake_case_							,						snake_case_    ):
																			__snake_case     			=       torch.cat(snake_case_							,						dim=0    )
												if isinstance(snake_case_							,						snake_case_    ):
																			__snake_case     			=       torch.cat(snake_case_							,						dim=0    )

												if do_classifier_free_guidance:
																			__snake_case     			=       image_embeds.repeat_interleave(snake_case_							,						dim=0    )
																			__snake_case     			=       negative_image_embeds.repeat_interleave(snake_case_							,						dim=0    )

												__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_    )
												__snake_case     			=       self.scheduler.timesteps

												__snake_case     			=       self.unet.config.in_channels

												__snake_case			,   __snake_case     			=       get_new_h_w(snake_case_							,						snake_case_							,						self.movq_scale_factor    )

												# create initial latent
												__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
																			__snake_case     			=       torch.cat([latents] * 2    ) if do_classifier_free_guidance else latents

																			__snake_case     			=       {"text_embeds": prompt_embeds, "image_embeds": image_embeds}
																			__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:
																										__snake_case			,   __snake_case     			=       noise_pred.split(latents.shape[1]							,						dim=1    )
																										__snake_case			,   __snake_case     			=       noise_pred.chunk(2    )
																										__snake_case			,   __snake_case     			=       variance_pred.chunk(2    )
																										__snake_case     			=       noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond)
																										__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"]
																			):
																										__snake_case			,   __snake_case     			=       noise_pred.split(latents.shape[1]							,						dim=1    )

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

												if output_type == "pil":
																			__snake_case     			=       self.numpy_to_pil(snake_case_    )

												if not return_dict:
																			return (image,)

												return ImagePipelineOutput(images=snake_case_    )


"""simple docstring"""







from collections import UserDict
from typing import List, Union

from ..utils import (
    add_end_docstrings,
    is_tf_available,
    is_torch_available,
    is_vision_available,
    logging,
    requires_backends,
)
from .base import PIPELINE_INIT_ARGS, Pipeline


if is_vision_available():
		from PIL import Image

		from ..image_utils import load_image

if is_torch_available():
		from ..models.auto.modeling_auto import MODEL_FOR_ZERO_SHOT_IMAGE_CLASSIFICATION_MAPPING

if is_tf_available():
		from ..models.auto.modeling_tf_auto import TF_MODEL_FOR_ZERO_SHOT_IMAGE_CLASSIFICATION_MAPPING
		from ..tf_utils import stable_softmax

_SCREAMING_SNAKE_CASE	       =      logging.get_logger(__name__)







@add_end_docstrings(lowercase__   )
class 							__magic_name__     (		lowercase__   ):





					def __init__(				self   :							Tuple							,						**snake_case_   :							str    ):
												super().__init__(**snake_case_    )

												requires_backends(self							,						"vision"    )
												self.check_model_type(
												    TF_MODEL_FOR_ZERO_SHOT_IMAGE_CLASSIFICATION_MAPPING
												    if self.framework == "tf"
												    else MODEL_FOR_ZERO_SHOT_IMAGE_CLASSIFICATION_MAPPING    )





					def __call__(				self   :							Optional[Any]							,						snake_case_   :							Union[str, List[str], "Image", List["Image"]]							,						**snake_case_   :							Optional[int]    ):
												return super().__call__(snake_case_							,						**snake_case_    )





					def 							lowerCAmelCase      (				self   :							Tuple							,						**snake_case_   :							Any    ):
												__snake_case     			=       {}
												if "candidate_labels" in kwargs:
																			__snake_case     			=       kwargs["candidate_labels"]
												if "hypothesis_template" in kwargs:
																			__snake_case     			=       kwargs["hypothesis_template"]

												return preprocess_params, {}, {}





					def 							lowerCAmelCase      (				self   :							int							,						snake_case_   :							List[str]							,						snake_case_   :							List[str]=None							,						snake_case_   :							int="This is a photo of {}."    ):
												__snake_case     			=       load_image(snake_case_    )
												__snake_case     			=       self.image_processor(images=[image]							,						return_tensors=self.framework    )
												__snake_case     			=       candidate_labels
												__snake_case     			=       [hypothesis_template.format(snake_case_    ) for x in candidate_labels]
												__snake_case     			=       self.tokenizer(snake_case_							,						return_tensors=self.framework							,						padding=snake_case_    )
												__snake_case     			=       [text_inputs]
												return inputs





					def 							lowerCAmelCase      (				self   :							Any							,						snake_case_   :							Dict    ):
												__snake_case     			=       model_inputs.pop("candidate_labels"    )
												__snake_case     			=       model_inputs.pop("text_inputs"    )
												if isinstance(text_inputs[0]							,						snake_case_    ):
																			__snake_case     			=       text_inputs[0]
												else:
																			# Batching case.
																			__snake_case     			=       text_inputs[0][0]

												__snake_case     			=       self.model(**snake_case_							,						**snake_case_    )

												__snake_case     			=       {
												    "candidate_labels": candidate_labels,
												    "logits": outputs.logits_per_image,
												}
												return model_outputs





					def 							lowerCAmelCase      (				self   :							Optional[int]							,						snake_case_   :							Any    ):
												__snake_case     			=       model_outputs.pop("candidate_labels"    )
												__snake_case     			=       model_outputs["logits"][0]
												if self.framework == "pt":
																			__snake_case     			=       logits.softmax(dim=-1    ).squeeze(-1    )
																			__snake_case     			=       probs.tolist()
																			if not isinstance(snake_case_							,						snake_case_    ):
																										__snake_case     			=       [scores]
												elif self.framework == "tf":
																			__snake_case     			=       stable_softmax(snake_case_							,						axis=-1    )
																			__snake_case     			=       probs.numpy().tolist()
												else:
																			raise ValueError(F'''Unsupported framework: {self.framework}'''    )

												__snake_case     			=       [
												    {"score": score, "label": candidate_label}
												    for score, candidate_label in sorted(zip(snake_case_							,						snake_case_    )							,						key=lambda snake_case_    : -x[0]    )
												]
												return result
