


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


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

__SCREAMING_SNAKE_CASE							:			Optional[Any]						=			"""
    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   _a (					_SCREAMING_SNAKE_CASE ,       _SCREAMING_SNAKE_CASE ,       _SCREAMING_SNAKE_CASE=8 )			-> List[str]:
							snake_case_							=						h // scale_factor**2
							if h % scale_factor**2 != 0:
														new_h += 1
							snake_case_							=						w // scale_factor**2
							if w % scale_factor**2 != 0:
														new_w += 1
							return new_h * scale_factor, new_w * scale_factor





class    __A   (snake_case__):





							'''simple docstring'''

							def __init__( self							:					List[str]						,				UpperCAmelCase_							:					Union[str, Any]						,				UpperCAmelCase_							:					Optional[int]						,				UpperCAmelCase_							:					Optional[int]						,				UpperCAmelCase_							:					str						,				UpperCAmelCase_							:					Union[str, Any]						,				)    ->List[str]:



														"""simple docstring"""



														super().__init__()

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



							def 						lowerCAmelCase				( self							:					Optional[Any]						,				UpperCAmelCase_							:					Dict						,				UpperCAmelCase_							:					str						,				UpperCAmelCase_							:					Tuple						,				UpperCAmelCase_							:					List[str]						,				UpperCAmelCase_							:					int						,				UpperCAmelCase_							:					Optional[Any]				)    ->List[Any]:



														"""simple docstring"""



														if latents is None:
																					snake_case_							=						randn_tensor(UpperCAmelCase_						,				generator=UpperCAmelCase_						,				device=UpperCAmelCase_						,				dtype=UpperCAmelCase_				)
														else:
																					if latents.shape != shape:
																												raise ValueError(F"""Unexpected latents shape, got {latents.shape}, expected {shape}"""				)
																					snake_case_							=						latents.to(UpperCAmelCase_				)

														snake_case_							=						latents * scheduler.init_noise_sigma
														return latents



							def 						lowerCAmelCase				( self							:					Optional[Any]						,				UpperCAmelCase_							:					Optional[int]						,				UpperCAmelCase_							:					List[str]						,				UpperCAmelCase_							:					Tuple						,				UpperCAmelCase_							:					str						,				UpperCAmelCase_							:					Optional[Any]=None						,				)    ->Union[str, Any]:



														"""simple docstring"""



														snake_case_							=						len(UpperCAmelCase_				) if isinstance(UpperCAmelCase_						,				UpperCAmelCase_				) else 1
														# get prompt text embeddings
														snake_case_							=						self.tokenizer(
														    UpperCAmelCase_						,				padding="""max_length"""						,				truncation=UpperCAmelCase_						,				max_length=77						,				return_attention_mask=UpperCAmelCase_						,				add_special_tokens=UpperCAmelCase_						,				return_tensors="""pt"""						,				)

														snake_case_							=						text_inputs.input_ids
														snake_case_							=						self.tokenizer(UpperCAmelCase_						,				padding="""longest"""						,				return_tensors="""pt"""				).input_ids

														if untruncated_ids.shape[-1] >= text_input_ids.shape[-1] and not torch.equal(UpperCAmelCase_						,				UpperCAmelCase_				):
																					snake_case_							=						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_							=						text_input_ids.to(UpperCAmelCase_				)
														snake_case_							=						text_inputs.attention_mask.to(UpperCAmelCase_				)

														snake_case_							=						self.text_encoder(
														    input_ids=UpperCAmelCase_						,				attention_mask=UpperCAmelCase_				)

														snake_case_							=						prompt_embeds.repeat_interleave(UpperCAmelCase_						,				dim=0				)
														snake_case_							=						text_encoder_hidden_states.repeat_interleave(UpperCAmelCase_						,				dim=0				)
														snake_case_							=						text_mask.repeat_interleave(UpperCAmelCase_						,				dim=0				)

														if do_classifier_free_guidance:
																					snake_case_							=						42
																					if negative_prompt is None:
																												snake_case_							=						[""""""] * batch_size
																					elif type(UpperCAmelCase_				) is not type(UpperCAmelCase_				):
																												raise TypeError(
																												    F"""`negative_prompt` should be the same type to `prompt`, but got {type(UpperCAmelCase_				)} !="""
																												    F""" {type(UpperCAmelCase_				)}."""				)
																					elif isinstance(UpperCAmelCase_						,				UpperCAmelCase_				):
																												snake_case_							=						[negative_prompt]
																					elif batch_size != len(UpperCAmelCase_				):
																												raise ValueError(
																												    F"""`negative_prompt`: {negative_prompt} has batch size {len(UpperCAmelCase_				)}, 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_							=						negative_prompt

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

																					snake_case_							=						self.text_encoder(
																					    input_ids=UpperCAmelCase_						,				attention_mask=UpperCAmelCase_				)

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

																					snake_case_							=						negative_prompt_embeds.shape[1]
																					snake_case_							=						negative_prompt_embeds.repeat(1						,				UpperCAmelCase_				)
																					snake_case_							=						negative_prompt_embeds.view(batch_size * num_images_per_prompt						,				UpperCAmelCase_				)

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

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

														return prompt_embeds, text_encoder_hidden_states, text_mask



							def 						lowerCAmelCase				( self							:					Dict						,				UpperCAmelCase_							:					List[Any]=0				)    ->Tuple:



														"""simple docstring"""



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

														snake_case_							=						torch.device(F"""cuda:{gpu_id}"""				)

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



							def 						lowerCAmelCase				( self							:					int						,				UpperCAmelCase_							:					Dict=0				)    ->Union[str, Any]:



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

														snake_case_							=						torch.device(F"""cuda:{gpu_id}"""				)

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

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

														if self.safety_checker is not None:
																					snake_case_							=						cpu_offload_with_hook(self.safety_checker						,				UpperCAmelCase_						,				prev_module_hook=UpperCAmelCase_				)

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



							@property
							# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline._execution_device
							def 						lowerCAmelCase				( self							:					Union[str, Any]				)    ->Optional[Any]:



														"""simple docstring"""



														if not hasattr(self.unet						,				"""_hf_hook"""				):
																					return self.device
														for module in self.unet.modules():
																					if (
																					    hasattr(UpperCAmelCase_						,				"""_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(UpperCAmelCase_				)
							def __call__( self							:					List[Any]						,				UpperCAmelCase_							:					List[Any]						,				UpperCAmelCase_							:					str						,				UpperCAmelCase_							:					str						,				UpperCAmelCase_							:					Optional[int] = None						,				UpperCAmelCase_							:					str = 512						,				UpperCAmelCase_							:					Tuple = 512						,				UpperCAmelCase_							:					Tuple = 100						,				UpperCAmelCase_							:					List[str] = 4.0						,				UpperCAmelCase_							:					int = 1						,				UpperCAmelCase_							:					Optional[Any] = None						,				UpperCAmelCase_							:					Optional[int] = None						,				UpperCAmelCase_							:					int = "pil"						,				UpperCAmelCase_							:					str = True						,				)    ->Dict:



														"""simple docstring"""



														if isinstance(UpperCAmelCase_						,				UpperCAmelCase_				):
																					snake_case_							=						1
														elif isinstance(UpperCAmelCase_						,				UpperCAmelCase_				):
																					snake_case_							=						len(UpperCAmelCase_				)
														else:
																					raise ValueError(F"""`prompt` has to be of type `str` or `list` but is {type(UpperCAmelCase_				)}"""				)

														snake_case_							=						self._execution_device

														snake_case_							=						batch_size * num_images_per_prompt
														snake_case_							=						guidance_scale > 1.0

														snake_case_							=						self._encode_prompt(
														    UpperCAmelCase_						,				UpperCAmelCase_						,				UpperCAmelCase_						,				UpperCAmelCase_						,				UpperCAmelCase_				)

														if isinstance(UpperCAmelCase_						,				UpperCAmelCase_				):
																					snake_case_							=						torch.cat(UpperCAmelCase_						,				dim=0				)
														if isinstance(UpperCAmelCase_						,				UpperCAmelCase_				):
																					snake_case_							=						torch.cat(UpperCAmelCase_						,				dim=0				)

														if do_classifier_free_guidance:
																					snake_case_							=						image_embeds.repeat_interleave(UpperCAmelCase_						,				dim=0				)
																					snake_case_							=						negative_image_embeds.repeat_interleave(UpperCAmelCase_						,				dim=0				)

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

														self.scheduler.set_timesteps(UpperCAmelCase_						,				device=UpperCAmelCase_				)
														snake_case_							=						self.scheduler.timesteps

														snake_case_							=						self.unet.config.in_channels

														snake_case_							=						get_new_h_w(UpperCAmelCase_						,				UpperCAmelCase_						,				self.movq_scale_factor				)

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

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

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

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

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

														if output_type == "pil":
																					snake_case_							=						self.numpy_to_pil(UpperCAmelCase_				)

														if not return_dict:
																					return (image,)

														return ImagePipelineOutput(images=UpperCAmelCase_				)








'''simple docstring'''




from operator import delitem, getitem, setitem

import pytest

from data_structures.hashing.hash_map import HashMap



def       __lowerCamelCase  (					_lowercase		)      ->			Optional[Any]:
							return getitem, k



def       __lowerCamelCase  (					_lowercase	,				_lowercase		)      ->			List[str]:
							return setitem, k, v



def       __lowerCamelCase  (					_lowercase		)      ->			int:
							return delitem, k



def       __lowerCamelCase  (					_lowercase	,				_lowercase	,				*_lowercase		)      ->			Optional[Any]:
							try:
														return fun(_lowercase	,				*_lowercase		), None
							except Exception as e:
														return None, e


a   :			List[str]					  =    (
    _set("""key_a""", """val_a"""),
    _set("""key_b""", """val_b"""),
)

a   :			List[Any]					  =    [
    _set("""key_a""", """val_a"""),
    _set("""key_a""", """val_b"""),
]

a   :			int					  =    [
    _set("""key_a""", """val_a"""),
    _set("""key_b""", """val_b"""),
    _del("""key_a"""),
    _del("""key_b"""),
    _set("""key_a""", """val_a"""),
    _del("""key_a"""),
]

a   :			List[Any]					  =    [
    _get("""key_a"""),
    _del("""key_a"""),
    _set("""key_a""", """val_a"""),
    _del("""key_a"""),
    _del("""key_a"""),
    _get("""key_a"""),
]

a   :			Tuple					  =    [
    *[_set(x, x) for x in range(5)],  # guaranteed upsize
]

a   :			Optional[Any]					  =    [
    *[_set(x, x) for x in range(5)],  # guaranteed upsize
    *[_del(x) for x in range(5)],
    _set("""key_a""", """val_b"""),
]



@pytest.mark.parametrize(
    """operations"""	,				(
        pytest.param(_add_items	,				id="""add items"""		),
        pytest.param(_overwrite_items	,				id="""overwrite items"""		),
        pytest.param(_delete_items	,				id="""delete items"""		),
        pytest.param(_access_absent_items	,				id="""access absent items"""		),
        pytest.param(_add_with_resize_up	,				id="""add with resize up"""		),
        pytest.param(_add_with_resize_down	,				id="""add with resize down"""		),
    )	,				)
def       __lowerCamelCase  (					_lowercase		)      ->			Optional[int]:
							UpperCAmelCase     :		List[str]							     =		HashMap(initial_block_size=4		)
							UpperCAmelCase     :		Dict							     =		{}
							for _, (fun, *args) in enumerate(_lowercase		):
														UpperCAmelCase     ,     UpperCAmelCase     :		Union[str, Any]							     =		_run_operation(_lowercase	,				_lowercase	,				*_lowercase		)
														UpperCAmelCase     ,     UpperCAmelCase     :		Any							     =		_run_operation(_lowercase	,				_lowercase	,				*_lowercase		)
														assert my_res == py_res
														assert str(_lowercase		) == str(_lowercase		)
														assert set(_lowercase		) == set(_lowercase		)
														assert len(_lowercase		) == len(_lowercase		)
														assert set(my.items()		) == set(py.items()		)



def       __lowerCamelCase  (					)      ->			List[Any]:
							def is_public(_lowercase		) -> bool:
														return not name.startswith("""_"""		)

							UpperCAmelCase     :		int							     =		{name for name in dir({}		) if is_public(_lowercase		)}
							UpperCAmelCase     :		Any							     =		{name for name in dir(HashMap()		) if is_public(_lowercase		)}

							assert dict_public_names > hash_public_names

