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


_lowercase:	Optional[Any]       			=		logging.get_logger(__name__)  # pylint: disable=invalid-name

_lowercase:	Optional[int]       			=		'''
    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  (		snake_case     , snake_case     , snake_case=8  ):
							_lowerCAmelCase							=							h // scale_factor**2
							if h % scale_factor**2 != 0:
														new_h += 1
							_lowerCAmelCase							=							w // scale_factor**2
							if w % scale_factor**2 != 0:
														new_w += 1
							return new_h * scale_factor, new_w * scale_factor



class        lowerCamelCase__				(   UpperCAmelCase				):





			def __init__(       self	:	List[str]  ,  lowercase__	:	MultilingualCLIP  ,  lowercase__	:	XLMRobertaTokenizer  ,  lowercase__	:	UNetaDConditionModel  ,  lowercase__	:	Union[DDIMScheduler, DDPMScheduler]  ,  lowercase__	:	VQModel  ,  ):
										super().__init__()

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





			def   SCREAMING_SNAKE_CASE__     (       self	:	Any  ,  lowercase__	:	Any  ,  lowercase__	:	Optional[int]  ,  lowercase__	:	List[Any]  ,  lowercase__	:	Tuple  ,  lowercase__	:	Union[str, Any]  ,  lowercase__	:	Tuple						):
										if latents is None:
																	_lowerCAmelCase							=							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}'						)
																	_lowerCAmelCase							=							latents.to(lowercase__						)

										_lowerCAmelCase							=							latents * scheduler.init_noise_sigma
										return latents





			def   SCREAMING_SNAKE_CASE__     (       self	:	Optional[Any]  ,  lowercase__	:	Optional[Any]  ,  lowercase__	:	List[Any]  ,  lowercase__	:	str  ,  lowercase__	:	int  ,  lowercase__	:	int=None  ,  ):
										_lowerCAmelCase							=							len(lowercase__						) if isinstance(lowercase__  ,  lowercase__						) else 1
										# get prompt text embeddings
										_lowerCAmelCase							=							self.tokenizer(
										    lowercase__  ,  padding='max_length'  ,  truncation=lowercase__  ,  max_length=77  ,  return_attention_mask=lowercase__  ,  add_special_tokens=lowercase__  ,  return_tensors='pt'  ,  )

										_lowerCAmelCase							=							text_inputs.input_ids
										_lowerCAmelCase							=							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__						):
																	_lowerCAmelCase							=							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}'						)

										_lowerCAmelCase							=							text_input_ids.to(lowercase__						)
										_lowerCAmelCase							=							text_inputs.attention_mask.to(lowercase__						)

										_lowerCAmelCase			,		_lowerCAmelCase							=							self.text_encoder(
										    input_ids=lowercase__  ,  attention_mask=lowercase__						)

										_lowerCAmelCase							=							prompt_embeds.repeat_interleave(lowercase__  ,  dim=0						)
										_lowerCAmelCase							=							text_encoder_hidden_states.repeat_interleave(lowercase__  ,  dim=0						)
										_lowerCAmelCase							=							text_mask.repeat_interleave(lowercase__  ,  dim=0						)

										if do_classifier_free_guidance:
																	_lowerCAmelCase							=							42
																	if negative_prompt is None:
																								_lowerCAmelCase							=							[''] * 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__						):
																								_lowerCAmelCase							=							[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:
																								_lowerCAmelCase							=							negative_prompt

																	_lowerCAmelCase							=							self.tokenizer(
																	    lowercase__  ,  padding='max_length'  ,  max_length=77  ,  truncation=lowercase__  ,  return_attention_mask=lowercase__  ,  add_special_tokens=lowercase__  ,  return_tensors='pt'  ,  )
																	_lowerCAmelCase							=							uncond_input.input_ids.to(lowercase__						)
																	_lowerCAmelCase							=							uncond_input.attention_mask.to(lowercase__						)

																	_lowerCAmelCase			,		_lowerCAmelCase							=							self.text_encoder(
																	    input_ids=lowercase__  ,  attention_mask=lowercase__						)

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

																	_lowerCAmelCase							=							negative_prompt_embeds.shape[1]
																	_lowerCAmelCase							=							negative_prompt_embeds.repeat(1  ,  lowercase__						)
																	_lowerCAmelCase							=							negative_prompt_embeds.view(batch_size * num_images_per_prompt  ,  lowercase__						)

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

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

										return prompt_embeds, text_encoder_hidden_states, text_mask





			def   SCREAMING_SNAKE_CASE__     (       self	:	List[str]  ,  lowercase__	:	Tuple=0						):
										if is_accelerate_available():
																	from accelerate import cpu_offload
										else:
																	raise ImportError('Please install accelerate via `pip install accelerate`'						)

										_lowerCAmelCase							=							torch.device(f'cuda:{gpu_id}'						)

										_lowerCAmelCase							=							[
										    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   SCREAMING_SNAKE_CASE__     (       self	:	Dict  ,  lowercase__	:	Tuple=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.'						)

										_lowerCAmelCase							=							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)

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

										if self.safety_checker is not None:
																	_lowerCAmelCase			,		_lowerCAmelCase							=							cpu_offload_with_hook(self.safety_checker  ,  lowercase__  ,  prev_module_hook=lowercase__						)

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





			@property
			# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline._execution_device
			def   SCREAMING_SNAKE_CASE__     (       self	:	Optional[int]						):
										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	:	List[str]  ,  lowercase__	:	Union[str, List[str]]  ,  lowercase__	:	Union[torch.FloatTensor, List[torch.FloatTensor]]  ,  lowercase__	:	Union[torch.FloatTensor, List[torch.FloatTensor]]  ,  lowercase__	:	Optional[Union[str, List[str]]] = None  ,  lowercase__	:	int = 5_12  ,  lowercase__	:	int = 5_12  ,  lowercase__	:	int = 1_00  ,  lowercase__	:	float = 4.0  ,  lowercase__	:	int = 1  ,  lowercase__	:	Optional[Union[torch.Generator, List[torch.Generator]]] = None  ,  lowercase__	:	Optional[torch.FloatTensor] = None  ,  lowercase__	:	Optional[str] = "pil"  ,  lowercase__	:	bool = True  ,  ):

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

										_lowerCAmelCase							=							self._execution_device

										_lowerCAmelCase							=							batch_size * num_images_per_prompt
										_lowerCAmelCase							=							guidance_scale > 1.0

										_lowerCAmelCase			,		_lowerCAmelCase			,		_lowerCAmelCase							=							self._encode_prompt(
										    lowercase__  ,  lowercase__  ,  lowercase__  ,  lowercase__  ,  lowercase__						)

										if isinstance(lowercase__  ,  lowercase__						):
																	_lowerCAmelCase							=							torch.cat(lowercase__  ,  dim=0						)
										if isinstance(lowercase__  ,  lowercase__						):
																	_lowerCAmelCase							=							torch.cat(lowercase__  ,  dim=0						)

										if do_classifier_free_guidance:
																	_lowerCAmelCase							=							image_embeds.repeat_interleave(lowercase__  ,  dim=0						)
																	_lowerCAmelCase							=							negative_image_embeds.repeat_interleave(lowercase__  ,  dim=0						)

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

										self.scheduler.set_timesteps(lowercase__  ,  device=lowercase__						)
										_lowerCAmelCase							=							self.scheduler.timesteps

										_lowerCAmelCase							=							self.unet.config.in_channels

										_lowerCAmelCase			,		_lowerCAmelCase							=							get_new_h_w(lowercase__  ,  lowercase__  ,  self.movq_scale_factor						)

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

																	_lowerCAmelCase							=							{'text_embeds': prompt_embeds, 'image_embeds': image_embeds}
																	_lowerCAmelCase							=							self.unet(
																	    sample=lowercase__  ,  timestep=lowercase__  ,  encoder_hidden_states=lowercase__  ,  added_cond_kwargs=lowercase__  ,  return_dict=lowercase__  ,  )[0]

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

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

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

										if not return_dict:
																	return (image,)

										return ImagePipelineOutput(images=lowercase__						)







from ...configuration_utils import PretrainedConfig
from ...utils import logging


_lowercase:	List[Any]       			=		logging.get_logger(__name__)

_lowercase:	int       			=		{
    '''alibaba-damo/mgp-str-base''': '''https://huggingface.co/alibaba-damo/mgp-str-base/resolve/main/config.json''',
}



