




from typing import List, Optional, Union

import numpy as np
import PIL
import torch
from PIL import Image

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


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

UpperCAmelCase__				     = "\n    Examples:\n        ```py\n        >>> from diffusers import KandinskyV22Img2ImgPipeline, KandinskyV22PriorPipeline\n        >>> from diffusers.utils import load_image\n        >>> import torch\n\n        >>> pipe_prior = KandinskyV22PriorPipeline.from_pretrained(\n        ...     \"kandinsky-community/kandinsky-2-2-prior\", torch_dtype=torch.float16\n        ... )\n        >>> pipe_prior.to(\"cuda\")\n\n        >>> prompt = \"A red cartoon frog, 4k\"\n        >>> image_emb, zero_image_emb = pipe_prior(prompt, return_dict=False)\n\n        >>> pipe = KandinskyV22Img2ImgPipeline.from_pretrained(\n        ...     \"kandinsky-community/kandinsky-2-2-decoder\", torch_dtype=torch.float16\n        ... )\n        >>> pipe.to(\"cuda\")\n\n        >>> init_image = load_image(\n        ...     \"https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main\"\n        ...     \"/kandinsky/frog.png\"\n        ... )\n\n        >>> image = pipe(\n        ...     image=init_image,\n        ...     image_embeds=image_emb,\n        ...     negative_image_embeds=zero_image_emb,\n        ...     height=768,\n        ...     width=768,\n        ...     num_inference_steps=100,\n        ...     strength=0.2,\n        ... ).images\n\n        >>> image[0].save(\"red_frog.png\")\n        ```\n"







def 		_A(       UpperCamelCase__							:			List[Any]				,       UpperCamelCase__							:			List[Any]				,       UpperCamelCase__							:			Any=8				)    -> List[str]:
							'''simple docstring'''


							__lowercase      =    height // scale_factor**2
							if height % scale_factor**2 != 0:
														new_height += 1
							__lowercase      =    width // scale_factor**2
							if width % scale_factor**2 != 0:
														new_width += 1
							return new_height * scale_factor, new_width * scale_factor







def 		_A(       UpperCamelCase__							:			int				,       UpperCamelCase__							:			Optional[int]=512				,       UpperCamelCase__							:			Dict=512				)    -> List[str]:
							'''simple docstring'''


							__lowercase      =    pil_image.resize((w, h)				,       resample=Image.BICUBIC				,       reducing_gap=1				)
							__lowercase      =    np.array(pil_image.convert('''RGB'''				)				)
							__lowercase      =    arr.astype(np.floataa				) / 127.5 - 1
							__lowercase      =    np.transpose(UpperCamelCase__				,       [2, 0, 1]				)
							__lowercase      =    torch.from_numpy(UpperCamelCase__				).unsqueeze(0				)
							return image







