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


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

lowerCAmelCase_     =	'''
    Examples:
        ```py
        >>> from diffusers import KandinskyPipeline, KandinskyPriorPipeline
        >>> import torch

        >>> pipe_prior = KandinskyPriorPipeline.from_pretrained("kandinsky-community/Kandinsky-2-1-prior")
        >>> pipe_prior.to("cuda")

        >>> prompt = "red cat, 4k photo"
        >>> out = pipe_prior(prompt)
        >>> image_emb = out.image_embeds
        >>> negative_image_emb = out.negative_image_embeds

        >>> pipe = KandinskyPipeline.from_pretrained("kandinsky-community/kandinsky-2-1")
        >>> pipe.to("cuda")

        >>> image = pipe(
        ...     prompt,
        ...     image_embeds=image_emb,
        ...     negative_image_embeds=negative_image_emb,
        ...     height=768,
        ...     width=768,
        ...     num_inference_steps=100,
        ... ).images

        >>> image[0].save("cat.png")
        ```
'''



def      lowerCamelCase_(__SCREAMING_SNAKE_CASE     ,	__SCREAMING_SNAKE_CASE     ,	__SCREAMING_SNAKE_CASE=8  )->    Union[str, Any]:
							_SCREAMING_SNAKE_CASE		:						int			  =    h // scale_factor**2
							if h % scale_factor**2 != 0:
														new_h += 1
							_SCREAMING_SNAKE_CASE		:						Optional[int]			  =    w // scale_factor**2
							if w % scale_factor**2 != 0:
														new_w += 1
							return new_h * scale_factor, new_w * scale_factor




