


'''simple docstring'''

from typing import Callable, List, Optional, Tuple, Union

import torch
from transformers import CLIPTextModel, CLIPTokenizer

from ...configuration_utils import ConfigMixin, register_to_config
from ...models import ModelMixin, TransformeraDModel, VQModel
from ...schedulers import VQDiffusionScheduler
from ...utils import logging
from ..pipeline_utils import DiffusionPipeline, ImagePipelineOutput


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


class       a__				( __A							,						__A      ):







				"""simple docstring"""



				@register_to_config
				def __init__(self       ,			__lowercase       ,			__lowercase = None       ,			__lowercase = None		):
										super().__init__()

										__lowerCAmelCase							       =				learnable

										if self.learnable:
																assert hidden_size is not None, "learnable=True requires `hidden_size` to be set"
																assert length is not None, "learnable=True requires `length` to be set"

																__lowerCAmelCase							       =				torch.zeros(__lowercase       ,			__lowercase		)
										else:
																__lowerCAmelCase							       =				None

										__lowerCAmelCase							       =				torch.nn.Parameter(__lowercase		)


class       a__				( __A      ):







				"""simple docstring"""



				__UpperCamelCase    :		VQModel
				__UpperCamelCase    :		CLIPTextModel
				__UpperCamelCase    :		CLIPTokenizer
				__UpperCamelCase    :		TransformeraDModel
				__UpperCamelCase    :		LearnedClassifierFreeSamplingEmbeddings
				__UpperCamelCase    :		VQDiffusionScheduler




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

										self.register_modules(
										    vqvae=__lowercase       ,			transformer=__lowercase       ,			text_encoder=__lowercase       ,			tokenizer=__lowercase       ,			scheduler=__lowercase       ,			learned_classifier_free_sampling_embeddings=__lowercase       ,			)




				def 					_snake_case  (self       ,			__lowercase       ,			__lowercase       ,			__lowercase		):
										__lowerCAmelCase							       =				len(__lowercase		) if isinstance(__lowercase       ,			__lowercase		) else 1

										# get prompt text embeddings
										__lowerCAmelCase							       =				self.tokenizer(
										    __lowercase       ,			padding='''max_length'''       ,			max_length=self.tokenizer.model_max_length       ,			return_tensors='''pt'''       ,			)
										__lowerCAmelCase							       =				text_inputs.input_ids

										if text_input_ids.shape[-1] > self.tokenizer.model_max_length:
																__lowerCAmelCase							       =				self.tokenizer.batch_decode(text_input_ids[:, self.tokenizer.model_max_length :]		)
																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}"""		)
																__lowerCAmelCase							       =				text_input_ids[:, : self.tokenizer.model_max_length]
										__lowerCAmelCase							       =				self.text_encoder(text_input_ids.to(self.device		)		)[0]

										# NOTE: This additional step of normalizing the text embeddings is from VQ-Diffusion.
										# While CLIP does normalize the pooled output of the text transformer when combining
										# the image and text embeddings, CLIP does not directly normalize the last hidden state.
										#
										# CLIP normalizing the pooled output.
										# https://github.com/huggingface/transformers/blob/d92e22d1f28324f513f3080e5c47c071a3916721/src/transformers/models/clip/modeling_clip.py#L1052-L1053
										__lowerCAmelCase							       =				prompt_embeds / prompt_embeds.norm(dim=-1       ,			keepdim=__lowercase		)

										# duplicate text embeddings for each generation per prompt
										__lowerCAmelCase							       =				prompt_embeds.repeat_interleave(__lowercase       ,			dim=0		)

										if do_classifier_free_guidance:
																if self.learned_classifier_free_sampling_embeddings.learnable:
																						__lowerCAmelCase							       =				self.learned_classifier_free_sampling_embeddings.embeddings
																						__lowerCAmelCase							       =				negative_prompt_embeds.unsqueeze(0		).repeat(__lowercase       ,			1       ,			1		)
																else:
																						__lowerCAmelCase							       =				[''''''] * batch_size

																						__lowerCAmelCase							       =				text_input_ids.shape[-1]
																						__lowerCAmelCase							       =				self.tokenizer(
																						    __lowercase       ,			padding='''max_length'''       ,			max_length=__lowercase       ,			truncation=__lowercase       ,			return_tensors='''pt'''       ,			)
																						__lowerCAmelCase							       =				self.text_encoder(uncond_input.input_ids.to(self.device		)		)[0]
																						# See comment for normalizing text embeddings
																						__lowerCAmelCase							       =				negative_prompt_embeds / negative_prompt_embeds.norm(dim=-1       ,			keepdim=__lowercase		)

																# duplicate unconditional embeddings for each generation per prompt, using mps friendly method
																__lowerCAmelCase							       =				negative_prompt_embeds.shape[1]
																__lowerCAmelCase							       =				negative_prompt_embeds.repeat(1       ,			__lowercase       ,			1		)
																__lowerCAmelCase							       =				negative_prompt_embeds.view(batch_size * num_images_per_prompt       ,			__lowercase       ,			-1		)

																# 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
																__lowerCAmelCase							       =				torch.cat([negative_prompt_embeds, prompt_embeds]		)

										return prompt_embeds




				@torch.no_grad()
				def __call__(self       ,			__lowercase       ,			__lowercase = 1_00       ,			__lowercase = 5.0       ,			__lowercase = 1.0       ,			__lowercase = 1       ,			__lowercase = None       ,			__lowercase = None       ,			__lowercase = "pil"       ,			__lowercase = True       ,			__lowercase = None       ,			__lowercase = 1       ,			):
										if isinstance(__lowercase       ,			__lowercase		):
																__lowerCAmelCase							       =				1
										elif isinstance(__lowercase       ,			__lowercase		):
																__lowerCAmelCase							       =				len(__lowercase		)
										else:
																raise ValueError(F"""`prompt` has to be of type `str` or `list` but is {type(__lowercase		)}"""		)

										__lowerCAmelCase							       =				batch_size * num_images_per_prompt

										__lowerCAmelCase							       =				guidance_scale > 1.0

										__lowerCAmelCase							       =				self._encode_prompt(__lowercase       ,			__lowercase       ,			__lowercase		)

										if (callback_steps is None) or (
										    callback_steps is not None and (not isinstance(__lowercase       ,			__lowercase		) or callback_steps <= 0)
										):
																raise ValueError(
																    F"""`callback_steps` has to be a positive integer but is {callback_steps} of type"""
																    F""" {type(__lowercase		)}."""		)

										# get the initial completely masked latents unless the user supplied it

										__lowerCAmelCase							       =				(batch_size, self.transformer.num_latent_pixels)
										if latents is None:
																__lowerCAmelCase							       =				self.transformer.num_vector_embeds - 1
																__lowerCAmelCase							       =				torch.full(__lowercase       ,			__lowercase		).to(self.device		)
										else:
																if latents.shape != latents_shape:
																						raise ValueError(F"""Unexpected latents shape, got {latents.shape}, expected {latents_shape}"""		)
																if (latents < 0).any() or (latents >= self.transformer.num_vector_embeds).any():
																						raise ValueError(
																						    '''Unexpected latents value(s). All latents be valid embedding indices i.e. in the range 0,'''
																						    F""" {self.transformer.num_vector_embeds - 1} (inclusive)."""		)
																__lowerCAmelCase							       =				latents.to(self.device		)

										# set timesteps
										self.scheduler.set_timesteps(__lowercase       ,			device=self.device		)

										__lowerCAmelCase							       =				self.scheduler.timesteps.to(self.device		)

										__lowerCAmelCase							       =				latents

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

																# predict the un-noised image
																# model_output == `log_p_x_0`
																__lowerCAmelCase							       =				self.transformer(__lowercase       ,			encoder_hidden_states=__lowercase       ,			timestep=__lowercase		).sample

																if do_classifier_free_guidance:
																						__lowerCAmelCase       ,					__lowerCAmelCase							       =				model_output.chunk(2		)
																						__lowerCAmelCase							       =				model_output_uncond + guidance_scale * (model_output_text - model_output_uncond)
																						model_output -= torch.logsumexp(__lowercase       ,			dim=1       ,			keepdim=__lowercase		)

																__lowerCAmelCase							       =				self.truncate(__lowercase       ,			__lowercase		)

																# remove `log(0)`'s (`-inf`s)
																__lowerCAmelCase							       =				model_output.clamp(-70		)

																# compute the previous noisy sample x_t -> x_t-1
																__lowerCAmelCase							       =				self.scheduler.step(__lowercase       ,			timestep=__lowercase       ,			sample=__lowercase       ,			generator=__lowercase		).prev_sample

																# call the callback, if provided
																if callback is not None and i % callback_steps == 0:
																						callback(__lowercase       ,			__lowercase       ,			__lowercase		)

										__lowerCAmelCase							       =				self.vqvae.config.vq_embed_dim
										__lowerCAmelCase							       =				(batch_size, self.transformer.height, self.transformer.width, embedding_channels)
										__lowerCAmelCase							       =				self.vqvae.quantize.get_codebook_entry(__lowercase       ,			shape=__lowercase		)
										__lowerCAmelCase							       =				self.vqvae.decode(__lowercase       ,			force_not_quantize=__lowercase		).sample

										__lowerCAmelCase							       =				(image / 2 + 0.5).clamp(0       ,			1		)
										__lowerCAmelCase							       =				image.cpu().permute(0       ,			2       ,			3       ,			1		).numpy()

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

										if not return_dict:
																return (image,)

										return ImagePipelineOutput(images=__lowercase		)




				def 					_snake_case  (self       ,			__lowercase       ,			__lowercase		):
										__lowerCAmelCase       ,					__lowerCAmelCase							       =				torch.sort(__lowercase       ,			1       ,			descending=__lowercase		)
										__lowerCAmelCase							       =				torch.exp(__lowercase		)
										__lowerCAmelCase							       =				sorted_p_x_0.cumsum(dim=1		) < truncation_rate

										# Ensure that at least the largest probability is not zeroed out
										__lowerCAmelCase							       =				torch.full_like(keep_mask[:, 0:1, :]       ,			__lowercase		)
										__lowerCAmelCase							       =				torch.cat((all_true, keep_mask)       ,			dim=1		)
										__lowerCAmelCase							       =				keep_mask[:, :-1, :]

										__lowerCAmelCase							       =				keep_mask.gather(1       ,			indices.argsort(1		)		)

										__lowerCAmelCase							       =				log_p_x_0.clone()

										__lowerCAmelCase							       =				-torch.inf  # -inf = log(0)

										return rv





'''simple docstring'''

from typing import Dict, List, Optional, Tuple, Union

import torch

from ...models import AutoencoderKL, TransformeraDModel
from ...schedulers import KarrasDiffusionSchedulers
from ...utils import randn_tensor
from ..pipeline_utils import DiffusionPipeline, ImagePipelineOutput


class       a__				( __A      ):







				"""simple docstring"""



				def __init__(self       ,			__lowercase       ,			__lowercase       ,			__lowercase       ,			__lowercase = None       ,			):
										super().__init__()
										self.register_modules(transformer=__lowercase       ,			vae=__lowercase       ,			scheduler=__lowercase		)

										# create a imagenet -> id dictionary for easier use
										__lowerCAmelCase							       =				{}
										if idalabel is not None:
																for key, value in idalabel.items():
																						for label in value.split(''','''		):
																												__lowerCAmelCase							       =				int(__lowercase		)
																__lowerCAmelCase							       =				dict(sorted(self.labels.items()		)		)




				def 					_snake_case  (self       ,			__lowercase		):

										if not isinstance(__lowercase       ,			__lowercase		):
																__lowerCAmelCase							       =				list(__lowercase		)

										for l in label:
																if l not in self.labels:
																						raise ValueError(
																						    F"""{l} does not exist. Please make sure to select one of the following labels: \n {self.labels}."""		)

										return [self.labels[l] for l in label]




				@torch.no_grad()
				def __call__(self       ,			__lowercase       ,			__lowercase = 4.0       ,			__lowercase = None       ,			__lowercase = 50       ,			__lowercase = "pil"       ,			__lowercase = True       ,			):
										__lowerCAmelCase							       =				len(__lowercase		)
										__lowerCAmelCase							       =				self.transformer.config.sample_size
										__lowerCAmelCase							       =				self.transformer.config.in_channels

										__lowerCAmelCase							       =				randn_tensor(
										    shape=(batch_size, latent_channels, latent_size, latent_size)       ,			generator=__lowercase       ,			device=self.device       ,			dtype=self.transformer.dtype       ,			)
										__lowerCAmelCase							       =				torch.cat([latents] * 2		) if guidance_scale > 1 else latents

										__lowerCAmelCase							       =				torch.tensor(__lowercase       ,			device=self.device		).reshape(-1		)
										__lowerCAmelCase							       =				torch.tensor([10_00] * batch_size       ,			device=self.device		)
										__lowerCAmelCase							       =				torch.cat([class_labels, class_null]       ,			0		) if guidance_scale > 1 else class_labels

										# set step values
										self.scheduler.set_timesteps(__lowercase		)

										for t in self.progress_bar(self.scheduler.timesteps		):
																if guidance_scale > 1:
																						__lowerCAmelCase							       =				latent_model_input[: len(__lowercase		) // 2]
																						__lowerCAmelCase							       =				torch.cat([half, half]       ,			dim=0		)
																__lowerCAmelCase							       =				self.scheduler.scale_model_input(__lowercase       ,			__lowercase		)

																__lowerCAmelCase							       =				t
																if not torch.is_tensor(__lowercase		):
																						# TODO: this requires sync between CPU and GPU. So try to pass timesteps as tensors if you can
																						# This would be a good case for the `match` statement (Python 3.10+)
																						__lowerCAmelCase							       =				latent_model_input.device.type == '''mps'''
																						if isinstance(__lowercase       ,			__lowercase		):
																												__lowerCAmelCase							       =				torch.floataa if is_mps else torch.floataa
																						else:
																												__lowerCAmelCase							       =				torch.intaa if is_mps else torch.intaa
																						__lowerCAmelCase							       =				torch.tensor([timesteps]       ,			dtype=__lowercase       ,			device=latent_model_input.device		)
																elif len(timesteps.shape		) == 0:
																						__lowerCAmelCase							       =				timesteps[None].to(latent_model_input.device		)
																# broadcast to batch dimension in a way that's compatible with ONNX/Core ML
																__lowerCAmelCase							       =				timesteps.expand(latent_model_input.shape[0]		)
																# predict noise model_output
																__lowerCAmelCase							       =				self.transformer(
																    __lowercase       ,			timestep=__lowercase       ,			class_labels=__lowercase		).sample

																# perform guidance
																if guidance_scale > 1:
																						__lowerCAmelCase       ,					__lowerCAmelCase							       =				noise_pred[:, :latent_channels], noise_pred[:, latent_channels:]
																						__lowerCAmelCase       ,					__lowerCAmelCase							       =				torch.split(__lowercase       ,			len(__lowercase		) // 2       ,			dim=0		)

																						__lowerCAmelCase							       =				uncond_eps + guidance_scale * (cond_eps - uncond_eps)
																						__lowerCAmelCase							       =				torch.cat([half_eps, half_eps]       ,			dim=0		)

																						__lowerCAmelCase							       =				torch.cat([eps, rest]       ,			dim=1		)

																# learned sigma
																if self.transformer.config.out_channels // 2 == latent_channels:
																						__lowerCAmelCase       ,					__lowerCAmelCase							       =				torch.split(__lowercase       ,			__lowercase       ,			dim=1		)
																else:
																						__lowerCAmelCase							       =				noise_pred

																# compute previous image: x_t -> x_t-1
																__lowerCAmelCase							       =				self.scheduler.step(__lowercase       ,			__lowercase       ,			__lowercase		).prev_sample

										if guidance_scale > 1:
																__lowerCAmelCase       ,					__lowerCAmelCase							       =				latent_model_input.chunk(2       ,			dim=0		)
										else:
																__lowerCAmelCase							       =				latent_model_input

										__lowerCAmelCase							       =				1 / self.vae.config.scaling_factor * latents
										__lowerCAmelCase							       =				self.vae.decode(__lowercase		).sample

										__lowerCAmelCase							       =				(samples / 2 + 0.5).clamp(0       ,			1		)

										# we always cast to float32 as this does not cause significant overhead and is compatible with bfloat16
										__lowerCAmelCase							       =				samples.cpu().permute(0       ,			2       ,			3       ,			1		).float().numpy()

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

										if not return_dict:
																return (samples,)

										return ImagePipelineOutput(images=__lowercase		)

