




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


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

__magic_name__				     =     """
    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					( __lowercase							,     __lowercase							,     __lowercase=8    ):





       """simple docstring"""



       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        SCREAMING_SNAKE_CASE__   (						UpperCAmelCase_      ):
   def __init__(       self			:				Tuple   ,	SCREAMING_SNAKE_CASE_			:				MultilingualCLIP   ,	SCREAMING_SNAKE_CASE_			:				XLMRobertaTokenizer   ,	SCREAMING_SNAKE_CASE_			:				UNetaDConditionModel   ,	SCREAMING_SNAKE_CASE_			:				Union[DDIMScheduler, DDPMScheduler]   ,	SCREAMING_SNAKE_CASE_			:				VQModel   ,	):
          super().__init__()

          self.register_modules(
              text_encoder=_snake_case   ,	tokenizer=_snake_case   ,	unet=_snake_case   ,	scheduler=_snake_case   ,	movq=_snake_case   ,	)
          lowerCamelCase__          =						2 ** (len(self.movq.config.block_out_channels							) - 1)
   def    __UpperCAmelCase       (       self			:				List[Any]   ,	SCREAMING_SNAKE_CASE_			:				List[str]   ,	SCREAMING_SNAKE_CASE_			:				int   ,	SCREAMING_SNAKE_CASE_			:				Tuple   ,	SCREAMING_SNAKE_CASE_			:				Any   ,	SCREAMING_SNAKE_CASE_			:				Tuple   ,	SCREAMING_SNAKE_CASE_			:				Optional[Any]							):


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

          lowerCamelCase__          =						latents * scheduler.init_noise_sigma
          return latents
   def    __UpperCAmelCase       (       self			:				Dict   ,	SCREAMING_SNAKE_CASE_			:				List[str]   ,	SCREAMING_SNAKE_CASE_			:				Any   ,	SCREAMING_SNAKE_CASE_			:				Union[str, Any]   ,	SCREAMING_SNAKE_CASE_			:				Any   ,	SCREAMING_SNAKE_CASE_			:				Tuple=None   ,	):
          lowerCamelCase__          =						len(_snake_case							) if isinstance(_snake_case   ,	_snake_case							) else 1
          # get prompt text embeddings
          lowerCamelCase__          =						self.tokenizer(
              _snake_case   ,	padding="""max_length"""   ,	truncation=_snake_case   ,	max_length=77   ,	return_attention_mask=_snake_case   ,	add_special_tokens=_snake_case   ,	return_tensors="""pt"""   ,	)

          lowerCamelCase__          =						text_inputs.input_ids
          lowerCamelCase__          =						self.tokenizer(_snake_case   ,	padding="""longest"""   ,	return_tensors="""pt"""							).input_ids

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

          lowerCamelCase__     ,	lowerCamelCase__          =						self.text_encoder(
              input_ids=_snake_case   ,	attention_mask=_snake_case							)

          lowerCamelCase__          =						prompt_embeds.repeat_interleave(_snake_case   ,	dim=0							)
          lowerCamelCase__          =						text_encoder_hidden_states.repeat_interleave(_snake_case   ,	dim=0							)
          lowerCamelCase__          =						text_mask.repeat_interleave(_snake_case   ,	dim=0							)

          if do_classifier_free_guidance:
                 lowerCamelCase__          =						42
                 if negative_prompt is None:
                        lowerCamelCase__          =						[""""""] * batch_size
                 elif type(_snake_case							) is not type(_snake_case							):
                        raise TypeError(
                            f"""`negative_prompt` should be the same type to `prompt`, but got {type(_snake_case							)} !="""
                            f""" {type(_snake_case							)}."""							)
                 elif isinstance(_snake_case   ,	_snake_case							):
                        lowerCamelCase__          =						[negative_prompt]
                 elif batch_size != len(_snake_case							):
                        raise ValueError(
                            f"""`negative_prompt`: {negative_prompt} has batch size {len(_snake_case							)}, 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(
                     _snake_case   ,	padding="""max_length"""   ,	max_length=77   ,	truncation=_snake_case   ,	return_attention_mask=_snake_case   ,	add_special_tokens=_snake_case   ,	return_tensors="""pt"""   ,	)
                 lowerCamelCase__          =						uncond_input.input_ids.to(_snake_case							)
                 lowerCamelCase__          =						uncond_input.attention_mask.to(_snake_case							)

                 lowerCamelCase__     ,	lowerCamelCase__          =						self.text_encoder(
                     input_ids=_snake_case   ,	attention_mask=_snake_case							)

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

                 lowerCamelCase__          =						negative_prompt_embeds.shape[1]
                 lowerCamelCase__          =						negative_prompt_embeds.repeat(1   ,	_snake_case							)
                 lowerCamelCase__          =						negative_prompt_embeds.view(batch_size * num_images_per_prompt   ,	_snake_case							)

                 lowerCamelCase__          =						uncond_text_encoder_hidden_states.shape[1]
                 lowerCamelCase__          =						uncond_text_encoder_hidden_states.repeat(1   ,	_snake_case   ,	1							)
                 lowerCamelCase__          =						uncond_text_encoder_hidden_states.view(
                     batch_size * num_images_per_prompt   ,	_snake_case   ,	-1							)
                 lowerCamelCase__          =						uncond_text_mask.repeat_interleave(_snake_case   ,	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    __UpperCAmelCase       (       self			:				Union[str, Any]   ,	SCREAMING_SNAKE_CASE_			:				int=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(_snake_case   ,	_snake_case							)
   def    __UpperCAmelCase       (       self			:				Any   ,	SCREAMING_SNAKE_CASE_			:				Optional[Any]=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=_snake_case							)
                 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(_snake_case   ,	_snake_case   ,	prev_module_hook=_snake_case							)

          if self.safety_checker is not None:
                 lowerCamelCase__     ,	lowerCamelCase__          =						cpu_offload_with_hook(self.safety_checker   ,	_snake_case   ,	prev_module_hook=_snake_case							)

          # We'll offload the last model manually.
          lowerCamelCase__          =						hook
   @property
   # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline._execution_device
   def    __UpperCAmelCase       (       self			:				Union[str, Any]							):


          if not hasattr(self.unet   ,	"""_hf_hook"""							):
                 return self.device
          for module in self.unet.modules():
                 if (
                     hasattr(_snake_case   ,	"""_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(_snake_case							)
   def __call__(       self			:				List[str]   ,	SCREAMING_SNAKE_CASE_			:				Union[str, List[str]]   ,	SCREAMING_SNAKE_CASE_			:				Union[torch.FloatTensor, List[torch.FloatTensor]]   ,	SCREAMING_SNAKE_CASE_			:				Union[torch.FloatTensor, List[torch.FloatTensor]]   ,	SCREAMING_SNAKE_CASE_			:				Optional[Union[str, List[str]]] = None   ,	SCREAMING_SNAKE_CASE_			:				int = 512   ,	SCREAMING_SNAKE_CASE_			:				int = 512   ,	SCREAMING_SNAKE_CASE_			:				int = 100   ,	SCREAMING_SNAKE_CASE_			:				float = 4.0   ,	SCREAMING_SNAKE_CASE_			:				int = 1   ,	SCREAMING_SNAKE_CASE_			:				Optional[Union[torch.Generator, List[torch.Generator]]] = None   ,	SCREAMING_SNAKE_CASE_			:				Optional[torch.FloatTensor] = None   ,	SCREAMING_SNAKE_CASE_			:				Optional[str] = "pil"   ,	SCREAMING_SNAKE_CASE_			:				bool = True   ,	):


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

          lowerCamelCase__          =						self._execution_device

          lowerCamelCase__          =						batch_size * num_images_per_prompt
          lowerCamelCase__          =						guidance_scale > 1.0

          lowerCamelCase__     ,	lowerCamelCase__     ,	lowerCamelCase__          =						self._encode_prompt(
              _snake_case   ,	_snake_case   ,	_snake_case   ,	_snake_case   ,	_snake_case							)

          if isinstance(_snake_case   ,	_snake_case							):
                 lowerCamelCase__          =						torch.cat(_snake_case   ,	dim=0							)
          if isinstance(_snake_case   ,	_snake_case							):
                 lowerCamelCase__          =						torch.cat(_snake_case   ,	dim=0							)

          if do_classifier_free_guidance:
                 lowerCamelCase__          =						image_embeds.repeat_interleave(_snake_case   ,	dim=0							)
                 lowerCamelCase__          =						negative_image_embeds.repeat_interleave(_snake_case   ,	dim=0							)

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

          self.scheduler.set_timesteps(_snake_case   ,	device=_snake_case							)
          lowerCamelCase__          =						self.scheduler.timesteps

          lowerCamelCase__          =						self.unet.config.in_channels

          lowerCamelCase__     ,	lowerCamelCase__          =						get_new_h_w(_snake_case   ,	_snake_case   ,	self.movq_scale_factor							)

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

          for i, t in enumerate(self.progress_bar(_snake_case							)							):
                 # 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=_snake_case   ,	timestep=_snake_case   ,	encoder_hidden_states=_snake_case   ,	added_cond_kwargs=_snake_case   ,	return_dict=_snake_case   ,	)[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(
                     _snake_case   ,	_snake_case   ,	_snake_case   ,	generator=_snake_case   ,	).prev_sample
          # post-processing
          lowerCamelCase__          =						self.movq.decode(_snake_case   ,	force_not_quantize=_snake_case							)["""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(_snake_case							)

          if not return_dict:
                 return (image,)

          return ImagePipelineOutput(images=_snake_case							)











import argparse
from collections import defaultdict

import yaml


SCREAMING_SNAKE_CASE__   			=   '''docs/source/en/_toctree.yml'''


def A    (							__UpperCamelCase       )    ->   Optional[Any]:
			A__	  =     defaultdict(__UpperCamelCase       )
			for doc in model_doc:
						counts[doc["local"]] += 1
			A__	  =     [key for key, value in counts.items() if value > 1]

			A__	  =     []
			for duplicate_key in duplicates:
						A__	  =     list({doc['title'] for doc in model_doc if doc['local'] == duplicate_key}       )
						if len(__UpperCamelCase       ) > 1:
									raise ValueError(
									    f'''{duplicate_key} is present several times in the documentation table of content at '''
									    '`docs/source/en/_toctree.yml` with different *Title* values. Choose one of those and remove the '
									    'others.'       )
						# Only add this once
						new_doc.append({'local': duplicate_key, 'title': titles[0]}       )

			# Add none duplicate-keys
			new_doc.extend([doc for doc in model_doc if counts[doc['local']] == 1]       )

			# Sort
			return sorted(__UpperCamelCase			,						key=lambda __UpperCamelCase       : s["title"].lower()       )


def A    (							__UpperCamelCase=False       )    ->   str:
			with open(__UpperCamelCase			,						encoding='utf-8'       ) as f:
						A__	  =     yaml.safe_load(f.read()       )

			# Get to the API doc
			A__	  =     0
			while content[api_idx]["title"] != "API":
						api_idx += 1
			A__	  =     content[api_idx]['sections']

			# Then to the model doc
			A__	  =     0
			while api_doc[model_idx]["title"] != "Models":
						model_idx += 1

			A__	  =     api_doc[model_idx]['sections']

			A__	  =     [(idx, section) for idx, section in enumerate(__UpperCamelCase       ) if 'sections' in section]
			A__	  =     False
			for idx, modality_doc in modalities_docs:
						A__	  =     modality_doc['sections']
						A__	  =     clean_model_doc_toc(__UpperCamelCase       )

						if old_modality_doc != new_modality_doc:
									A__	  =     True
									if overwrite:
												A__	  =     new_modality_doc

			if diff:
						if overwrite:
									A__	  =     model_doc
									A__	  =     api_doc
									with open(__UpperCamelCase			,						'w'			,						encoding='utf-8'       ) as f:
												f.write(yaml.dump(__UpperCamelCase			,						allow_unicode=__UpperCamelCase       )       )
						else:
									raise ValueError(
									    'The model doc part of the table of content is not properly sorted, run `make style` to fix this.'       )


if __name__ == "__main__":
							SCREAMING_SNAKE_CASE__   			=   argparse.ArgumentParser()
							parser.add_argument('''--fix_and_overwrite''', action='''store_true''', help='''Whether to fix inconsistencies.''')
							SCREAMING_SNAKE_CASE__   			=   parser.parse_args()

							check_model_doc(args.fix_and_overwrite)




