
from typing import List, Optional, Union

import torch

from ...models import UNetaDConditionModel, VQModel
from ...pipelines import DiffusionPipeline
from ...pipelines.pipeline_utils import ImagePipelineOutput
from ...schedulers import DDPMScheduler
from ...utils import (
    is_accelerate_available,
    is_accelerate_version,
    logging,
    randn_tensor,
    replace_example_docstring,
)


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

__lowerCAmelCase     :	List[Any]       			=  '\n    Examples:\n        ```py\n        >>> from diffusers import KandinskyV22Pipeline, KandinskyV22PriorPipeline\n        >>> import torch\n\n        >>> pipe_prior = KandinskyV22PriorPipeline.from_pretrained("kandinsky-community/kandinsky-2-2-prior")\n        >>> pipe_prior.to("cuda")\n        >>> prompt = "red cat, 4k photo"\n        >>> out = pipe_prior(prompt)\n        >>> image_emb = out.image_embeds\n        >>> zero_image_emb = out.negative_image_embeds\n        >>> pipe = KandinskyV22Pipeline.from_pretrained("kandinsky-community/kandinsky-2-2-decoder")\n        >>> pipe.to("cuda")\n        >>> image = pipe(\n        ...     image_embeds=image_emb,\n        ...     negative_image_embeds=zero_image_emb,\n        ...     height=768,\n        ...     width=768,\n        ...     num_inference_steps=50,\n        ... ).images\n        >>> image[0].save("cat.png")\n        ```\n'

def a__       (			A_,	A_,	A_=8   ):






  '''simple docstring'''





  __magic_name__								=			height // scale_factor**2
  if height % scale_factor**2 != 0:
    new_height += 1
  __magic_name__								=			width // scale_factor**2
  if width % scale_factor**2 != 0:
    new_width += 1
  return new_height * scale_factor, new_width * scale_factor