class a (     __SCREAMING_SNAKE_CASE  ):
						"""simple docstring"""

						def __init__( self :     int				,							lowerCamelCase__ :     UNetaDConditionModel				,							lowerCamelCase__ :     DDPMScheduler				,							lowerCamelCase__ :     VQModel				,							)    ->						Any:
													"""simple docstring"""


													super().__init__()

													self.register_modules(
													    unet=lowerCamelCase__				,							scheduler=lowerCamelCase__				,							movq=lowerCamelCase__				,							)
													__lowercase      =    2 ** (len(self.movq.config.block_out_channels				) - 1)






						def     UpperCAmelCase_      ( self :     List[str]				,							lowerCamelCase__ :     Union[str, Any]				,							lowerCamelCase__ :     Union[str, Any]				,							lowerCamelCase__ :     str				)    ->						List[str]:
													"""simple docstring"""


													__lowercase      =    min(int(num_inference_steps * strength				)				,							lowerCamelCase__				)

													__lowercase      =    max(num_inference_steps - init_timestep				,							0				)
													__lowercase      =    self.scheduler.timesteps[t_start:]

													return timesteps, num_inference_steps - t_start






						def     UpperCAmelCase_      ( self :     List[str]				,							lowerCamelCase__ :     Union[str, Any]				,							lowerCamelCase__ :     Union[str, Any]				,							lowerCamelCase__ :     Dict				,							lowerCamelCase__ :     Optional[Any]				,							lowerCamelCase__ :     List[str]				,							lowerCamelCase__ :     Tuple				,							lowerCamelCase__ :     Dict=None				)    ->						str:
													"""simple docstring"""


													if not isinstance(lowerCamelCase__				,							(torch.Tensor, PIL.Image.Image, list)				):
																				raise ValueError(
																				    f'`image` has to be of type `torch.Tensor`, `PIL.Image.Image` or list but is {type(lowerCamelCase__				)}'				)

													__lowercase      =    image.to(device=lowerCamelCase__				,							dtype=lowerCamelCase__				)

													__lowercase      =    batch_size * num_images_per_prompt

													if image.shape[1] == 4:
																				__lowercase      =    image

													else:
																				if isinstance(lowerCamelCase__				,							lowerCamelCase__				) and len(lowerCamelCase__				) != batch_size:
																											raise ValueError(
																											    f'You have passed a list of generators of length {len(lowerCamelCase__				)}, but requested an effective batch'
																											    f' size of {batch_size}. Make sure the batch size matches the length of the generators.'				)

																				elif isinstance(lowerCamelCase__				,							lowerCamelCase__				):
																											__lowercase      =    [
																											    self.movq.encode(image[i : i + 1]				).latent_dist.sample(generator[i]				) for i in range(lowerCamelCase__				)
																											]
																											__lowercase      =    torch.cat(lowerCamelCase__				,							dim=0				)
																				else:
																											__lowercase      =    self.movq.encode(lowerCamelCase__				).latent_dist.sample(lowerCamelCase__				)

																				__lowercase      =    self.movq.config.scaling_factor * init_latents

													__lowercase      =    torch.cat([init_latents]				,							dim=0				)

													__lowercase      =    init_latents.shape
													__lowercase      =    randn_tensor(lowerCamelCase__				,							generator=lowerCamelCase__				,							device=lowerCamelCase__				,							dtype=lowerCamelCase__				)

													# get latents
													__lowercase      =    self.scheduler.add_noise(lowerCamelCase__				,							lowerCamelCase__				,							lowerCamelCase__				)

													__lowercase      =    init_latents

													return latents






						def     UpperCAmelCase_      ( self :     Optional[Any]				,							lowerCamelCase__ :     int=0				)    ->						Tuple:
													"""simple docstring"""


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

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

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






						def     UpperCAmelCase_      ( self :     Optional[int]				,							lowerCamelCase__ :     Tuple=0				)    ->						str:
													"""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.'''				)

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

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

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






						@property
						# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline._execution_device
						def     UpperCAmelCase_      ( self :     List[Any]				)    ->						Dict:
													"""simple docstring"""


													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 :     int				,							lowerCamelCase__ :     Union[torch.FloatTensor, List[torch.FloatTensor]]				,							lowerCamelCase__ :     Union[torch.FloatTensor, PIL.Image.Image, List[torch.FloatTensor], List[PIL.Image.Image]]				,							lowerCamelCase__ :     Union[torch.FloatTensor, List[torch.FloatTensor]]				,							lowerCamelCase__ :     int = 512				,							lowerCamelCase__ :     int = 512				,							lowerCamelCase__ :     int = 100				,							lowerCamelCase__ :     float = 4.0				,							lowerCamelCase__ :     float = 0.3				,							lowerCamelCase__ :     int = 1				,							lowerCamelCase__ :     Optional[Union[torch.Generator, List[torch.Generator]]] = None				,							lowerCamelCase__ :     Optional[str] = "pil"				,							lowerCamelCase__ :     bool = True				,							)    ->						List[str]:
													"""simple docstring"""


													__lowercase      =    self._execution_device

													__lowercase      =    guidance_scale > 1.0

													if isinstance(lowerCamelCase__				,							lowerCamelCase__				):
																				__lowercase      =    torch.cat(lowerCamelCase__				,							dim=0				)
													__lowercase      =    image_embeds.shape[0]
													if isinstance(lowerCamelCase__				,							lowerCamelCase__				):
																				__lowercase      =    torch.cat(lowerCamelCase__				,							dim=0				)

													if do_classifier_free_guidance:
																				__lowercase      =    image_embeds.repeat_interleave(lowerCamelCase__				,							dim=0				)
																				__lowercase      =    negative_image_embeds.repeat_interleave(lowerCamelCase__				,							dim=0				)

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

													if not isinstance(lowerCamelCase__				,							lowerCamelCase__				):
																				__lowercase      =    [image]
													if not all(isinstance(lowerCamelCase__				,							(PIL.Image.Image, torch.Tensor)				) for i in image				):
																				raise ValueError(
																				    f'Input is in incorrect format: {[type(lowerCamelCase__				) for i in image]}. Currently, we only support  PIL image and pytorch tensor'				)

													__lowercase      =    torch.cat([prepare_image(lowerCamelCase__				,							lowerCamelCase__				,							lowerCamelCase__				) for i in image]				,							dim=0				)
													__lowercase      =    image.to(dtype=image_embeds.dtype				,							device=lowerCamelCase__				)

													__lowercase      =    self.movq.encode(lowerCamelCase__				)['''latents''']
													__lowercase      =    latents.repeat_interleave(lowerCamelCase__				,							dim=0				)
													self.scheduler.set_timesteps(lowerCamelCase__				,							device=lowerCamelCase__				)
													__lowercase     ,     __lowercase      =    self.get_timesteps(lowerCamelCase__				,							lowerCamelCase__				,							lowerCamelCase__				)
													__lowercase      =    timesteps[:1].repeat(batch_size * num_images_per_prompt				)
													__lowercase     ,     __lowercase      =    downscale_height_and_width(lowerCamelCase__				,							lowerCamelCase__				,							self.movq_scale_factor				)
													__lowercase      =    self.prepare_latents(
													    lowerCamelCase__				,							lowerCamelCase__				,							lowerCamelCase__				,							lowerCamelCase__				,							image_embeds.dtype				,							lowerCamelCase__				,							lowerCamelCase__				)
													for i, t in enumerate(self.progress_bar(lowerCamelCase__				)				):
																				# expand the latents if we are doing classifier free guidance
																				__lowercase      =    torch.cat([latents] * 2				) if do_classifier_free_guidance else latents

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

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

																				# compute the previous noisy sample x_t -> x_t-1
																				__lowercase      =    self.scheduler.step(
																				    lowerCamelCase__				,							lowerCamelCase__				,							lowerCamelCase__				,							generator=lowerCamelCase__				,							)[0]

													# post-processing
													__lowercase      =    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"]:
																				__lowercase      =    image * 0.5 + 0.5
																				__lowercase      =    image.clamp(0				,							1				)
																				__lowercase      =    image.cpu().permute(0				,							2				,							3				,							1				).float().numpy()

													if output_type == "pil":
																				__lowercase      =    self.numpy_to_pil(lowerCamelCase__				)

													if not return_dict:
																				return (image,)

													return ImagePipelineOutput(images=lowerCamelCase__				)












import tensorflow as tf

from ...tf_utils import shape_list







class a (     tf.keras.layers.Layer  ):
						"""simple docstring"""

						def __init__( self :     int				,							lowerCamelCase__ :     Union[str, Any]				,							lowerCamelCase__ :     Tuple				,							lowerCamelCase__ :     List[Any]				,							lowerCamelCase__ :     Optional[Any]				,							lowerCamelCase__ :     Tuple=1				,							lowerCamelCase__ :     Optional[Any]=False				,							**lowerCamelCase__ :     Optional[int]				)    ->						str:
													"""simple docstring"""


													super().__init__(**lowerCamelCase__				)

													__lowercase      =    vocab_size
													__lowercase      =    d_embed
													__lowercase      =    d_proj

													__lowercase      =    cutoffs + [vocab_size]
													__lowercase      =    [0] + self.cutoffs
													__lowercase      =    div_val

													__lowercase      =    self.cutoffs[0]
													__lowercase      =    len(self.cutoffs				) - 1
													__lowercase      =    self.shortlist_size + self.n_clusters
													__lowercase      =    keep_order

													__lowercase      =    []
													__lowercase      =    []






						def     UpperCAmelCase_      ( self :     Optional[int]				,							lowerCamelCase__ :     Tuple				)    ->						Tuple:
													"""simple docstring"""


													if self.n_clusters > 0:
																				__lowercase      =    self.add_weight(
																				    shape=(self.n_clusters, self.d_embed)				,							initializer='''zeros'''				,							trainable=lowerCamelCase__				,							name='''cluster_weight'''				)
																				__lowercase      =    self.add_weight(
																				    shape=(self.n_clusters,)				,							initializer='''zeros'''				,							trainable=lowerCamelCase__				,							name='''cluster_bias'''				)

													if self.div_val == 1:
																				for i in range(len(self.cutoffs				)				):
																											if self.d_proj != self.d_embed:
																																		__lowercase      =    self.add_weight(
																																		    shape=(self.d_embed, self.d_proj)				,							initializer='''zeros'''				,							trainable=lowerCamelCase__				,							name=f'out_projs_._{i}'				,							)
																																		self.out_projs.append(lowerCamelCase__				)
																											else:
																																		self.out_projs.append(lowerCamelCase__				)
																											__lowercase      =    self.add_weight(
																											    shape=(self.vocab_size, self.d_embed)				,							initializer='''zeros'''				,							trainable=lowerCamelCase__				,							name=f'out_layers_._{i}_._weight'				,							)
																											__lowercase      =    self.add_weight(
																											    shape=(self.vocab_size,)				,							initializer='''zeros'''				,							trainable=lowerCamelCase__				,							name=f'out_layers_._{i}_._bias'				,							)
																											self.out_layers.append((weight, bias)				)
													else:
																				for i in range(len(self.cutoffs				)				):
																											__lowercase     ,     __lowercase      =    self.cutoff_ends[i], self.cutoff_ends[i + 1]
																											__lowercase      =    self.d_embed // (self.div_val**i)

																											__lowercase      =    self.add_weight(
																											    shape=(d_emb_i, self.d_proj)				,							initializer='''zeros'''				,							trainable=lowerCamelCase__				,							name=f'out_projs_._{i}'				)
																											self.out_projs.append(lowerCamelCase__				)
																											__lowercase      =    self.add_weight(
																											    shape=(r_idx - l_idx, d_emb_i)				,							initializer='''zeros'''				,							trainable=lowerCamelCase__				,							name=f'out_layers_._{i}_._weight'				,							)
																											__lowercase      =    self.add_weight(
																											    shape=(r_idx - l_idx,)				,							initializer='''zeros'''				,							trainable=lowerCamelCase__				,							name=f'out_layers_._{i}_._bias'				,							)
																											self.out_layers.append((weight, bias)				)
													super().build(lowerCamelCase__				)






						@staticmethod
						def     UpperCAmelCase_      ( lowerCamelCase__ :     List[str]				,							lowerCamelCase__ :     int				,							lowerCamelCase__ :     List[Any]				,							lowerCamelCase__ :     List[str]=None				)    ->						Tuple:
													"""simple docstring"""


													__lowercase      =    x
													if proj is not None:
																				__lowercase      =    tf.einsum('''ibd,ed->ibe'''				,							lowerCamelCase__				,							lowerCamelCase__				)
													return tf.einsum('''ibd,nd->ibn'''				,							lowerCamelCase__				,							lowerCamelCase__				) + b






						@staticmethod
						def     UpperCAmelCase_      ( lowerCamelCase__ :     Tuple				,							lowerCamelCase__ :     List[str]				)    ->						str:
													"""simple docstring"""


													__lowercase      =    shape_list(lowerCamelCase__				)
													__lowercase      =    tf.range(lp_size[0]				,							dtype=target.dtype				)
													__lowercase      =    tf.stack([r, target]				,							1				)
													return tf.gather_nd(lowerCamelCase__				,							lowerCamelCase__				)






						def     UpperCAmelCase_      ( self :     Optional[int]				,							lowerCamelCase__ :     Tuple				,							lowerCamelCase__ :     str				,							lowerCamelCase__ :     Optional[int]=True				,							lowerCamelCase__ :     Optional[int]=False				)    ->						List[str]:
													"""simple docstring"""


													__lowercase      =    0
													if self.n_clusters == 0:
																				__lowercase      =    self._logit(lowerCamelCase__				,							self.out_layers[0][0]				,							self.out_layers[0][1]				,							self.out_projs[0]				)
																				if target is not None:
																											__lowercase      =    tf.nn.sparse_softmax_cross_entropy_with_logits(labels=lowerCamelCase__				,							logits=lowerCamelCase__				)
																				__lowercase      =    tf.nn.log_softmax(lowerCamelCase__				,							axis=-1				)
													else:
																				__lowercase      =    shape_list(lowerCamelCase__				)
																				__lowercase      =    []
																				__lowercase      =    tf.zeros(hidden_sizes[:2]				)
																				for i in range(len(self.cutoffs				)				):
																											__lowercase     ,     __lowercase      =    self.cutoff_ends[i], self.cutoff_ends[i + 1]
																											if target is not None:
																																		__lowercase      =    (target >= l_idx) & (target < r_idx)
																																		__lowercase      =    tf.where(lowerCamelCase__				)
																																		__lowercase      =    tf.boolean_mask(lowerCamelCase__				,							lowerCamelCase__				) - l_idx

																											if self.div_val == 1:
																																		__lowercase      =    self.out_layers[0][0][l_idx:r_idx]
																																		__lowercase      =    self.out_layers[0][1][l_idx:r_idx]
																											else:
																																		__lowercase      =    self.out_layers[i][0]
																																		__lowercase      =    self.out_layers[i][1]

																											if i == 0:
																																		__lowercase      =    tf.concat([cur_W, self.cluster_weight]				,							0				)
																																		__lowercase      =    tf.concat([cur_b, self.cluster_bias]				,							0				)

																																		__lowercase      =    self._logit(lowerCamelCase__				,							lowerCamelCase__				,							lowerCamelCase__				,							self.out_projs[0]				)
																																		__lowercase      =    tf.nn.log_softmax(lowerCamelCase__				)
																																		out.append(head_logprob[..., : self.cutoffs[0]]				)
																																		if target is not None:
																																									__lowercase      =    tf.boolean_mask(lowerCamelCase__				,							lowerCamelCase__				)
																																									__lowercase      =    self._gather_logprob(lowerCamelCase__				,							lowerCamelCase__				)
																											else:
																																		__lowercase      =    self._logit(lowerCamelCase__				,							lowerCamelCase__				,							lowerCamelCase__				,							self.out_projs[i]				)
																																		__lowercase      =    tf.nn.log_softmax(lowerCamelCase__				)
																																		__lowercase      =    self.cutoffs[0] + i - 1  # No probability for the head cluster
																																		__lowercase      =    head_logprob[..., cluster_prob_idx, None] + tail_logprob
																																		out.append(lowerCamelCase__				)
																																		if target is not None:
																																									__lowercase      =    tf.boolean_mask(lowerCamelCase__				,							lowerCamelCase__				)
																																									__lowercase      =    tf.boolean_mask(lowerCamelCase__				,							lowerCamelCase__				)
																																									__lowercase      =    self._gather_logprob(lowerCamelCase__				,							lowerCamelCase__				)
																																									cur_logprob += cur_head_logprob[:, self.cutoff_ends[1] + i - 1]
																											if target is not None:
																																		loss += tf.scatter_nd(lowerCamelCase__				,							-cur_logprob				,							shape_list(lowerCamelCase__				)				)
																				__lowercase      =    tf.concat(lowerCamelCase__				,							axis=-1				)

													if target is not None:
																				if return_mean:
																											__lowercase      =    tf.reduce_mean(lowerCamelCase__				)
																				# Add the training-time loss value to the layer using `self.add_loss()`.
																				self.add_loss(lowerCamelCase__				)

																				# Log the loss as a metric (we could log arbitrary metrics,
																				# including different metrics for training and inference.
																				self.add_metric(lowerCamelCase__				,							name=self.name				,							aggregation='''mean''' if return_mean else ''''''				)

													return out






