






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


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

lowerCAmelCase_           = '\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      snake_case		(							UpperCAmelCase	:       Optional[int],						UpperCAmelCase	:       Optional[int],						UpperCAmelCase	:       List[Any]=8   ):
				A	 =      height // scale_factor**2
				if height % scale_factor**2 != 0:
								new_height += 1
				A	 =      width // scale_factor**2
				if width % scale_factor**2 != 0:
								new_width += 1
				return new_height * scale_factor, new_width * scale_factor



class      UpperCamelCase (				snake_case__ ):







		"""simple docstring"""



		def __init__(  self		:       Optional[Any]						,_SCREAMING_SNAKE_CASE		:       UNetaDConditionModel						,_SCREAMING_SNAKE_CASE		:       DDPMScheduler						,_SCREAMING_SNAKE_CASE		:       VQModel						,)    ->   int:


						'''simple docstring'''
						super().__init__()

						self.register_modules(
						    unet=_SCREAMING_SNAKE_CASE						,scheduler=_SCREAMING_SNAKE_CASE						,movq=_SCREAMING_SNAKE_CASE						,)
						A	 =      2 ** (len(self.movq.config.block_out_channels		) - 1)
		def       A(  self		:       int						,_SCREAMING_SNAKE_CASE		:       int						,_SCREAMING_SNAKE_CASE		:       Optional[int]						,_SCREAMING_SNAKE_CASE		:       int						,_SCREAMING_SNAKE_CASE		:       Optional[int]						,_SCREAMING_SNAKE_CASE		:       Any						,_SCREAMING_SNAKE_CASE		:       int		)    ->   Optional[Any]:


						'''simple docstring'''
						if latents is None:
										A	 =      randn_tensor(_SCREAMING_SNAKE_CASE						,generator=_SCREAMING_SNAKE_CASE						,device=_SCREAMING_SNAKE_CASE						,dtype=_SCREAMING_SNAKE_CASE		)
						else:
										if latents.shape != shape:
														raise ValueError(f'Unexpected latents shape, got {latents.shape}, expected {shape}'		)
										A	 =      latents.to(_SCREAMING_SNAKE_CASE		)

						A	 =      latents * scheduler.init_noise_sigma
						return latents
		def       A(  self		:       Any						,_SCREAMING_SNAKE_CASE		:       int=0		)    ->   Optional[int]:


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

						A	 =      torch.device(f'cuda:{gpu_id}'		)

						A	 =      [
						    self.unet,
						    self.movq,
						]
						for cpu_offloaded_model in models:
										if cpu_offloaded_model is not None:
														cpu_offload(_SCREAMING_SNAKE_CASE						,_SCREAMING_SNAKE_CASE		)
		def       A(  self		:       Tuple						,_SCREAMING_SNAKE_CASE		:       List[str]=0		)    ->   Optional[Any]:


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

						A	 =      torch.device(f'cuda:{gpu_id}'		)

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

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

						# We'll offload the last model manually.
						A	 =      hook
		@property
		# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline._execution_device
		def       A(  self		:       Dict		)    ->   Optional[int]:


						'''simple docstring'''
						if not hasattr(self.unet						,'_hf_hook'		):
										return self.device
						for module in self.unet.modules():
										if (
										    hasattr(_SCREAMING_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(_SCREAMING_SNAKE_CASE		)
		def __call__(  self		:       List[str]						,_SCREAMING_SNAKE_CASE		:       Union[torch.FloatTensor, List[torch.FloatTensor]]						,_SCREAMING_SNAKE_CASE		:       Union[torch.FloatTensor, List[torch.FloatTensor]]						,_SCREAMING_SNAKE_CASE		:       int = 5_1_2						,_SCREAMING_SNAKE_CASE		:       int = 5_1_2						,_SCREAMING_SNAKE_CASE		:       int = 1_0_0						,_SCREAMING_SNAKE_CASE		:       float = 4.0						,_SCREAMING_SNAKE_CASE		:       int = 1						,_SCREAMING_SNAKE_CASE		:       Optional[Union[torch.Generator, List[torch.Generator]]] = None						,_SCREAMING_SNAKE_CASE		:       Optional[torch.FloatTensor] = None						,_SCREAMING_SNAKE_CASE		:       Optional[str] = "pil"						,_SCREAMING_SNAKE_CASE		:       bool = True						,)    ->   Optional[Any]:


						'''simple docstring'''
						A	 =      self._execution_device

						A	 =      guidance_scale > 1.0

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

						if do_classifier_free_guidance:
										A	 =      image_embeds.repeat_interleave(_SCREAMING_SNAKE_CASE						,dim=0		)
										A	 =      negative_image_embeds.repeat_interleave(_SCREAMING_SNAKE_CASE						,dim=0		)

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

						self.scheduler.set_timesteps(_SCREAMING_SNAKE_CASE						,device=_SCREAMING_SNAKE_CASE		)
						A	 =      self.scheduler.timesteps

						A	 =      self.unet.config.in_channels

						A ,		A	 =      downscale_height_and_width(_SCREAMING_SNAKE_CASE						,_SCREAMING_SNAKE_CASE						,self.movq_scale_factor		)

						# create initial latent
						A	 =      self.prepare_latents(
						    (batch_size, num_channels_latents, height, width)						,image_embeds.dtype						,_SCREAMING_SNAKE_CASE						,_SCREAMING_SNAKE_CASE						,_SCREAMING_SNAKE_CASE						,self.scheduler						,)

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

										A	 =      {'image_embeds': image_embeds}
										A	 =      self.unet(
										    sample=_SCREAMING_SNAKE_CASE						,timestep=_SCREAMING_SNAKE_CASE						,encoder_hidden_states=_SCREAMING_SNAKE_CASE						,added_cond_kwargs=_SCREAMING_SNAKE_CASE						,return_dict=_SCREAMING_SNAKE_CASE						,)[0]

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

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

						if output_type == "pil":
										A	 =      self.numpy_to_pil(_SCREAMING_SNAKE_CASE		)

						if not return_dict:
										return (image,)

						return ImagePipelineOutput(images=_SCREAMING_SNAKE_CASE		)













from collections.abc import Callable
class      UpperCamelCase :







		"""simple docstring"""



		def __init__(  self		:       Tuple						,_SCREAMING_SNAKE_CASE		:       Callable | None = None		)    ->   None:


						'''simple docstring'''
						# Stores actual heap items.
						A	 =      []
						# Stores indexes of each item for supporting updates and deletion.
						A	 =      {}
						# Stores current size of heap.
						A	 =      0
						# Stores function used to evaluate the score of an item on which basis ordering
						# will be done.
						A	 =      key or (lambda _SCREAMING_SNAKE_CASE		: x)
		def       A(  self		:       Tuple						,_SCREAMING_SNAKE_CASE		:       int		)    ->   int | None:


						'''simple docstring'''
						return int((i - 1) / 2		) if i > 0 else None
		def       A(  self		:       Union[str, Any]						,_SCREAMING_SNAKE_CASE		:       int		)    ->   int | None:


						'''simple docstring'''
						A	 =      int(2 * i + 1		)
						return left if 0 < left < self.size else None
		def       A(  self		:       Union[str, Any]						,_SCREAMING_SNAKE_CASE		:       int		)    ->   int | None:


						'''simple docstring'''
						A	 =      int(2 * i + 2		)
						return right if 0 < right < self.size else None
		def       A(  self		:       List[str]						,_SCREAMING_SNAKE_CASE		:       int						,_SCREAMING_SNAKE_CASE		:       int		)    ->   None:


						'''simple docstring'''
						A ,		A	 =      (
						    self.pos_map[self.arr[j][0]],
						    self.pos_map[self.arr[i][0]],
						)
						# Then swap the items in the list.
						A ,		A	 =      self.arr[j], self.arr[i]
		def       A(  self		:       List[str]						,_SCREAMING_SNAKE_CASE		:       int						,_SCREAMING_SNAKE_CASE		:       int		)    ->   bool:


						'''simple docstring'''
						return self.arr[i][1] < self.arr[j][1]
		def       A(  self		:       int						,_SCREAMING_SNAKE_CASE		:       int		)    ->   int:


						'''simple docstring'''
						A	 =      self._left(_SCREAMING_SNAKE_CASE		)
						A	 =      self._right(_SCREAMING_SNAKE_CASE		)
						A	 =      i

						if left is not None and not self._cmp(_SCREAMING_SNAKE_CASE						,_SCREAMING_SNAKE_CASE		):
										A	 =      left
						if right is not None and not self._cmp(_SCREAMING_SNAKE_CASE						,_SCREAMING_SNAKE_CASE		):
										A	 =      right

						return valid_parent
		def       A(  self		:       Any						,_SCREAMING_SNAKE_CASE		:       int		)    ->   None:


						'''simple docstring'''
						A	 =      self._parent(_SCREAMING_SNAKE_CASE		)
						while parent is not None and not self._cmp(_SCREAMING_SNAKE_CASE						,_SCREAMING_SNAKE_CASE		):
										self._swap(_SCREAMING_SNAKE_CASE						,_SCREAMING_SNAKE_CASE		)
										A ,		A	 =      parent, self._parent(_SCREAMING_SNAKE_CASE		)
		def       A(  self		:       List[Any]						,_SCREAMING_SNAKE_CASE		:       int		)    ->   None:


						'''simple docstring'''
						A	 =      self._get_valid_parent(_SCREAMING_SNAKE_CASE		)
						while valid_parent != index:
										self._swap(_SCREAMING_SNAKE_CASE						,_SCREAMING_SNAKE_CASE		)
										A ,		A	 =      valid_parent, self._get_valid_parent(_SCREAMING_SNAKE_CASE		)
		def       A(  self		:       Optional[Any]						,_SCREAMING_SNAKE_CASE		:       int						,_SCREAMING_SNAKE_CASE		:       int		)    ->   None:


						'''simple docstring'''
						if item not in self.pos_map:
										return
						A	 =      self.pos_map[item]
						A	 =      [item, self.key(_SCREAMING_SNAKE_CASE		)]
						# Make sure heap is right in both up and down direction.
						# Ideally only one of them will make any change.
						self._heapify_up(_SCREAMING_SNAKE_CASE		)
						self._heapify_down(_SCREAMING_SNAKE_CASE		)
		def       A(  self		:       int						,_SCREAMING_SNAKE_CASE		:       int		)    ->   None:


						'''simple docstring'''
						if item not in self.pos_map:
										return
						A	 =      self.pos_map[item]
						del self.pos_map[item]
						A	 =      self.arr[self.size - 1]
						A	 =      index
						self.size -= 1
						# Make sure heap is right in both up and down direction. Ideally only one
						# of them will make any change- so no performance loss in calling both.
						if self.size > index:
										self._heapify_up(_SCREAMING_SNAKE_CASE		)
										self._heapify_down(_SCREAMING_SNAKE_CASE		)
		def       A(  self		:       Optional[Any]						,_SCREAMING_SNAKE_CASE		:       int						,_SCREAMING_SNAKE_CASE		:       int		)    ->   None:


						'''simple docstring'''
						A	 =      len(self.arr		)
						if arr_len == self.size:
										self.arr.append([item, self.key(_SCREAMING_SNAKE_CASE		)]		)
						else:
										A	 =      [item, self.key(_SCREAMING_SNAKE_CASE		)]
						A	 =      self.size
						self.size += 1
						self._heapify_up(self.size - 1		)
		def       A(  self		:       str		)    ->   tuple | None:


						'''simple docstring'''
						return self.arr[0] if self.size else None







		def       A(  self		:       Any		)    ->   tuple | None:


						'''simple docstring'''
						A	 =      self.get_top()
						if top_item_tuple:
										self.delete_item(top_item_tuple[0]		)
						return top_item_tuple







def      snake_case		(							):
				pass


if __name__ == "__main__":
					import doctest

					doctest.testmod()





