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


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

__snake_case							   =  '''\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     a   (    __a				,      __a				,      __a=8     )						->			str:



					'''simple docstring'''






					UpperCamelCase__     :Optional[Any]							      =				h // scale_factor**2
					if h % scale_factor**2 != 0:
										new_h += 1
					UpperCamelCase__     :Optional[Any]							      =				w // scale_factor**2
					if w % scale_factor**2 != 0:
										new_w += 1
					return new_h * scale_factor, new_w * scale_factor







class      lowercase			(							_UpperCAmelCase      ):

		"""simple docstring"""

		def __init__(	self					,     UpperCamelCase_					,     UpperCamelCase_					,     UpperCamelCase_					,     UpperCamelCase_					,     UpperCamelCase_					,     ):

							'''simple docstring'''


							super().__init__()

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




		def 					lowerCAmelCase__       (	self					,     UpperCamelCase_					,     UpperCamelCase_					,     UpperCamelCase_					,     UpperCamelCase_					,     UpperCamelCase_					,     UpperCamelCase_							):

							'''simple docstring'''


							if latents is None:
												UpperCamelCase__     :Tuple							      =				randn_tensor(lowercase__					,     generator=lowercase__					,     device=lowercase__					,     dtype=lowercase__							)
							else:
												if latents.shape != shape:
																	raise ValueError(F'''Unexpected latents shape, got {latents.shape}, expected {shape}'''							)
												UpperCamelCase__     :List[str]							      =				latents.to(lowercase__							)

							UpperCamelCase__     :int							      =				latents * scheduler.init_noise_sigma
							return latents




		def 					lowerCAmelCase__       (	self					,     UpperCamelCase_					,     UpperCamelCase_					,     UpperCamelCase_					,     UpperCamelCase_					,     UpperCamelCase_=None					,     ):

							'''simple docstring'''


							UpperCamelCase__     :int							      =				len(lowercase__							) if isinstance(lowercase__					,     lowercase__							) else 1
							# get prompt text embeddings
							UpperCamelCase__     :Any							      =				self.tokenizer(
							    lowercase__					,     padding='''max_length'''					,     truncation=lowercase__					,     max_length=77					,     return_attention_mask=lowercase__					,     add_special_tokens=lowercase__					,     return_tensors='''pt'''					,     )

							UpperCamelCase__     :Optional[Any]							      =				text_inputs.input_ids
							UpperCamelCase__     :int							      =				self.tokenizer(lowercase__					,     padding='''longest'''					,     return_tensors='''pt'''							).input_ids

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

							UpperCamelCase__     :Union[str, Any]							      =				text_input_ids.to(lowercase__							)
							UpperCamelCase__     :Tuple							      =				text_inputs.attention_mask.to(lowercase__							)

							UpperCamelCase__     :Optional[int]							      =				self.text_encoder(
							    input_ids=lowercase__					,     attention_mask=lowercase__							)

							UpperCamelCase__     :Union[str, Any]							      =				prompt_embeds.repeat_interleave(lowercase__					,     dim=0							)
							UpperCamelCase__     :Tuple							      =				text_encoder_hidden_states.repeat_interleave(lowercase__					,     dim=0							)
							UpperCamelCase__     :Optional[int]							      =				text_mask.repeat_interleave(lowercase__					,     dim=0							)

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

												UpperCamelCase__     :str							      =				self.tokenizer(
												    lowercase__					,     padding='''max_length'''					,     max_length=77					,     truncation=lowercase__					,     return_attention_mask=lowercase__					,     add_special_tokens=lowercase__					,     return_tensors='''pt'''					,     )
												UpperCamelCase__     :List[Any]							      =				uncond_input.input_ids.to(lowercase__							)
												UpperCamelCase__     :Optional[int]							      =				uncond_input.attention_mask.to(lowercase__							)

												UpperCamelCase__     :Union[str, Any]							      =				self.text_encoder(
												    input_ids=lowercase__					,     attention_mask=lowercase__							)

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

												UpperCamelCase__     :List[Any]							      =				negative_prompt_embeds.shape[1]
												UpperCamelCase__     :Dict							      =				negative_prompt_embeds.repeat(1					,     lowercase__							)
												UpperCamelCase__     :Tuple							      =				negative_prompt_embeds.view(batch_size * num_images_per_prompt					,     lowercase__							)

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

												UpperCamelCase__     :Dict							      =				torch.cat([uncond_text_mask, text_mask]							)

							return prompt_embeds, text_encoder_hidden_states, text_mask




		def 					lowerCAmelCase__       (	self					,     UpperCamelCase_=0							):

							'''simple docstring'''


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

							UpperCamelCase__     :Optional[Any]							      =				torch.device(F'''cuda:{gpu_id}'''							)

							UpperCamelCase__     :Optional[int]							      =				[
							    self.unet,
							    self.text_encoder,
							    self.movq,
							]
							for cpu_offloaded_model in models:
												if cpu_offloaded_model is not None:
																	cpu_offload(lowercase__					,     lowercase__							)




		def 					lowerCAmelCase__       (	self					,     UpperCamelCase_=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.'''							)

							UpperCamelCase__     :List[str]							      =				torch.device(F'''cuda:{gpu_id}'''							)

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

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

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

							# We'll offload the last model manually.
							UpperCamelCase__     :int							      =				hook




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

							'''simple docstring'''


							if not hasattr(self.unet					,     '''_hf_hook'''							):
												return self.device
							for module in self.unet.modules():
												if (
												    hasattr(lowercase__					,     '''_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(lowercase__							)
		def __call__(	self					,     UpperCamelCase_					,     UpperCamelCase_					,     UpperCamelCase_					,     UpperCamelCase_ = None					,     UpperCamelCase_ = 512					,     UpperCamelCase_ = 512					,     UpperCamelCase_ = 100					,     UpperCamelCase_ = 4.0					,     UpperCamelCase_ = 1					,     UpperCamelCase_ = None					,     UpperCamelCase_ = None					,     UpperCamelCase_ = "pil"					,     UpperCamelCase_ = True					,     ):

							'''simple docstring'''


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

							UpperCamelCase__     :Optional[Any]							      =				self._execution_device

							UpperCamelCase__     :Dict							      =				batch_size * num_images_per_prompt
							UpperCamelCase__     :Optional[int]							      =				guidance_scale > 1.0

							UpperCamelCase__     :str							      =				self._encode_prompt(
							    lowercase__					,     lowercase__					,     lowercase__					,     lowercase__					,     lowercase__							)

							if isinstance(lowercase__					,     lowercase__							):
												UpperCamelCase__     :int							      =				torch.cat(lowercase__					,     dim=0							)
							if isinstance(lowercase__					,     lowercase__							):
												UpperCamelCase__     :List[Any]							      =				torch.cat(lowercase__					,     dim=0							)

							if do_classifier_free_guidance:
												UpperCamelCase__     :Union[str, Any]							      =				image_embeds.repeat_interleave(lowercase__					,     dim=0							)
												UpperCamelCase__     :List[str]							      =				negative_image_embeds.repeat_interleave(lowercase__					,     dim=0							)

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

							self.scheduler.set_timesteps(lowercase__					,     device=lowercase__							)
							UpperCamelCase__     :Tuple							      =				self.scheduler.timesteps

							UpperCamelCase__     :Tuple							      =				self.unet.config.in_channels

							UpperCamelCase__     :Dict							      =				get_new_h_w(lowercase__					,     lowercase__					,     self.movq_scale_factor							)

							# create initial latent
							UpperCamelCase__     :Tuple							      =				self.prepare_latents(
							    (batch_size, num_channels_latents, height, width)					,     text_encoder_hidden_states.dtype					,     lowercase__					,     lowercase__					,     lowercase__					,     self.scheduler					,     )

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

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

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

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

							if output_type == "pil":
												UpperCamelCase__     :Dict							      =				self.numpy_to_pil(lowercase__							)

							if not return_dict:
												return (image,)

							return ImagePipelineOutput(images=lowercase__							)



'''simple docstring'''
from .integrations import (
    is_optuna_available,
    is_ray_available,
    is_sigopt_available,
    is_wandb_available,
    run_hp_search_optuna,
    run_hp_search_ray,
    run_hp_search_sigopt,
    run_hp_search_wandb,
)
from .trainer_utils import (
    HPSearchBackend,
    default_hp_space_optuna,
    default_hp_space_ray,
    default_hp_space_sigopt,
    default_hp_space_wandb,
)
from .utils import logging


__snake_case							   =  logging.get_logger(__name__)
class      lowercase			:

		"""simple docstring"""

		_a									=       42
		_a									=       None




		@staticmethod
		def 					lowerCAmelCase__       (	):

							'''simple docstring'''


							raise NotImplementedError




		def 					lowerCAmelCase__       (	self					,     UpperCamelCase_					,     UpperCamelCase_					,     UpperCamelCase_					,     **UpperCamelCase_							):

							'''simple docstring'''


							raise NotImplementedError




		def 					lowerCAmelCase__       (	self					,     UpperCamelCase_							):

							'''simple docstring'''


							raise NotImplementedError




		def 					lowerCAmelCase__       (	self							):

							'''simple docstring'''


							if not self.is_available():
												raise RuntimeError(
												    F'''You picked the {self.name} backend, but it is not installed. Run {self.pip_install()}.'''							)




		@classmethod
		def 					lowerCAmelCase__       (	cls							):

							'''simple docstring'''


							return F'''`pip install {cls.pip_package or cls.name}`'''



class      lowercase			(							A__      ):

		"""simple docstring"""

		_a									=       'optuna'




		@staticmethod
		def 					lowerCAmelCase__       (	):

							'''simple docstring'''


							return is_optuna_available()




		def 					lowerCAmelCase__       (	self					,     UpperCamelCase_					,     UpperCamelCase_					,     UpperCamelCase_					,     **UpperCamelCase_							):

							'''simple docstring'''


							return run_hp_search_optuna(UpperCamelCase_					,     UpperCamelCase_					,     UpperCamelCase_					,     **UpperCamelCase_							)




		def 					lowerCAmelCase__       (	self					,     UpperCamelCase_							):

							'''simple docstring'''


							return default_hp_space_optuna(UpperCamelCase_							)



class      lowercase			(							A__      ):

		"""simple docstring"""

		_a									=       'ray'
		_a									=       '\'ray[tune]\''




		@staticmethod
		def 					lowerCAmelCase__       (	):

							'''simple docstring'''


							return is_ray_available()




		def 					lowerCAmelCase__       (	self					,     UpperCamelCase_					,     UpperCamelCase_					,     UpperCamelCase_					,     **UpperCamelCase_							):

							'''simple docstring'''


							return run_hp_search_ray(UpperCamelCase_					,     UpperCamelCase_					,     UpperCamelCase_					,     **UpperCamelCase_							)




		def 					lowerCAmelCase__       (	self					,     UpperCamelCase_							):

							'''simple docstring'''


							return default_hp_space_ray(UpperCamelCase_							)



class      lowercase			(							A__      ):

		"""simple docstring"""

		_a									=       'sigopt'




		@staticmethod
		def 					lowerCAmelCase__       (	):

							'''simple docstring'''


							return is_sigopt_available()




		def 					lowerCAmelCase__       (	self					,     UpperCamelCase_					,     UpperCamelCase_					,     UpperCamelCase_					,     **UpperCamelCase_							):

							'''simple docstring'''


							return run_hp_search_sigopt(UpperCamelCase_					,     UpperCamelCase_					,     UpperCamelCase_					,     **UpperCamelCase_							)




		def 					lowerCAmelCase__       (	self					,     UpperCamelCase_							):

							'''simple docstring'''


							return default_hp_space_sigopt(UpperCamelCase_							)



class      lowercase			(							A__      ):

		"""simple docstring"""

		_a									=       'wandb'




		@staticmethod
		def 					lowerCAmelCase__       (	):

							'''simple docstring'''


							return is_wandb_available()




		def 					lowerCAmelCase__       (	self					,     UpperCamelCase_					,     UpperCamelCase_					,     UpperCamelCase_					,     **UpperCamelCase_							):

							'''simple docstring'''


							return run_hp_search_wandb(UpperCamelCase_					,     UpperCamelCase_					,     UpperCamelCase_					,     **UpperCamelCase_							)




		def 					lowerCAmelCase__       (	self					,     UpperCamelCase_							):

							'''simple docstring'''


							return default_hp_space_wandb(UpperCamelCase_							)



__snake_case							   =  {
    HPSearchBackend(backend.name): backend for backend in [OptunaBackend, RayTuneBackend, SigOptBackend, WandbBackend]
}


def     a   (    )						->			str:



					'''simple docstring'''






					UpperCamelCase__     :List[Any]							      =				[backend for backend in ALL_HYPERPARAMETER_SEARCH_BACKENDS.values() if backend.is_available()]
					if len(__a     ) > 0:
										UpperCamelCase__     :Tuple							      =				available_backends[0].name
										if len(__a     ) > 1:
															logger.info(
															    f'''{len(__a     )} hyperparameter search backends available. Using {name} as the default.'''     )
										return name
					raise RuntimeError(
					    '''No hyperparameter search backend available.\n'''
					    + '''\n'''.join(
					        f''' - To install {backend.name} run {backend.pip_install()}'''
					        for backend in ALL_HYPERPARAMETER_SEARCH_BACKENDS.values()     )     )