




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


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

a      =							'\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 							lowercase					(snake_case__   : Optional[Any]      ,       snake_case__   : int      ,       snake_case__   : Optional[int]=8 )     ->  List[str]:


				'''simple docstring'''






				lowerCAmelCase  =				h // scale_factor**2
				if h % scale_factor**2 != 0:
								new_h += 1
				lowerCAmelCase  =				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					):


						def __init__(							self				:			int						,						lowerCAmelCase				:			Optional[Any]						,						lowerCAmelCase				:			Union[str, Any]						,						lowerCAmelCase				:			List[Any]						,						lowerCAmelCase				:			Union[str, Any]						,						lowerCAmelCase				:			str						,						):
										super().__init__()

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


						def   __lowercase (							self				:			int						,						lowerCAmelCase				:			str						,						lowerCAmelCase				:			Union[str, Any]						,						lowerCAmelCase				:			Tuple						,						lowerCAmelCase				:			Optional[Any]						,						lowerCAmelCase				:			Dict						,						lowerCAmelCase				:			Optional[int]					):
										if latents is None:
														lowerCAmelCase  =				randn_tensor(__a						,						generator=__a						,						device=__a						,						dtype=__a					)
										else:
														if latents.shape != shape:
																		raise ValueError(f'''Unexpected latents shape, got {latents.shape}, expected {shape}'''					)
														lowerCAmelCase  =				latents.to(__a					)

										lowerCAmelCase  =				latents * scheduler.init_noise_sigma
										return latents


						def   __lowercase (							self				:			Union[str, Any]						,						lowerCAmelCase				:			Dict						,						lowerCAmelCase				:			Dict						,						lowerCAmelCase				:			List[Any]						,						lowerCAmelCase				:			Optional[Any]						,						lowerCAmelCase				:			List[Any]=None						,						):
										lowerCAmelCase  =				len(__a					) if isinstance(__a						,						__a					) else 1
										# get prompt text embeddings
										lowerCAmelCase  =				self.tokenizer(
										    __a						,						padding="""max_length"""						,						truncation=__a						,						max_length=77						,						return_attention_mask=__a						,						add_special_tokens=__a						,						return_tensors="""pt"""						,						)

										lowerCAmelCase  =				text_inputs.input_ids
										lowerCAmelCase  =				self.tokenizer(__a						,						padding="""longest"""						,						return_tensors="""pt"""					).input_ids

										if untruncated_ids.shape[-1] >= text_input_ids.shape[-1] and not torch.equal(__a						,						__a					):
														lowerCAmelCase  =				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}'''					)

										lowerCAmelCase  =				text_input_ids.to(__a					)
										lowerCAmelCase  =				text_inputs.attention_mask.to(__a					)

										lowerCAmelCase ,						lowerCAmelCase  =				self.text_encoder(
										    input_ids=__a						,						attention_mask=__a					)

										lowerCAmelCase  =				prompt_embeds.repeat_interleave(__a						,						dim=0					)
										lowerCAmelCase  =				text_encoder_hidden_states.repeat_interleave(__a						,						dim=0					)
										lowerCAmelCase  =				text_mask.repeat_interleave(__a						,						dim=0					)

										if do_classifier_free_guidance:
														lowerCAmelCase  =				42
														if negative_prompt is None:
																		lowerCAmelCase  =				[""""""] * batch_size
														elif type(__a					) is not type(__a					):
																		raise TypeError(
																		    f'''`negative_prompt` should be the same type to `prompt`, but got {type(__a					)} !='''
																		    f''' {type(__a					)}.'''					)
														elif isinstance(__a						,						__a					):
																		lowerCAmelCase  =				[negative_prompt]
														elif batch_size != len(__a					):
																		raise ValueError(
																		    f'''`negative_prompt`: {negative_prompt} has batch size {len(__a					)}, but `prompt`:'''
																		    f''' {prompt} has batch size {batch_size}. Please make sure that passed `negative_prompt` matches'''
																		    """ the batch size of `prompt`."""					)
														else:
																		lowerCAmelCase  =				negative_prompt

														lowerCAmelCase  =				self.tokenizer(
														    __a						,						padding="""max_length"""						,						max_length=77						,						truncation=__a						,						return_attention_mask=__a						,						add_special_tokens=__a						,						return_tensors="""pt"""						,						)
														lowerCAmelCase  =				uncond_input.input_ids.to(__a					)
														lowerCAmelCase  =				uncond_input.attention_mask.to(__a					)

														lowerCAmelCase ,						lowerCAmelCase  =				self.text_encoder(
														    input_ids=__a						,						attention_mask=__a					)

														# duplicate unconditional embeddings for each generation per prompt, using mps friendly method

														lowerCAmelCase  =				negative_prompt_embeds.shape[1]
														lowerCAmelCase  =				negative_prompt_embeds.repeat(1						,						__a					)
														lowerCAmelCase  =				negative_prompt_embeds.view(batch_size * num_images_per_prompt						,						__a					)

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

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

										return prompt_embeds, text_encoder_hidden_states, text_mask


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

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

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


						def   __lowercase (							self				:			Dict						,						lowerCAmelCase				:			Optional[int]=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."""					)

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

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

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

										if self.safety_checker is not None:
														lowerCAmelCase ,						lowerCAmelCase  =				cpu_offload_with_hook(self.safety_checker						,						__a						,						prev_module_hook=__a					)

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


						@property
						# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline._execution_device
						def   __lowercase (							self				:			str					):
										if not hasattr(self.unet						,						"""_hf_hook"""					):
														return self.device
										for module in self.unet.modules():
														if (
														    hasattr(__a						,						"""_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(__a					)
						def __call__(							self				:			Optional[Any]						,						lowerCAmelCase				:			List[Any]						,						lowerCAmelCase				:			Optional[int]						,						lowerCAmelCase				:			Any						,						lowerCAmelCase				:			Any = None						,						lowerCAmelCase				:			List[Any] = 512						,						lowerCAmelCase				:			List[str] = 512						,						lowerCAmelCase				:			Union[str, Any] = 100						,						lowerCAmelCase				:			Tuple = 4.0						,						lowerCAmelCase				:			str = 1						,						lowerCAmelCase				:			Optional[int] = None						,						lowerCAmelCase				:			Optional[int] = None						,						lowerCAmelCase				:			List[str] = "pil"						,						lowerCAmelCase				:			List[str] = True						,						):

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

										lowerCAmelCase  =				self._execution_device

										lowerCAmelCase  =				batch_size * num_images_per_prompt
										lowerCAmelCase  =				guidance_scale > 1.0

										lowerCAmelCase ,						lowerCAmelCase ,						lowerCAmelCase  =				self._encode_prompt(
										    __a						,						__a						,						__a						,						__a						,						__a					)

										if isinstance(__a						,						__a					):
														lowerCAmelCase  =				torch.cat(__a						,						dim=0					)
										if isinstance(__a						,						__a					):
														lowerCAmelCase  =				torch.cat(__a						,						dim=0					)

										if do_classifier_free_guidance:
														lowerCAmelCase  =				image_embeds.repeat_interleave(__a						,						dim=0					)
														lowerCAmelCase  =				negative_image_embeds.repeat_interleave(__a						,						dim=0					)

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

										self.scheduler.set_timesteps(__a						,						device=__a					)
										lowerCAmelCase  =				self.scheduler.timesteps

										lowerCAmelCase  =				self.unet.config.in_channels

										lowerCAmelCase ,						lowerCAmelCase  =				get_new_h_w(__a						,						__a						,						self.movq_scale_factor					)

										# create initial latent
										lowerCAmelCase  =				self.prepare_latents(
										    (batch_size, num_channels_latents, height, width)						,						text_encoder_hidden_states.dtype						,						__a						,						__a						,						__a						,						self.scheduler						,						)

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

														lowerCAmelCase  =				{"""text_embeds""": prompt_embeds, """image_embeds""": image_embeds}
														lowerCAmelCase  =				self.unet(
														    sample=__a						,						timestep=__a						,						encoder_hidden_states=__a						,						added_cond_kwargs=__a						,						return_dict=__a						,						)[0]

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

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

										if output_type == "pil":
														lowerCAmelCase  =				self.numpy_to_pil(__a					)

										if not return_dict:
														return (image,)

										return ImagePipelineOutput(images=__a					)











"""simple docstring"""







from typing import TYPE_CHECKING

from ...utils import OptionalDependencyNotAvailable, _LazyModule, is_tokenizers_available, is_torch_available


a      =							{
    'configuration_biogpt': ['BIOGPT_PRETRAINED_CONFIG_ARCHIVE_MAP', 'BioGptConfig'],
    'tokenization_biogpt': ['BioGptTokenizer'],
}

try:
		if not is_torch_available():
				raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
		pass
else:
		a      =							[
		    'BIOGPT_PRETRAINED_MODEL_ARCHIVE_LIST',
		    'BioGptForCausalLM',
		    'BioGptForTokenClassification',
		    'BioGptForSequenceClassification',
		    'BioGptModel',
		    'BioGptPreTrainedModel',
		]


if TYPE_CHECKING:
		from .configuration_biogpt import BIOGPT_PRETRAINED_CONFIG_ARCHIVE_MAP, BioGptConfig
		from .tokenization_biogpt import BioGptTokenizer

		try:
				if not is_torch_available():
						raise OptionalDependencyNotAvailable()
		except OptionalDependencyNotAvailable:
				pass
		else:
				from .modeling_biogpt import (
				    BIOGPT_PRETRAINED_MODEL_ARCHIVE_LIST,
				    BioGptForCausalLM,
				    BioGptForSequenceClassification,
				    BioGptForTokenClassification,
				    BioGptModel,
				    BioGptPreTrainedModel,
				)


else:
		import sys

		a      =							_LazyModule(__name__, globals()['__file__'], _import_structure, module_spec=__spec__)