class 			UpperCAmelCase_     (					_A	):


  '''simple docstring'''







  def __init__(					self     :					str     ,				UpperCamelCase__     :					UNetaDConditionModel     ,				UpperCamelCase__     :					DDPMScheduler     ,				UpperCamelCase__     :					VQModel     ,				)     ->      Optional[Any]:







    """simple docstring"""






    super().__init__()

    self.register_modules(
        unet=UpperCamelCase__     ,				scheduler=UpperCamelCase__     ,				movq=UpperCamelCase__     ,				)
    __magic_name__								=			2 ** (len(self.movq.config.block_out_channels       ) - 1)
  def 				_lowercase				(					self     :					List[str]     ,				UpperCamelCase__     :					List[str]     ,				UpperCamelCase__     :					Optional[Any]     ,				UpperCamelCase__     :					List[str]     ,				UpperCamelCase__     :					str     ,				UpperCamelCase__     :					Union[str, Any]     ,				UpperCamelCase__     :					Dict       )     ->      Any:







    """simple docstring"""






    if latents is None:
      __magic_name__								=			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}'''       )
      __magic_name__								=			latents.to(UpperCamelCase__       )

    __magic_name__								=			latents * scheduler.init_noise_sigma
    return latents
  def 				_lowercase				(					self     :					List[str]     ,				UpperCamelCase__     :					Any=0       )     ->      Tuple:







    """simple docstring"""






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

    __magic_name__								=			torch.device(F'''cuda:{gpu_id}'''       )

    __magic_name__								=			[
        self.unet,
        self.movq,
    ]
    for cpu_offloaded_model in models:
      if cpu_offloaded_model is not None:
        cpu_offload(UpperCamelCase__     ,				UpperCamelCase__       )
  def 				_lowercase				(					self     :					List[str]     ,				UpperCamelCase__     :					str=0       )     ->      Tuple:







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

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

    __magic_name__								=			None
    for cpu_offloaded_model in [self.unet, self.movq]:
      __magic_name__	,       __magic_name__								=			cpu_offload_with_hook(UpperCamelCase__     ,				UpperCamelCase__     ,				prev_module_hook=UpperCamelCase__       )

    # We'll offload the last model manually.
    __magic_name__								=			hook
  @property
  # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline._execution_device
  def 				_lowercase				(					self     :					List[Any]       )     ->      Union[str, 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     :					Tuple     ,				UpperCamelCase__     :					Union[torch.FloatTensor, List[torch.FloatTensor]]     ,				UpperCamelCase__     :					Union[torch.FloatTensor, List[torch.FloatTensor]]     ,				UpperCamelCase__     :					int = 512     ,				UpperCamelCase__     :					int = 512     ,				UpperCamelCase__     :					int = 100     ,				UpperCamelCase__     :					float = 4.0     ,				UpperCamelCase__     :					int = 1     ,				UpperCamelCase__     :					Optional[Union[torch.Generator, List[torch.Generator]]] = None     ,				UpperCamelCase__     :					Optional[torch.FloatTensor] = None     ,				UpperCamelCase__     :					Optional[str] = "pil"     ,				UpperCamelCase__     :					bool = True     ,				)     ->      Optional[Any]:







    """simple docstring"""






    __magic_name__								=			self._execution_device

    __magic_name__								=			guidance_scale > 1.0

    if isinstance(UpperCamelCase__     ,				UpperCamelCase__       ):
      __magic_name__								=			torch.cat(UpperCamelCase__     ,				dim=0       )
    __magic_name__								=			image_embeds.shape[0] * num_images_per_prompt
    if isinstance(UpperCamelCase__     ,				UpperCamelCase__       ):
      __magic_name__								=			torch.cat(UpperCamelCase__     ,				dim=0       )

    if do_classifier_free_guidance:
      __magic_name__								=			image_embeds.repeat_interleave(UpperCamelCase__     ,				dim=0       )
      __magic_name__								=			negative_image_embeds.repeat_interleave(UpperCamelCase__     ,				dim=0       )

    __magic_name__								=			torch.cat([negative_image_embeds, image_embeds]     ,				dim=0       ).to(dtype=self.unet.dtype     ,				device=UpperCamelCase__       )

    self.scheduler.set_timesteps(UpperCamelCase__     ,				device=UpperCamelCase__       )
    __magic_name__								=			self.scheduler.timesteps

    __magic_name__								=			self.unet.config.in_channels

    __magic_name__	,       __magic_name__								=			downscale_height_and_width(UpperCamelCase__     ,				UpperCamelCase__     ,				self.movq_scale_factor       )

    # create initial latent
    __magic_name__								=			self.prepare_latents(
        (batch_size, num_channels_latents, height, width)     ,				image_embeds.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
      __magic_name__								=			torch.cat([latents] * 2       ) if do_classifier_free_guidance else latents

      __magic_name__								=			{"""image_embeds""": image_embeds}
      __magic_name__								=			self.unet(
          sample=UpperCamelCase__     ,				timestep=UpperCamelCase__     ,				encoder_hidden_states=UpperCamelCase__     ,				added_cond_kwargs=UpperCamelCase__     ,				return_dict=UpperCamelCase__     ,				)[0]

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

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

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

    if not return_dict:
      return (image,)

    return ImagePipelineOutput(images=UpperCamelCase__       )



import argparse
import collections
import os
import re

from transformers.utils import direct_transformers_import


# All paths are set with the intent you should run this script from the root of the repo with the command
# python utils/check_table.py
__lowerCAmelCase     :	Dict       			=  'src/transformers'
__lowerCAmelCase     :	List[Any]       			=  'docs/source/en'
__lowerCAmelCase     :	str       			=  '.'

def a__       (			A_,	A_,	A_   ):






  '''simple docstring'''





  with open(A_,	"""r""",	encoding="""utf-8""",	newline="""\n"""   ) as f:
    __magic_name__								=			f.readlines()
  # Find the start prompt.
  __magic_name__								=			0
  while not lines[start_index].startswith(A_   ):
    start_index += 1
  start_index += 1

  __magic_name__								=			start_index
  while not lines[end_index].startswith(A_   ):
    end_index += 1
  end_index -= 1

  while len(lines[start_index]   ) <= 1:
    start_index += 1
  while len(lines[end_index]   ) <= 1:
    end_index -= 1
  end_index += 1
  return "".join(lines[start_index:end_index]   ), start_index, end_index, lines


# Add here suffixes that are used to identify models, separated by |
__lowerCAmelCase     :	str       			=  'Model|Encoder|Decoder|ForConditionalGeneration'
# Regexes that match TF/Flax/PT model names.
__lowerCAmelCase     :	List[str]       			=  re.compile(R'TF(.*)(?:Model|Encoder|Decoder|ForConditionalGeneration)')
__lowerCAmelCase     :	Any       			=  re.compile(R'Flax(.*)(?:Model|Encoder|Decoder|ForConditionalGeneration)')
# Will match any TF or Flax model too so need to be in an else branch afterthe two previous regexes.
__lowerCAmelCase     :	Optional[int]       			=  re.compile(R'(.*)(?:Model|Encoder|Decoder|ForConditionalGeneration)')


# This is to make sure the transformers module imported is the one in the repo.
__lowerCAmelCase     :	Union[str, Any]       			=  direct_transformers_import(TRANSFORMERS_PATH)

def a__       (			A_   ):






  '''simple docstring'''





  __magic_name__								=			re.finditer(""".+?(?:(?<=[a-z])(?=[A-Z])|(?<=[A-Z])(?=[A-Z][a-z])|$)""",	A_   )
  return [m.group(0   ) for m in matches]

def a__       (			A_,	A_   ):






  '''simple docstring'''





  __magic_name__								=			2 if text == """✅""" or text == """❌""" else len(A_   )
  __magic_name__								=			(width - text_length) // 2
  __magic_name__								=			width - text_length - left_indent
  return " " * left_indent + text + " " * right_indent

def a__       (			):






  '''simple docstring'''





  __magic_name__								=			transformers_module.models.auto.configuration_auto.CONFIG_MAPPING_NAMES
  __magic_name__								=			{
      name: config_maping_names[code]
      for code, name in transformers_module.MODEL_NAMES_MAPPING.items()
      if code in config_maping_names
  }
  __magic_name__								=			{name: config.replace("""Config""",	""""""   ) for name, config in model_name_to_config.items()}

  # Dictionaries flagging if each model prefix has a slow/fast tokenizer, backend in PT/TF/Flax.
  __magic_name__								=			collections.defaultdict(A_   )
  __magic_name__								=			collections.defaultdict(A_   )
  __magic_name__								=			collections.defaultdict(A_   )
  __magic_name__								=			collections.defaultdict(A_   )
  __magic_name__								=			collections.defaultdict(A_   )

  # Let's lookup through all transformers object (once).
  for attr_name in dir(A_   ):
    __magic_name__								=			None
    if attr_name.endswith("""Tokenizer"""   ):
      __magic_name__								=			slow_tokenizers
      __magic_name__								=			attr_name[:-9]
    elif attr_name.endswith("""TokenizerFast"""   ):
      __magic_name__								=			fast_tokenizers
      __magic_name__								=			attr_name[:-13]
    elif _re_tf_models.match(A_   ) is not None:
      __magic_name__								=			tf_models
      __magic_name__								=			_re_tf_models.match(A_   ).groups()[0]
    elif _re_flax_models.match(A_   ) is not None:
      __magic_name__								=			flax_models
      __magic_name__								=			_re_flax_models.match(A_   ).groups()[0]
    elif _re_pt_models.match(A_   ) is not None:
      __magic_name__								=			pt_models
      __magic_name__								=			_re_pt_models.match(A_   ).groups()[0]

    if lookup_dict is not None:
      while len(A_   ) > 0:
        if attr_name in model_name_to_prefix.values():
          __magic_name__								=			True
          break
        # Try again after removing the last word in the name
        __magic_name__								=			"""""".join(camel_case_split(A_   )[:-1]   )

    # Let's build that table!
  __magic_name__								=			list(model_name_to_config.keys()   )
  model_names.sort(key=str.lower   )
  __magic_name__								=			["""Model""", """Tokenizer slow""", """Tokenizer fast""", """PyTorch support""", """TensorFlow support""", """Flax Support"""]
  # We'll need widths to properly display everything in the center (+2 is to leave one extra space on each side).
  __magic_name__								=			[len(A_   ) + 2 for c in columns]
  __magic_name__								=			max([len(A_   ) for name in model_names]   ) + 2

  # Build the table per se
  __magic_name__								=			"""|""" + """|""".join([_center_text(A_,	A_   ) for c, w in zip(A_,	A_   )]   ) + """|\n"""
  # Use ":-----:" format to center-aligned table cell texts
  table += "|" + "|".join([""":""" + """-""" * (w - 2) + """:""" for w in widths]   ) + "|\n"

  __magic_name__								=			{True: """✅""", False: """❌"""}
  for name in model_names:
    __magic_name__								=			model_name_to_prefix[name]
    __magic_name__								=			[
        name,
        check[slow_tokenizers[prefix]],
        check[fast_tokenizers[prefix]],
        check[pt_models[prefix]],
        check[tf_models[prefix]],
        check[flax_models[prefix]],
    ]
    table += "|" + "|".join([_center_text(A_,	A_   ) for l, w in zip(A_,	A_   )]   ) + "|\n"
  return table

def a__       (			A_=False   ):






  '''simple docstring'''





  __magic_name__	,       __magic_name__	,       __magic_name__	,       __magic_name__								=			_find_text_in_file(
      filename=os.path.join(A_,	"""index.md"""   ),	start_prompt="""<!--This table is updated automatically from the auto modules""",	end_prompt="""<!-- End table-->""",	)
  __magic_name__								=			get_model_table_from_auto_modules()

  if current_table != new_table:
    if overwrite:
      with open(os.path.join(A_,	"""index.md"""   ),	"""w""",	encoding="""utf-8""",	newline="""\n"""   ) as f:
        f.writelines(lines[:start_index] + [new_table] + lines[end_index:]   )
    else:
      raise ValueError(
          """The model table in the `index.md` has not been updated. Run `make fix-copies` to fix this."""   )


if __name__ == "__main__":
    __lowerCAmelCase     :	Tuple       			=  argparse.ArgumentParser()
    parser.add_argument('--fix_and_overwrite', action='store_true', help='Whether to fix inconsistencies.')
    __lowerCAmelCase     :	Optional[Any]       			=  parser.parse_args()

    check_model_table(args.fix_and_overwrite)