class  _snake_case		(    __snake_case ):







				"""simple docstring"""






				def __init__(      self							:							int      ,    _A							:							MultilingualCLIP      ,    _A							:							XLMRobertaTokenizer      ,    _A							:							UNetaDConditionModel      ,    _A							:							Union[DDIMScheduler, DDPMScheduler]      ,    _A							:							VQModel      ,    ):



											"""simple docstring"""







											super().__init__()

											self.register_modules(
											    text_encoder=_A      ,    tokenizer=_A      ,    unet=_A      ,    scheduler=_A      ,    movq=_A      ,    )
											_SCREAMING_SNAKE_CASE		:						Dict			  =    2 ** (len(self.movq.config.block_out_channels) - 1)





				def 		_lowerCAmelCase		(      self							:							Optional[Any]      ,    _A							:							Tuple      ,    _A							:							Optional[Any]      ,    _A							:							Optional[Any]      ,    _A							:							Optional[int]      ,    _A							:							Tuple      ,    _A							:							Tuple):



											"""simple docstring"""







											if latents is None:
																		_SCREAMING_SNAKE_CASE		:						List[Any]			  =    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}""")
																		_SCREAMING_SNAKE_CASE		:						Union[str, Any]			  =    latents.to(_A)

											_SCREAMING_SNAKE_CASE		:						int			  =    latents * scheduler.init_noise_sigma
											return latents





				def 		_lowerCAmelCase		(      self							:							Optional[int]      ,    _A							:							str      ,    _A							:							List[Any]      ,    _A							:							Tuple      ,    _A							:							Optional[int]      ,    _A							:							Any=None      ,    ):



											"""simple docstring"""







											_SCREAMING_SNAKE_CASE		:						Union[str, Any]			  =    len(_A) if isinstance(_A      ,    _A) else 1
											# get prompt text embeddings
											_SCREAMING_SNAKE_CASE		:						Union[str, Any]			  =    self.tokenizer(
											    _A      ,    padding="""max_length"""      ,    truncation=_A      ,    max_length=7_7      ,    return_attention_mask=_A      ,    add_special_tokens=_A      ,    return_tensors="""pt"""      ,    )

											_SCREAMING_SNAKE_CASE		:						List[str]			  =    text_inputs.input_ids
											_SCREAMING_SNAKE_CASE		:						List[Any]			  =    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):
																		_SCREAMING_SNAKE_CASE		:						Union[str, Any]			  =    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		:						Optional[Any]			  =    text_input_ids.to(_A)
											_SCREAMING_SNAKE_CASE		:						Union[str, Any]			  =    text_inputs.attention_mask.to(_A)

											_SCREAMING_SNAKE_CASE  ,				_SCREAMING_SNAKE_CASE		:						Dict			  =    self.text_encoder(
											    input_ids=_A      ,    attention_mask=_A)

											_SCREAMING_SNAKE_CASE		:						Optional[Any]			  =    prompt_embeds.repeat_interleave(_A      ,    dim=0)
											_SCREAMING_SNAKE_CASE		:						Tuple			  =    text_encoder_hidden_states.repeat_interleave(_A      ,    dim=0)
											_SCREAMING_SNAKE_CASE		:						Optional[Any]			  =    text_mask.repeat_interleave(_A      ,    dim=0)

											if do_classifier_free_guidance:
																		_SCREAMING_SNAKE_CASE		:						List[str]
																		if negative_prompt is None:
																									_SCREAMING_SNAKE_CASE		:						Dict			  =    [""""""] * 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):
																									_SCREAMING_SNAKE_CASE		:						Optional[int]			  =    [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:
																									_SCREAMING_SNAKE_CASE		:						int			  =    negative_prompt

																		_SCREAMING_SNAKE_CASE		:						Optional[Any]			  =    self.tokenizer(
																		    _A      ,    padding="""max_length"""      ,    max_length=7_7      ,    truncation=_A      ,    return_attention_mask=_A      ,    add_special_tokens=_A      ,    return_tensors="""pt"""      ,    )
																		_SCREAMING_SNAKE_CASE		:						List[Any]			  =    uncond_input.input_ids.to(_A)
																		_SCREAMING_SNAKE_CASE		:						List[str]			  =    uncond_input.attention_mask.to(_A)

																		_SCREAMING_SNAKE_CASE  ,				_SCREAMING_SNAKE_CASE		:						Optional[int]			  =    self.text_encoder(
																		    input_ids=_A      ,    attention_mask=_A)

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

																		_SCREAMING_SNAKE_CASE		:						Optional[int]			  =    negative_prompt_embeds.shape[1]
																		_SCREAMING_SNAKE_CASE		:						List[str]			  =    negative_prompt_embeds.repeat(1      ,    _A)
																		_SCREAMING_SNAKE_CASE		:						List[Any]			  =    negative_prompt_embeds.view(batch_size * num_images_per_prompt      ,    _A)

																		_SCREAMING_SNAKE_CASE		:						Any			  =    uncond_text_encoder_hidden_states.shape[1]
																		_SCREAMING_SNAKE_CASE		:						Optional[int]			  =    uncond_text_encoder_hidden_states.repeat(1      ,    _A      ,    1)
																		_SCREAMING_SNAKE_CASE		:						List[str]			  =    uncond_text_encoder_hidden_states.view(
																		    batch_size * num_images_per_prompt      ,    _A      ,    -1)
																		_SCREAMING_SNAKE_CASE		:						Dict			  =    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
																		_SCREAMING_SNAKE_CASE		:						str			  =    torch.cat([negative_prompt_embeds, prompt_embeds])
																		_SCREAMING_SNAKE_CASE		:						List[str]			  =    torch.cat([uncond_text_encoder_hidden_states, text_encoder_hidden_states])

																		_SCREAMING_SNAKE_CASE		:						str			  =    torch.cat([uncond_text_mask, text_mask])

											return prompt_embeds, text_encoder_hidden_states, text_mask





				def 		_lowerCAmelCase		(      self							:							Dict      ,    _A							:							Union[str, Any]=0):



											"""simple docstring"""







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

											_SCREAMING_SNAKE_CASE		:						Dict			  =    torch.device(f"""cuda:{gpu_id}""")

											_SCREAMING_SNAKE_CASE		:						Tuple			  =    [
											    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 		_lowerCAmelCase		(      self							:							List[Any]      ,    _A							:							Optional[int]=0):



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

											_SCREAMING_SNAKE_CASE		:						Optional[int]			  =    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)

											_SCREAMING_SNAKE_CASE		:						Any			  =    None
											for cpu_offloaded_model in [self.text_encoder, self.unet, self.movq]:
																		_SCREAMING_SNAKE_CASE  ,				_SCREAMING_SNAKE_CASE		:						List[Any]			  =    cpu_offload_with_hook(_A      ,    _A      ,    prev_module_hook=_A)

											if self.safety_checker is not None:
																		_SCREAMING_SNAKE_CASE  ,				_SCREAMING_SNAKE_CASE		:						Dict			  =    cpu_offload_with_hook(self.safety_checker      ,    _A      ,    prev_module_hook=_A)

											# We'll offload the last model manually.
											_SCREAMING_SNAKE_CASE		:						List[str]			  =    hook





				@property
				# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline._execution_device
				def 		_lowerCAmelCase		(      self							:							List[Any]):



											"""simple docstring"""







											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[int]      ,    _A							:							Union[str, List[str]]      ,    _A							:							Union[torch.FloatTensor, List[torch.FloatTensor]]      ,    _A							:							Union[torch.FloatTensor, List[torch.FloatTensor]]      ,    _A							:							Optional[Union[str, List[str]]] = None      ,    _A							:							int = 5_1_2      ,    _A							:							int = 5_1_2      ,    _A							:							int = 1_0_0      ,    _A							:							float = 4.0      ,    _A							:							int = 1      ,    _A							:							Optional[Union[torch.Generator, List[torch.Generator]]] = None      ,    _A							:							Optional[torch.FloatTensor] = None      ,    _A							:							Optional[str] = "pil"      ,    _A							:							bool = True      ,    ):



											"""simple docstring"""







											if isinstance(_A      ,    _A):
																		_SCREAMING_SNAKE_CASE		:						List[Any]			  =    1
											elif isinstance(_A      ,    _A):
																		_SCREAMING_SNAKE_CASE		:						Any			  =    len(_A)
											else:
																		raise ValueError(f"""`prompt` has to be of type `str` or `list` but is {type(_A)}""")

											_SCREAMING_SNAKE_CASE		:						Tuple			  =    self._execution_device

											_SCREAMING_SNAKE_CASE		:						Any			  =    batch_size * num_images_per_prompt
											_SCREAMING_SNAKE_CASE		:						Optional[int]			  =    guidance_scale > 1.0

											_SCREAMING_SNAKE_CASE  ,				_SCREAMING_SNAKE_CASE  ,				_SCREAMING_SNAKE_CASE		:						List[str]			  =    self._encode_prompt(
											    _A      ,    _A      ,    _A      ,    _A      ,    _A)

											if isinstance(_A      ,    _A):
																		_SCREAMING_SNAKE_CASE		:						str			  =    torch.cat(_A      ,    dim=0)
											if isinstance(_A      ,    _A):
																		_SCREAMING_SNAKE_CASE		:						Tuple			  =    torch.cat(_A      ,    dim=0)

											if do_classifier_free_guidance:
																		_SCREAMING_SNAKE_CASE		:						str			  =    image_embeds.repeat_interleave(_A      ,    dim=0)
																		_SCREAMING_SNAKE_CASE		:						Optional[int]			  =    negative_image_embeds.repeat_interleave(_A      ,    dim=0)

											_SCREAMING_SNAKE_CASE		:						List[str]			  =    torch.cat([negative_image_embeds, image_embeds]      ,    dim=0).to(
											    dtype=prompt_embeds.dtype      ,    device=_A)

											self.scheduler.set_timesteps(_A      ,    device=_A)
											_SCREAMING_SNAKE_CASE		:						Any			  =    self.scheduler.timesteps

											_SCREAMING_SNAKE_CASE		:						Optional[Any]			  =    self.unet.config.in_channels

											_SCREAMING_SNAKE_CASE  ,				_SCREAMING_SNAKE_CASE		:						Dict			  =    get_new_h_w(_A      ,    _A      ,    self.movq_scale_factor)

											# create initial latent
											_SCREAMING_SNAKE_CASE		:						List[Any]			  =    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
																		_SCREAMING_SNAKE_CASE		:						Optional[Any]			  =    torch.cat([latents] * 2) if do_classifier_free_guidance else latents

																		_SCREAMING_SNAKE_CASE		:						Union[str, Any]			  =    {"""text_embeds""": prompt_embeds, """image_embeds""": image_embeds}
																		_SCREAMING_SNAKE_CASE		:						List[Any]			  =    self.unet(
																		    sample=_A      ,    timestep=_A      ,    encoder_hidden_states=_A      ,    added_cond_kwargs=_A      ,    return_dict=_A      ,    )[0]

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

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

											if output_type == "pil":
																		_SCREAMING_SNAKE_CASE		:						List[str]			  =    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_flax_available,
    is_tf_available,
    is_torch_available,
)


lowerCAmelCase_     =	{'''configuration_unispeech''': ['''UNISPEECH_PRETRAINED_CONFIG_ARCHIVE_MAP''', '''UniSpeechConfig''']}

try:
			if not is_torch_available():
						raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
			pass
else:
			lowerCAmelCase_     =	[
			    '''UNISPEECH_PRETRAINED_MODEL_ARCHIVE_LIST''',
			    '''UniSpeechForCTC''',
			    '''UniSpeechForPreTraining''',
			    '''UniSpeechForSequenceClassification''',
			    '''UniSpeechModel''',
			    '''UniSpeechPreTrainedModel''',
			]

if TYPE_CHECKING:
			from .configuration_unispeech import UNISPEECH_PRETRAINED_CONFIG_ARCHIVE_MAP, UniSpeechConfig

			try:
						if not is_torch_available():
									raise OptionalDependencyNotAvailable()
			except OptionalDependencyNotAvailable:
						pass
			else:
						from .modeling_unispeech import (
						    UNISPEECH_PRETRAINED_MODEL_ARCHIVE_LIST,
						    UniSpeechForCTC,
						    UniSpeechForPreTraining,
						    UniSpeechForSequenceClassification,
						    UniSpeechModel,
						    UniSpeechPreTrainedModel,
						)

else:
			import sys

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