class        lowerCamelCase__				(   UpperCAmelCase				):
			UpperCamelCase__    							="mgp-str"





			def __init__(       self	:	int  ,  lowercase__	:	List[Any]=[32, 1_28]  ,  lowercase__	:	str=4  ,  lowercase__	:	List[str]=3  ,  lowercase__	:	Tuple=27  ,  lowercase__	:	Optional[Any]=38  ,  lowercase__	:	Optional[Any]=5_02_57  ,  lowercase__	:	Dict=3_05_22  ,  lowercase__	:	Dict=7_68  ,  lowercase__	:	Tuple=12  ,  lowercase__	:	int=12  ,  lowercase__	:	int=4.0  ,  lowercase__	:	Tuple=True  ,  lowercase__	:	List[Any]=False  ,  lowercase__	:	int=1e-5  ,  lowercase__	:	int=0.0  ,  lowercase__	:	Optional[Any]=0.0  ,  lowercase__	:	List[Any]=0.0  ,  lowercase__	:	Union[str, Any]=False  ,  lowercase__	:	str=0.0_2  ,  **lowercase__	:	Optional[int]  ,  ):
										super().__init__(**lowercase__						)

										_lowerCAmelCase							=							image_size
										_lowerCAmelCase							=							patch_size
										_lowerCAmelCase							=							num_channels
										_lowerCAmelCase							=							max_token_length
										_lowerCAmelCase							=							num_character_labels
										_lowerCAmelCase							=							num_bpe_labels
										_lowerCAmelCase							=							num_wordpiece_labels
										_lowerCAmelCase							=							hidden_size
										_lowerCAmelCase							=							num_hidden_layers
										_lowerCAmelCase							=							num_attention_heads
										_lowerCAmelCase							=							mlp_ratio
										_lowerCAmelCase							=							distilled
										_lowerCAmelCase							=							layer_norm_eps
										_lowerCAmelCase							=							drop_rate
										_lowerCAmelCase							=							qkv_bias
										_lowerCAmelCase							=							attn_drop_rate
										_lowerCAmelCase							=							drop_path_rate
										_lowerCAmelCase							=							output_aa_attentions
										_lowerCAmelCase							=							initializer_range






