

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_ :				Tuple						 =    logging.get_logger(__name__)  # pylint: disable=invalid-name

A_ :				str						 =    '\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      (lowercase_:				Optional[Any]					,   lowercase_:				Any					,   lowercase_:				Tuple=8      )  ->   Optional[Any]:
	A__							:					int       =	h // scale_factor**2
	if h % scale_factor**2 != 0:
		new_h += 1
	A__							:					List[Any]       =	w // scale_factor**2
	if w % scale_factor**2 != 0:
		new_w += 1
	return new_h * scale_factor, new_w * scale_factor



class 					_a(__a       ):


					'''simple docstring'''






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

						self.register_modules(
						    text_encoder=lowerCAmelCase_					, tokenizer=lowerCAmelCase_					, unet=lowerCAmelCase_					, scheduler=lowerCAmelCase_					, movq=lowerCAmelCase_					, )
						A__							:					int       =	2 ** (len(self.movq.config.block_out_channels    ) - 1)






					def 						__A    (				self					, A__					, A__					, A__					, A__					, A__					, A__    ):







						if latents is None:
							A__							:					Union[str, Any]       =	randn_tensor(lowerCAmelCase_					, generator=lowerCAmelCase_					, device=lowerCAmelCase_					, dtype=lowerCAmelCase_    )
						else:
							if latents.shape != shape:
								raise ValueError(F"""Unexpected latents shape, got {latents.shape}, expected {shape}"""    )
							A__							:					int       =	latents.to(lowerCAmelCase_    )

						A__							:					Union[str, Any]       =	latents * scheduler.init_noise_sigma
						return latents






					def 						__A    (				self					, A__					, A__					, A__					, A__					, A__=None					, ):
						A__							:					int       =	len(lowerCAmelCase_    ) if isinstance(lowerCAmelCase_					, lowerCAmelCase_    ) else 1
						# get prompt text embeddings
						A__							:					Optional[Any]       =	self.tokenizer(
						    lowerCAmelCase_					, padding="""max_length"""					, truncation=lowerCAmelCase_					, max_length=77					, return_attention_mask=lowerCAmelCase_					, add_special_tokens=lowerCAmelCase_					, return_tensors="""pt"""					, )

						A__							:					Optional[int]       =	text_inputs.input_ids
						A__							:					str       =	self.tokenizer(lowerCAmelCase_					, padding="""longest"""					, return_tensors="""pt"""    ).input_ids

						if untruncated_ids.shape[-1] >= text_input_ids.shape[-1] and not torch.equal(lowerCAmelCase_					, lowerCAmelCase_    ):
							A__							:					int       =	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}"""    )

						A__							:					List[str]       =	text_input_ids.to(lowerCAmelCase_    )
						A__							:					Dict       =	text_inputs.attention_mask.to(lowerCAmelCase_    )

						A__     ,     A__							:					int       =	self.text_encoder(
						    input_ids=lowerCAmelCase_					, attention_mask=lowerCAmelCase_    )

						A__							:					int       =	prompt_embeds.repeat_interleave(lowerCAmelCase_					, dim=0    )
						A__							:					int       =	text_encoder_hidden_states.repeat_interleave(lowerCAmelCase_					, dim=0    )
						A__							:					int       =	text_mask.repeat_interleave(lowerCAmelCase_					, dim=0    )

						if do_classifier_free_guidance:
							A__							:					int       =	42
							if negative_prompt is None:
								A__							:					Union[str, Any]       =	[""""""] * batch_size
							elif type(lowerCAmelCase_    ) is not type(lowerCAmelCase_    ):
								raise TypeError(
								    F"""`negative_prompt` should be the same type to `prompt`, but got {type(lowerCAmelCase_    )} !="""
								    F""" {type(lowerCAmelCase_    )}."""    )
							elif isinstance(lowerCAmelCase_					, lowerCAmelCase_    ):
								A__							:					Optional[Any]       =	[negative_prompt]
							elif batch_size != len(lowerCAmelCase_    ):
								raise ValueError(
								    F"""`negative_prompt`: {negative_prompt} has batch size {len(lowerCAmelCase_    )}, but `prompt`:"""
								    F""" {prompt} has batch size {batch_size}. Please make sure that passed `negative_prompt` matches"""
								    """ the batch size of `prompt`."""    )
							else:
								A__							:					List[Any]       =	negative_prompt

							A__							:					Optional[Any]       =	self.tokenizer(
							    lowerCAmelCase_					, padding="""max_length"""					, max_length=77					, truncation=lowerCAmelCase_					, return_attention_mask=lowerCAmelCase_					, add_special_tokens=lowerCAmelCase_					, return_tensors="""pt"""					, )
							A__							:					Optional[Any]       =	uncond_input.input_ids.to(lowerCAmelCase_    )
							A__							:					Any       =	uncond_input.attention_mask.to(lowerCAmelCase_    )

							A__     ,     A__							:					Dict       =	self.text_encoder(
							    input_ids=lowerCAmelCase_					, attention_mask=lowerCAmelCase_    )

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

							A__							:					Optional[Any]       =	negative_prompt_embeds.shape[1]
							A__							:					Any       =	negative_prompt_embeds.repeat(1					, lowerCAmelCase_    )
							A__							:					Tuple       =	negative_prompt_embeds.view(batch_size * num_images_per_prompt					, lowerCAmelCase_    )

							A__							:					List[str]       =	uncond_text_encoder_hidden_states.shape[1]
							A__							:					str       =	uncond_text_encoder_hidden_states.repeat(1					, lowerCAmelCase_					, 1    )
							A__							:					Tuple       =	uncond_text_encoder_hidden_states.view(
							    batch_size * num_images_per_prompt					, lowerCAmelCase_					, -1    )
							A__							:					Tuple       =	uncond_text_mask.repeat_interleave(lowerCAmelCase_					, 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
							A__							:					List[Any]       =	torch.cat([negative_prompt_embeds, prompt_embeds]    )
							A__							:					List[str]       =	torch.cat([uncond_text_encoder_hidden_states, text_encoder_hidden_states]    )

							A__							:					Tuple       =	torch.cat([uncond_text_mask, text_mask]    )

						return prompt_embeds, text_encoder_hidden_states, text_mask






					def 						__A    (				self					, A__=0    ):







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

						A__							:					Optional[int]       =	torch.device(F"""cuda:{gpu_id}"""    )

						A__							:					List[str]       =	[
						    self.unet,
						    self.text_encoder,
						    self.movq,
						]
						for cpu_offloaded_model in models:
							if cpu_offloaded_model is not None:
								cpu_offload(lowerCAmelCase_					, lowerCAmelCase_    )






					def 						__A    (				self					, A__=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."""    )

						A__							:					Union[str, Any]       =	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)

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

						if self.safety_checker is not None:
							A__     ,     A__							:					Dict       =	cpu_offload_with_hook(self.safety_checker					, lowerCAmelCase_					, prev_module_hook=lowerCAmelCase_    )

						# We'll offload the last model manually.
						A__							:					Any       =	hook






					@property
					# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline._execution_device
					def 						__A    (				self    ):







						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					, A__					, A__					, A__					, A__ = None					, A__ = 512					, A__ = 512					, A__ = 100					, A__ = 4.0					, A__ = 1					, A__ = None					, A__ = None					, A__ = "pil"					, A__ = True					, ):







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

						A__							:					Union[str, Any]       =	self._execution_device

						A__							:					Tuple       =	batch_size * num_images_per_prompt
						A__							:					int       =	guidance_scale > 1.0

						A__     ,     A__     ,     A__							:					Union[str, Any]       =	self._encode_prompt(
						    lowerCAmelCase_					, lowerCAmelCase_					, lowerCAmelCase_					, lowerCAmelCase_					, lowerCAmelCase_    )

						if isinstance(lowerCAmelCase_					, lowerCAmelCase_    ):
							A__							:					int       =	torch.cat(lowerCAmelCase_					, dim=0    )
						if isinstance(lowerCAmelCase_					, lowerCAmelCase_    ):
							A__							:					Dict       =	torch.cat(lowerCAmelCase_					, dim=0    )

						if do_classifier_free_guidance:
							A__							:					List[Any]       =	image_embeds.repeat_interleave(lowerCAmelCase_					, dim=0    )
							A__							:					Optional[int]       =	negative_image_embeds.repeat_interleave(lowerCAmelCase_					, dim=0    )

						A__							:					Tuple       =	torch.cat([negative_image_embeds, image_embeds]					, dim=0    ).to(
						    dtype=prompt_embeds.dtype					, device=lowerCAmelCase_    )

						self.scheduler.set_timesteps(lowerCAmelCase_					, device=lowerCAmelCase_    )
						A__							:					List[Any]       =	self.scheduler.timesteps

						A__							:					str       =	self.unet.config.in_channels

						A__     ,     A__							:					List[str]       =	get_new_h_w(lowerCAmelCase_					, lowerCAmelCase_					, self.movq_scale_factor    )

						# create initial latent
						A__							:					Dict       =	self.prepare_latents(
						    (batch_size, num_channels_latents, height, width)					, text_encoder_hidden_states.dtype					, lowerCAmelCase_					, lowerCAmelCase_					, lowerCAmelCase_					, self.scheduler					, )

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

							A__							:					Any       =	{"""text_embeds""": prompt_embeds, """image_embeds""": image_embeds}
							A__							:					Dict       =	self.unet(
							    sample=lowerCAmelCase_					, timestep=lowerCAmelCase_					, encoder_hidden_states=lowerCAmelCase_					, added_cond_kwargs=lowerCAmelCase_					, return_dict=lowerCAmelCase_					, )[0]

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

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

						if output_type == "pil":
							A__							:					Tuple       =	self.numpy_to_pil(lowerCAmelCase_    )

						if not return_dict:
							return (image,)

						return ImagePipelineOutput(images=lowerCAmelCase_    )






