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


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

__magic_name__       =			"\n    Examples:\n        ```py\n        >>> from diffusers import KandinskyPipeline, KandinskyPriorPipeline\n        >>> import torch\n\n        >>> pipe_prior = KandinskyPriorPipeline.from_pretrained(\"kandinsky-community/Kandinsky-2-1-prior\")\n        >>> pipe_prior.to(\"cuda\")\n\n        >>> prompt = \"red cat, 4k photo\"\n        >>> out = pipe_prior(prompt)\n        >>> image_emb = out.image_embeds\n        >>> negative_image_emb = out.negative_image_embeds\n\n        >>> pipe = KandinskyPipeline.from_pretrained(\"kandinsky-community/kandinsky-2-1\")\n        >>> pipe.to(\"cuda\")\n\n        >>> image = pipe(\n        ...     prompt,\n        ...     image_embeds=image_emb,\n        ...     negative_image_embeds=negative_image_emb,\n        ...     height=768,\n        ...     width=768,\n        ...     num_inference_steps=100,\n        ... ).images\n\n        >>> image[0].save(\"cat.png\")\n        ```\n"



def 						_lowerCAmelCase				(				UpperCamelCase_  ,				UpperCamelCase_  ,				UpperCamelCase_=8				):
		__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 						SCREAMING_SNAKE_CASE_ (						__a ):




		"""simple docstring"""





		def __init__(     self							,  lowerCAmelCase__							,  lowerCAmelCase__							,  lowerCAmelCase__							,  lowerCAmelCase__							,  lowerCAmelCase__							,  ):
				super().__init__()

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




		def 	snake_case_     (     self							,  lowerCAmelCase__							,  lowerCAmelCase__							,  lowerCAmelCase__							,  lowerCAmelCase__							,  lowerCAmelCase__							,  lowerCAmelCase__):
				if latents is None:
						__SCREAMING_SNAKE_CASE            = 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}")
						__SCREAMING_SNAKE_CASE            = latents.to(lowerCAmelCase__)

				__SCREAMING_SNAKE_CASE            = latents * scheduler.init_noise_sigma
				return latents




		def 	snake_case_     (     self							,  lowerCAmelCase__							,  lowerCAmelCase__							,  lowerCAmelCase__							,  lowerCAmelCase__							,  lowerCAmelCase__=None							,  ):
				__SCREAMING_SNAKE_CASE            = len(lowerCAmelCase__) if isinstance(lowerCAmelCase__							,  lowerCAmelCase__) else 1
				# get prompt text embeddings
				__SCREAMING_SNAKE_CASE            = self.tokenizer(
				    lowerCAmelCase__							,  padding="""max_length"""							,  truncation=lowerCAmelCase__							,  max_length=7_7							,  return_attention_mask=lowerCAmelCase__							,  add_special_tokens=lowerCAmelCase__							,  return_tensors="""pt"""							,  )

				__SCREAMING_SNAKE_CASE            = text_inputs.input_ids
				__SCREAMING_SNAKE_CASE            = self.tokenizer(lowerCAmelCase__							,  padding="""longest"""							,  return_tensors="""pt""").input_ids

				if untruncated_ids.shape[-1] >= text_input_ids.shape[-1] and not torch.equal(lowerCAmelCase__							,  lowerCAmelCase__):
						__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(lowerCAmelCase__)
				__SCREAMING_SNAKE_CASE            = text_inputs.attention_mask.to(lowerCAmelCase__)

				__SCREAMING_SNAKE_CASE       ,__SCREAMING_SNAKE_CASE            = self.text_encoder(
				    input_ids=lowerCAmelCase__							,  attention_mask=lowerCAmelCase__)

				__SCREAMING_SNAKE_CASE            = prompt_embeds.repeat_interleave(lowerCAmelCase__							,  dim=0)
				__SCREAMING_SNAKE_CASE            = text_encoder_hidden_states.repeat_interleave(lowerCAmelCase__							,  dim=0)
				__SCREAMING_SNAKE_CASE            = text_mask.repeat_interleave(lowerCAmelCase__							,  dim=0)

				if do_classifier_free_guidance:
						__SCREAMING_SNAKE_CASE            = 42
						if negative_prompt is None:
								__SCREAMING_SNAKE_CASE            = [""""""] * batch_size
						elif type(lowerCAmelCase__) is not type(lowerCAmelCase__):
								raise TypeError(
								    f"`negative_prompt` should be the same type to `prompt`, but got {type(lowerCAmelCase__)} !="
								    f" {type(lowerCAmelCase__)}.")
						elif isinstance(lowerCAmelCase__							,  lowerCAmelCase__):
								__SCREAMING_SNAKE_CASE            = [negative_prompt]
						elif batch_size != len(lowerCAmelCase__):
								raise ValueError(
								    f"`negative_prompt`: {negative_prompt} has batch size {len(lowerCAmelCase__)}, 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(
						    lowerCAmelCase__							,  padding="""max_length"""							,  max_length=7_7							,  truncation=lowerCAmelCase__							,  return_attention_mask=lowerCAmelCase__							,  add_special_tokens=lowerCAmelCase__							,  return_tensors="""pt"""							,  )
						__SCREAMING_SNAKE_CASE            = uncond_input.input_ids.to(lowerCAmelCase__)
						__SCREAMING_SNAKE_CASE            = uncond_input.attention_mask.to(lowerCAmelCase__)

						__SCREAMING_SNAKE_CASE       ,__SCREAMING_SNAKE_CASE            = self.text_encoder(
						    input_ids=lowerCAmelCase__							,  attention_mask=lowerCAmelCase__)

						# 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							,  lowerCAmelCase__)
						__SCREAMING_SNAKE_CASE            = negative_prompt_embeds.view(batch_size * num_images_per_prompt							,  lowerCAmelCase__)

						__SCREAMING_SNAKE_CASE            = uncond_text_encoder_hidden_states.shape[1]
						__SCREAMING_SNAKE_CASE            = uncond_text_encoder_hidden_states.repeat(1							,  lowerCAmelCase__							,  1)
						__SCREAMING_SNAKE_CASE            = uncond_text_encoder_hidden_states.view(
						    batch_size * num_images_per_prompt							,  lowerCAmelCase__							,  -1)
						__SCREAMING_SNAKE_CASE            = uncond_text_mask.repeat_interleave(lowerCAmelCase__							,  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 	snake_case_     (     self							,  lowerCAmelCase__=0):
				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(lowerCAmelCase__							,  lowerCAmelCase__)




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

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

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

				if self.safety_checker is not None:
						__SCREAMING_SNAKE_CASE       ,__SCREAMING_SNAKE_CASE            = cpu_offload_with_hook(self.safety_checker							,  lowerCAmelCase__							,  prev_module_hook=lowerCAmelCase__)

				# 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 	snake_case_     (     self):
				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							,  lowerCAmelCase__							,  lowerCAmelCase__							,  lowerCAmelCase__							,  lowerCAmelCase__ = None							,  lowerCAmelCase__ = 5_1_2							,  lowerCAmelCase__ = 5_1_2							,  lowerCAmelCase__ = 1_0_0							,  lowerCAmelCase__ = 4.0							,  lowerCAmelCase__ = 1							,  lowerCAmelCase__ = None							,  lowerCAmelCase__ = None							,  lowerCAmelCase__ = "pil"							,  lowerCAmelCase__ = True							,  ):

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

				__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       ,__SCREAMING_SNAKE_CASE       ,__SCREAMING_SNAKE_CASE            = self._encode_prompt(
				    lowerCAmelCase__							,  lowerCAmelCase__							,  lowerCAmelCase__							,  lowerCAmelCase__							,  lowerCAmelCase__)

				if isinstance(lowerCAmelCase__							,  lowerCAmelCase__):
						__SCREAMING_SNAKE_CASE            = torch.cat(lowerCAmelCase__							,  dim=0)
				if isinstance(lowerCAmelCase__							,  lowerCAmelCase__):
						__SCREAMING_SNAKE_CASE            = torch.cat(lowerCAmelCase__							,  dim=0)

				if do_classifier_free_guidance:
						__SCREAMING_SNAKE_CASE            = image_embeds.repeat_interleave(lowerCAmelCase__							,  dim=0)
						__SCREAMING_SNAKE_CASE            = negative_image_embeds.repeat_interleave(lowerCAmelCase__							,  dim=0)

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

				self.scheduler.set_timesteps(lowerCAmelCase__							,  device=lowerCAmelCase__)
				__SCREAMING_SNAKE_CASE            = self.scheduler.timesteps

				__SCREAMING_SNAKE_CASE            = self.unet.config.in_channels

				__SCREAMING_SNAKE_CASE       ,__SCREAMING_SNAKE_CASE            = get_new_h_w(lowerCAmelCase__							,  lowerCAmelCase__							,  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							,  lowerCAmelCase__							,  lowerCAmelCase__							,  lowerCAmelCase__							,  self.scheduler							,  )

				for i, t in enumerate(self.progress_bar(lowerCAmelCase__)):
						# 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=lowerCAmelCase__							,  timestep=lowerCAmelCase__							,  encoder_hidden_states=lowerCAmelCase__							,  added_cond_kwargs=lowerCAmelCase__							,  return_dict=lowerCAmelCase__							,  )[0]

						if do_classifier_free_guidance:
								__SCREAMING_SNAKE_CASE       ,__SCREAMING_SNAKE_CASE            = noise_pred.split(latents.shape[1]							,  dim=1)
								__SCREAMING_SNAKE_CASE       ,__SCREAMING_SNAKE_CASE            = noise_pred.chunk(2)
								__SCREAMING_SNAKE_CASE       ,__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       ,__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(
						    lowerCAmelCase__							,  lowerCAmelCase__							,  lowerCAmelCase__							,  generator=lowerCAmelCase__							,  ).prev_sample
				# post-processing
				__SCREAMING_SNAKE_CASE            = 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"]:
						__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(lowerCAmelCase__)

				if not return_dict:
						return (image,)

				return ImagePipelineOutput(images=lowerCAmelCase__)






"""simple docstring"""





def 						_lowerCAmelCase				(				UpperCamelCase_				):

		if isinstance(UpperCamelCase_  ,				UpperCamelCase_				):
				raise TypeError("""'float' object cannot be interpreted as an integer"""				)
		if isinstance(UpperCamelCase_  ,				UpperCamelCase_				):
				raise TypeError("""'str' object cannot be interpreted as an integer"""				)

		if num == 0:
				return "0b0"

		__SCREAMING_SNAKE_CASE            = False

		if num < 0:
				__SCREAMING_SNAKE_CASE            = True
				__SCREAMING_SNAKE_CASE            = -num

		__SCREAMING_SNAKE_CASE            = []
		while num > 0:
				binary.insert(0  ,				num % 2				)
				num >>= 1

		if negative:
				return "-0b" + "".join(str(UpperCamelCase_				) for e in binary				)

		return "0b" + "".join(str(UpperCamelCase_				) for e in binary				)


if __name__ == "__main__":
							import doctest

							doctest.testmod()




