"""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  :  List[Any]=				logging.get_logger(__name__)  # pylint: disable=invalid-name

_a  :  List[Any]=				"\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 						__UpperCAmelCase       (					UpperCAmelCase_					:       Dict    ,							UpperCAmelCase_					:       List[Any]    ,							UpperCAmelCase_					:       List[Any]=8						)							->    Tuple:

	'''simple docstring'''





	__snake_case    :							Tuple   	=				h // scale_factor**2
	if h % scale_factor**2 != 0:
		new_h += 1
	__snake_case    :							Tuple   	=				w // scale_factor**2
	if w % scale_factor**2 != 0:
		new_w += 1
	return new_h * scale_factor, new_w * scale_factor

class 	UpperCamelCase    (					lowercase		):



				def __init__(self       :   str			,						_A       :   MultilingualCLIP			,						_A       :   XLMRobertaTokenizer			,						_A       :   UNetaDConditionModel			,						_A       :   Union[DDIMScheduler, DDPMScheduler]			,						_A       :   VQModel			,						)					->			int:
					super().__init__()

					self.register_modules(
					    text_encoder=_A			,						tokenizer=_A			,						unet=_A			,						scheduler=_A			,						movq=_A			,						)
					__snake_case    :							List[Any]   	=				2 ** (len(self.movq.config.block_out_channels) - 1)



				def       _lowercase  (self       :   Any			,						_A       :   Optional[Any]			,						_A       :   Dict			,						_A       :   List[Any]			,						_A       :   Union[str, Any]			,						_A       :   Optional[int]			,						_A       :   Union[str, Any])					->			Tuple:
					if latents is None:
						__snake_case    :							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}")
						__snake_case    :							Tuple   	=				latents.to(_A)

					__snake_case    :							Optional[Any]   	=				latents * scheduler.init_noise_sigma
					return latents



				def       _lowercase  (self       :   str			,						_A       :   Tuple			,						_A       :   List[Any]			,						_A       :   Optional[int]			,						_A       :   int			,						_A       :   Optional[int]=None			,						)					->			Optional[int]:
					__snake_case    :							Dict   	=				len(_A) if isinstance(_A			,						_A) else 1
					# get prompt text embeddings
					__snake_case    :							Dict   	=				self.tokenizer(
					    _A			,						padding='max_length'			,						truncation=_A			,						max_length=77			,						return_attention_mask=_A			,						add_special_tokens=_A			,						return_tensors='pt'			,						)

					__snake_case    :							Tuple   	=				text_inputs.input_ids
					__snake_case    :							Union[str, 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):
						__snake_case    :							Tuple   	=				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}")

					__snake_case    :							List[str]   	=				text_input_ids.to(_A)
					__snake_case    :							Optional[Any]   	=				text_inputs.attention_mask.to(_A)

					__snake_case	,       __snake_case    :							List[Any]   	=				self.text_encoder(
					    input_ids=_A			,						attention_mask=_A)

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

					if do_classifier_free_guidance:
						__snake_case    :							List[str]
						if negative_prompt is None:
							__snake_case    :							List[Any]   	=				[''] * 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):
							__snake_case    :							str   	=				[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:
							__snake_case    :							int   	=				negative_prompt

						__snake_case    :							int   	=				self.tokenizer(
						    _A			,						padding='max_length'			,						max_length=77			,						truncation=_A			,						return_attention_mask=_A			,						add_special_tokens=_A			,						return_tensors='pt'			,						)
						__snake_case    :							Any   	=				uncond_input.input_ids.to(_A)
						__snake_case    :							Any   	=				uncond_input.attention_mask.to(_A)

						__snake_case	,       __snake_case    :							str   	=				self.text_encoder(
						    input_ids=_A			,						attention_mask=_A)

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

						__snake_case    :							Any   	=				negative_prompt_embeds.shape[1]
						__snake_case    :							Any   	=				negative_prompt_embeds.repeat(1			,						_A)
						__snake_case    :							List[Any]   	=				negative_prompt_embeds.view(batch_size * num_images_per_prompt			,						_A)

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

						__snake_case    :							Union[str, Any]   	=				torch.cat([uncond_text_mask, text_mask])

					return prompt_embeds, text_encoder_hidden_states, text_mask



				def       _lowercase  (self       :   Union[str, Any]			,						_A       :   List[Any]=0)					->			Optional[Any]:
					if is_accelerate_available():
						from accelerate import cpu_offload
					else:
						raise ImportError('Please install accelerate via `pip install accelerate`')

					__snake_case    :							Any   	=				torch.device(f"cuda:{gpu_id}")

					__snake_case    :							Optional[Any]   	=				[
					    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       :   List[str]			,						_A       :   Optional[Any]=0)					->			Union[str, Any]:
					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.')

					__snake_case    :							str   	=				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)

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

					if self.safety_checker is not None:
						__snake_case	,       __snake_case    :							Optional[Any]   	=				cpu_offload_with_hook(self.safety_checker			,						_A			,						prev_module_hook=_A)

					# We'll offload the last model manually.
					__snake_case    :							Optional[Any]   	=				hook



				@property
				# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline._execution_device
				def       _lowercase  (self       :   int)					->			Optional[int]:
					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       :   str			,						_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_12			,						_A       :   int = 5_12			,						_A       :   int = 1_00			,						_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			,						)					->			Any:

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

					__snake_case    :							List[Any]   	=				self._execution_device

					__snake_case    :							Union[str, Any]   	=				batch_size * num_images_per_prompt
					__snake_case    :							Dict   	=				guidance_scale > 1.0

					__snake_case	,       __snake_case	,       __snake_case    :							Any   	=				self._encode_prompt(
					    _A			,						_A			,						_A			,						_A			,						_A)

					if isinstance(_A			,						_A):
						__snake_case    :							Optional[int]   	=				torch.cat(_A			,						dim=0)
					if isinstance(_A			,						_A):
						__snake_case    :							Union[str, Any]   	=				torch.cat(_A			,						dim=0)

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

					__snake_case    :							Any   	=				torch.cat([negative_image_embeds, image_embeds]			,						dim=0).to(
					    dtype=prompt_embeds.dtype			,						device=_A)

					self.scheduler.set_timesteps(_A			,						device=_A)
					__snake_case    :							Optional[Any]   	=				self.scheduler.timesteps

					__snake_case    :							str   	=				self.unet.config.in_channels

					__snake_case	,       __snake_case    :							List[str]   	=				get_new_h_w(_A			,						_A			,						self.movq_scale_factor)

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

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

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

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

					if output_type == "pil":
						__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,
)


_a  :  Optional[Any]=				{
    "configuration_wav2vec2": ["WAV_2_VEC_2_PRETRAINED_CONFIG_ARCHIVE_MAP", "Wav2Vec2Config"],
    "feature_extraction_wav2vec2": ["Wav2Vec2FeatureExtractor"],
    "processing_wav2vec2": ["Wav2Vec2Processor"],
    "tokenization_wav2vec2": ["Wav2Vec2CTCTokenizer", "Wav2Vec2Tokenizer"],
}


try:
	if not is_torch_available():
		raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
	pass
else:
	_a  :  Optional[Any]=				[
	    "WAV_2_VEC_2_PRETRAINED_MODEL_ARCHIVE_LIST",
	    "Wav2Vec2ForAudioFrameClassification",
	    "Wav2Vec2ForCTC",
	    "Wav2Vec2ForMaskedLM",
	    "Wav2Vec2ForPreTraining",
	    "Wav2Vec2ForSequenceClassification",
	    "Wav2Vec2ForXVector",
	    "Wav2Vec2Model",
	    "Wav2Vec2PreTrainedModel",
	]

try:
	if not is_tf_available():
		raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
	pass
else:
	_a  :  List[str]=				[
	    "TF_WAV_2_VEC_2_PRETRAINED_MODEL_ARCHIVE_LIST",
	    "TFWav2Vec2ForCTC",
	    "TFWav2Vec2Model",
	    "TFWav2Vec2PreTrainedModel",
	    "TFWav2Vec2ForSequenceClassification",
	]

try:
	if not is_flax_available():
		raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
	pass
else:
	_a  :  List[Any]=				[
	    "FlaxWav2Vec2ForCTC",
	    "FlaxWav2Vec2ForPreTraining",
	    "FlaxWav2Vec2Model",
	    "FlaxWav2Vec2PreTrainedModel",
	]


if TYPE_CHECKING:
	from .configuration_wavaveca import WAV_2_VEC_2_PRETRAINED_CONFIG_ARCHIVE_MAP, WavaVecaConfig
	from .feature_extraction_wavaveca import WavaVecaFeatureExtractor
	from .processing_wavaveca import WavaVecaProcessor
	from .tokenization_wavaveca import WavaVecaCTCTokenizer, WavaVecaTokenizer

	try:
		if not is_torch_available():
			raise OptionalDependencyNotAvailable()
	except OptionalDependencyNotAvailable:
		pass
	else:
		from .modeling_wavaveca import (
		    WAV_2_VEC_2_PRETRAINED_MODEL_ARCHIVE_LIST,
		    WavaVecaForAudioFrameClassification,
		    WavaVecaForCTC,
		    WavaVecaForMaskedLM,
		    WavaVecaForPreTraining,
		    WavaVecaForSequenceClassification,
		    WavaVecaForXVector,
		    WavaVecaModel,
		    WavaVecaPreTrainedModel,
		)

	try:
		if not is_tf_available():
			raise OptionalDependencyNotAvailable()
	except OptionalDependencyNotAvailable:
		pass
	else:
		from .modeling_tf_wavaveca import (
		    TF_WAV_2_VEC_2_PRETRAINED_MODEL_ARCHIVE_LIST,
		    TFWavaVecaForCTC,
		    TFWavaVecaForSequenceClassification,
		    TFWavaVecaModel,
		    TFWavaVecaPreTrainedModel,
		)

	try:
		if not is_flax_available():
			raise OptionalDependencyNotAvailable()
	except OptionalDependencyNotAvailable:
		pass
	else:
		from .modeling_tf_wavaveca import (
		    FlaxWavaVecaForCTC,
		    FlaxWavaVecaForPreTraining,
		    FlaxWavaVecaModel,
		    FlaxWavaVecaPreTrainedModel,
		)


else:
	import sys

	_a  :  Tuple=				_LazyModule(__name__, globals()["__file__"], _import_structure, module_spec=__spec__)