import itertools
import random
import unittest

import numpy as np

from transformers import ASTFeatureExtractor
from transformers.testing_utils import require_torch, require_torchaudio
from transformers.utils.import_utils import is_torch_available

from ...test_sequence_feature_extraction_common import SequenceFeatureExtractionTestMixin


A_ :				Dict						 =    random.Random()

if is_torch_available():
	import torch

def 		UpperCamelCase      (lowercase_:				Tuple					,   lowercase_:				Tuple=1.0					,   lowercase_:				Dict=None					,   lowercase_:				int=None      )  ->   str:
	if rng is None:
		A__							:					Optional[Any]       =	global_rng

	A__							:					List[str]       =	[]
	for batch_idx in range(shape[0]      ):
		values.append([]      )
		for _ in range(shape[1]      ):
			values[-1].append(rng.random() * scale      )

	return values



class 					_a		(unittest.TestCase       ):


					'''simple docstring'''






					def __init__(				self					, A__					, A__=7					, A__=400					, A__=2000					, A__=1					, A__=0.0					, A__=1_6000					, A__=True					, A__=True					, ):
						A__							:					Any       =	parent
						A__							:					Optional[int]       =	batch_size
						A__							:					Union[str, Any]       =	min_seq_length
						A__							:					Dict       =	max_seq_length
						A__							:					Tuple       =	(self.max_seq_length - self.min_seq_length) // (self.batch_size - 1)
						A__							:					str       =	feature_size
						A__							:					Optional[int]       =	padding_value
						A__							:					List[str]       =	sampling_rate
						A__							:					List[str]       =	return_attention_mask
						A__							:					int       =	do_normalize






					def 						__A    (				self    ):
						return {
						    "feature_size": self.feature_size,
						    "padding_value": self.padding_value,
						    "sampling_rate": self.sampling_rate,
						    "return_attention_mask": self.return_attention_mask,
						    "do_normalize": self.do_normalize,
						}






					def 						__A    (				self					, A__=False					, A__=False    ):
						def _flatten(A__    ):
							return list(itertools.chain(*A__    )    )

						if equal_length:
							A__							:					Dict       =	floats_list((self.batch_size, self.max_seq_length)    )
						else:
							# make sure that inputs increase in size
							A__							:					Union[str, Any]       =	[
							    _flatten(floats_list((x, self.feature_size)    )    )
							    for x in range(self.min_seq_length					, self.max_seq_length					, self.seq_length_diff    )
							]

						if numpify:
							A__							:					Optional[int]       =	[np.asarray(A__    ) for x in speech_inputs]

						return speech_inputs



@require_torch
@require_torchaudio
class 					_a		(__magic_name__	,     unittest.TestCase       ):


					'''simple docstring'''






					UpperCAmelCase__:     int													= ASTFeatureExtractor






					def 						__A    (				self    ):
						A__							:					Optional[Any]       =	ASTFeatureExtractionTester(self    )






					def 						__A    (				self    ):
						# Tests that all call wrap to encode_plus and batch_encode_plus
						A__							:					Optional[int]       =	self.feature_extraction_class(**self.feat_extract_tester.prepare_feat_extract_dict()    )
						# create three inputs of length 800, 1000, and 1200
						A__							:					List[Any]       =	[floats_list((1, x)    )[0] for x in range(800					, 1400					, 200    )]
						A__							:					Optional[Any]       =	[np.asarray(A__    ) for speech_input in speech_inputs]

						# Test not batched input
						A__							:					Tuple       =	feat_extract(speech_inputs[0]					, return_tensors="""np"""    ).input_values
						A__							:					Tuple       =	feat_extract(np_speech_inputs[0]					, return_tensors="""np"""    ).input_values
						self.assertTrue(np.allclose(A__					, A__					, atol=1e-3    )    )

						# Test batched
						A__							:					Tuple       =	feat_extract(A__					, padding=A__					, return_tensors="""np"""    ).input_values
						A__							:					Tuple       =	feat_extract(A__					, padding=A__					, return_tensors="""np"""    ).input_values
						for enc_seq_a, enc_seq_a in zip(A__					, A__    ):
							self.assertTrue(np.allclose(A__					, A__					, atol=1e-3    )    )

						# Test 2-D numpy arrays are batched.
						A__							:					int       =	[floats_list((1, x)    )[0] for x in (800, 800, 800)]
						A__							:					List[str]       =	np.asarray(A__    )
						A__							:					Union[str, Any]       =	feat_extract(A__					, return_tensors="""np"""    ).input_values
						A__							:					Optional[Any]       =	feat_extract(A__					, return_tensors="""np"""    ).input_values
						for enc_seq_a, enc_seq_a in zip(A__					, A__    ):
							self.assertTrue(np.allclose(A__					, A__					, atol=1e-3    )    )






					@require_torch
					def 						__A    (				self    ):
						import torch

						A__							:					Union[str, Any]       =	self.feature_extraction_class(**self.feat_extract_tester.prepare_feat_extract_dict()    )
						A__							:					Tuple       =	np.random.rand(100    ).astype(np.floataa    )
						A__							:					Tuple       =	np_speech_inputs.tolist()

						for inputs in [py_speech_inputs, np_speech_inputs]:
							A__							:					List[str]       =	feature_extractor.pad([{"""input_values""": inputs}]					, return_tensors="""np"""    )
							self.assertTrue(np_processed.input_values.dtype == np.floataa    )
							A__							:					Any       =	feature_extractor.pad([{"""input_values""": inputs}]					, return_tensors="""pt"""    )
							self.assertTrue(pt_processed.input_values.dtype == torch.floataa    )






					def 						__A    (				self					, A__    ):
						from datasets import load_dataset

						A__							:					str       =	load_dataset("""hf-internal-testing/librispeech_asr_dummy"""					, """clean"""					, split="""validation"""    )
						# automatic decoding with librispeech
						A__							:					str       =	ds.sort("""id"""    ).select(range(A__    )    )[:num_samples]["""audio"""]

						return [x["array"] for x in speech_samples]






					@require_torch
					def 						__A    (				self    ):
						# fmt: off
						A__							:					Optional[Any]       =	torch.tensor(
						    [-0.9_8_9_4, -1.2_7_7_6, -0.9_0_6_6, -1.2_7_7_6, -0.9_3_4_9, -1.2_6_0_9, -1.0_3_8_6, -1.2_7_7_6,
						     -1.1_5_6_1, -1.2_7_7_6, -1.2_0_5_2, -1.2_7_2_3, -1.2_1_9_0, -1.2_1_3_2, -1.2_7_7_6, -1.1_1_3_3,
						     -1.1_9_5_3, -1.1_3_4_3, -1.1_5_8_4, -1.2_2_0_3, -1.1_7_7_0, -1.2_4_7_4, -1.2_3_8_1, -1.1_9_3_6,
						     -0.9_2_7_0, -0.8_3_1_7, -0.8_0_4_9, -0.7_7_0_6, -0.7_5_6_5, -0.7_8_6_9]    )
						# fmt: on

						A__							:					Any       =	self._load_datasamples(1    )
						A__							:					Tuple       =	ASTFeatureExtractor()
						A__							:					Dict       =	feature_extractor(A__					, return_tensors="""pt"""    ).input_values
						self.assertEquals(input_values.shape					, (1, 1024, 128)    )
						self.assertTrue(torch.allclose(input_values[0, 0, :30]					, A__					, atol=1e-4    )    )



