



'''simple docstring'''





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,
)


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

_UpperCamelCase											=		'\n    Examples:\n        ```py\n        >>> import torch\n        >>> import numpy as np\n\n        >>> from diffusers import KandinskyV22PriorPipeline, KandinskyV22ControlnetPipeline\n        >>> from transformers import pipeline\n        >>> from diffusers.utils import load_image\n\n\n        >>> def make_hint(image, depth_estimator):\n        ...     image = depth_estimator(image)["depth"]\n        ...     image = np.array(image)\n        ...     image = image[:, :, None]\n        ...     image = np.concatenate([image, image, image], axis=2)\n        ...     detected_map = torch.from_numpy(image).float() / 255.0\n        ...     hint = detected_map.permute(2, 0, 1)\n        ...     return hint\n\n\n        >>> depth_estimator = pipeline("depth-estimation")\n\n        >>> pipe_prior = KandinskyV22PriorPipeline.from_pretrained(\n        ...     "kandinsky-community/kandinsky-2-2-prior", torch_dtype=torch.float16\n        ... )\n        >>> pipe_prior = pipe_prior.to("cuda")\n\n        >>> pipe = KandinskyV22ControlnetPipeline.from_pretrained(\n        ...     "kandinsky-community/kandinsky-2-2-controlnet-depth", torch_dtype=torch.float16\n        ... )\n        >>> pipe = pipe.to("cuda")\n\n\n        >>> img = load_image(\n        ...     "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main"\n        ...     "/kandinsky/cat.png"\n        ... ).resize((768, 768))\n\n        >>> hint = make_hint(img, depth_estimator).unsqueeze(0).half().to("cuda")\n\n        >>> prompt = "A robot, 4k photo"\n        >>> negative_prior_prompt = "lowres, text, error, cropped, worst quality, low quality, jpeg artifacts, ugly, duplicate, morbid, mutilated, out of frame, extra fingers, mutated hands, poorly drawn hands, poorly drawn face, mutation, deformed, blurry, dehydrated, bad anatomy, bad proportions, extra limbs, cloned face, disfigured, gross proportions, malformed limbs, missing arms, missing legs, extra arms, extra legs, fused fingers, too many fingers, long neck, username, watermark, signature"\n\n        >>> generator = torch.Generator(device="cuda").manual_seed(43)\n\n        >>> image_emb, zero_image_emb = pipe_prior(\n        ...     prompt=prompt, negative_prompt=negative_prior_prompt, generator=generator\n        ... ).to_tuple()\n\n        >>> images = pipe(\n        ...     image_embeds=image_emb,\n        ...     negative_image_embeds=zero_image_emb,\n        ...     hint=hint,\n        ...     num_inference_steps=50,\n        ...     generator=generator,\n        ...     height=768,\n        ...     width=768,\n        ... ).images\n\n        >>> images[0].save("robot_cat.png")\n        ```\n'

def 			a_	(				_lowerCAmelCase			,_lowerCAmelCase			,_lowerCAmelCase=8			) -> Any:
  __lowerCamelCase			:			Optional[int]       			=					height // scale_factor**2
  if height % scale_factor**2 != 0:
    new_height += 1
  __lowerCamelCase			:			List[str]       			=					width // scale_factor**2
  if width % scale_factor**2 != 0:
    new_width += 1
  return new_height * scale_factor, new_width * scale_factor






class    lowerCamelCase_       (							SCREAMING_SNAKE_CASE_	):
  """simple docstring"""




  def __init__(					self :      Tuple     ,      _a :      UNetaDConditionModel     ,      _a :      DDPMScheduler     ,      _a :      VQModel     ,      )   ->   Dict:
    super().__init__()

    self.register_modules(
        unet=_a     ,      scheduler=_a     ,      movq=_a     ,      )
    __lowerCamelCase			:			Any       			=					2 ** (len(self.movq.config.block_out_channels  ) - 1)




  def 		_lowercase	(					self :      int     ,      _a :      Dict     ,      _a :      str     ,      _a :      Optional[int]     ,      _a :      Any     ,      _a :      Optional[Any]     ,      _a :      Optional[Any]  )   ->   int:
    if latents is None:
      __lowerCamelCase			:			int       			=					randn_tensor(_a     ,      generator=_a     ,      device=_a     ,      dtype=_a  )
    else:
      if latents.shape != shape:
        raise ValueError(f'Unexpected latents shape, got {latents.shape}, expected {shape}'  )
      __lowerCamelCase			:			Union[str, Any]       			=					latents.to(_a  )

    __lowerCamelCase			:			Union[str, Any]       			=					latents * scheduler.init_noise_sigma
    return latents




  def 		_lowercase	(					self :      Any     ,      _a :      List[Any]=0  )   ->   int:
    if is_accelerate_available():
      from accelerate import cpu_offload
    else:
      raise ImportError('Please install accelerate via `pip install accelerate`'  )

    __lowerCamelCase			:			List[Any]       			=					torch.device(f'cuda:{gpu_id}'  )

    __lowerCamelCase			:			List[str]       			=					[
        self.unet,
        self.movq,
    ]
    for cpu_offloaded_model in models:
      if cpu_offloaded_model is not None:
        cpu_offload(_a     ,      _a  )




  def 		_lowercase	(					self :      Union[str, Any]     ,      _a :      Dict=0  )   ->   str:
    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			:			Any       			=					torch.device(f'cuda:{gpu_id}'  )

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

    __lowerCamelCase			:			Optional[int]       			=					None
    for cpu_offloaded_model in [self.unet, self.movq]:
      __lowerCamelCase			,__lowerCamelCase			:			Optional[Any]       			=					cpu_offload_with_hook(_a     ,      _a     ,      prev_module_hook=_a  )

    # We'll offload the last model manually.
    __lowerCamelCase			:			Optional[int]       			=					hook




  @property
  # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline._execution_device
  def 		_lowercase	(					self :      Optional[Any]  )   ->   int:
    if not hasattr(self.unet     ,      '_hf_hook'  ):
      return self.device
    for module in self.unet.modules():
      if (
          hasattr(_a     ,      '_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(_a  )
  def __call__(					self :      List[str]     ,      _a :      Union[torch.FloatTensor, List[torch.FloatTensor]]     ,      _a :      Union[torch.FloatTensor, List[torch.FloatTensor]]     ,      _a :      torch.FloatTensor     ,      _a :      int = 512     ,      _a :      int = 512     ,      _a :      int = 100     ,      _a :      float = 4.0     ,      _a :      int = 1     ,      _a :      Optional[Union[torch.Generator, List[torch.Generator]]] = None     ,      _a :      Optional[torch.FloatTensor] = None     ,      _a :      Optional[str] = "pil"     ,      _a :      bool = True     ,      )   ->   str:
    __lowerCamelCase			:			str       			=					self._execution_device

    __lowerCamelCase			:			Optional[int]       			=					guidance_scale > 1.0

    if isinstance(_a     ,      _a  ):
      __lowerCamelCase			:			Tuple       			=					torch.cat(_a     ,      dim=0  )
    if isinstance(_a     ,      _a  ):
      __lowerCamelCase			:			Tuple       			=					torch.cat(_a     ,      dim=0  )
    if isinstance(_a     ,      _a  ):
      __lowerCamelCase			:			Any       			=					torch.cat(_a     ,      dim=0  )

    __lowerCamelCase			:			Optional[int]       			=					image_embeds.shape[0] * num_images_per_prompt

    if do_classifier_free_guidance:
      __lowerCamelCase			:			Union[str, Any]       			=					image_embeds.repeat_interleave(_a     ,      dim=0  )
      __lowerCamelCase			:			Dict       			=					negative_image_embeds.repeat_interleave(_a     ,      dim=0  )
      __lowerCamelCase			:			Tuple       			=					hint.repeat_interleave(_a     ,      dim=0  )

    __lowerCamelCase			:			Dict       			=					torch.cat([negative_image_embeds, image_embeds]     ,      dim=0  ).to(dtype=self.unet.dtype     ,      device=_a  )
    __lowerCamelCase			:			int       			=					torch.cat([hint, hint]     ,      dim=0  ).to(dtype=self.unet.dtype     ,      device=_a  )
    self.scheduler.set_timesteps(_a     ,      device=_a  )
    __lowerCamelCase			:			int       			=					self.scheduler.timesteps

    __lowerCamelCase			:			Optional[int]       			=					self.movq.config.latent_channels

    __lowerCamelCase			,__lowerCamelCase			:			List[str]       			=					downscale_height_and_width(_a     ,      _a     ,      self.movq_scale_factor  )

    # create initial latent
    __lowerCamelCase			:			Dict       			=					self.prepare_latents(
        (batch_size, num_channels_latents, height, width)     ,      image_embeds.dtype     ,      _a     ,      _a     ,      _a     ,      self.scheduler     ,      )

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

      __lowerCamelCase			:			Tuple       			=					{'image_embeds': image_embeds, 'hint': hint}
      __lowerCamelCase			:			Optional[Any]       			=					self.unet(
          sample=_a     ,      timestep=_a     ,      encoder_hidden_states=_a     ,      added_cond_kwargs=_a     ,      return_dict=_a     ,      )[0]

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

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

    if output_type == "pil":
      __lowerCamelCase			:			List[str]       			=					self.numpy_to_pil(_a  )

    if not return_dict:
      return (image,)

    return ImagePipelineOutput(images=_a  )











'''simple docstring'''





import argparse
from collections import defaultdict

import yaml


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

def 			a_	(				_lowerCAmelCase			) -> Any:
  __lowerCamelCase			:			Optional[int]       			=					defaultdict(_lowerCAmelCase			)
  __lowerCamelCase			:			Optional[Any]       			=					[]
  __lowerCamelCase			:			Union[str, Any]       			=					[]
  for doc in doc_list:
    if "local" in doc:
      counts[doc["local"]] += 1

    if doc["title"].lower() == "overview":
      overview_doc.append({'local': doc['local'], 'title': doc['title']}			)
    else:
      new_doc_list.append(_lowerCAmelCase			)

  __lowerCamelCase			:			Dict       			=					new_doc_list
  __lowerCamelCase			:			Optional[Any]       			=					[key for key, value in counts.items() if value > 1]

  __lowerCamelCase			:			int       			=					[]
  for duplicate_key in duplicates:
    __lowerCamelCase			:			int       			=					list({doc['title'] for doc in doc_list if doc['local'] == duplicate_key}			)
    if len(_lowerCAmelCase			) > 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 doc_list if 'local' not in counts or counts[doc['local']] == 1]			)
  __lowerCamelCase			:			Dict       			=					sorted(_lowerCAmelCase			,key=lambda _lowerCAmelCase			: s["title"].lower()			)

  # "overview" gets special treatment and is always first
  if len(_lowerCAmelCase			) > 1:
    raise ValueError('{doc_list} has two \'overview\' docs which is not allowed.'			)

  overview_doc.extend(_lowerCAmelCase			)

  # Sort
  return overview_doc

def 			a_	(				_lowerCAmelCase=False			) -> Optional[Any]:
  with open(_lowerCAmelCase			,encoding='utf-8'			) as f:
    __lowerCamelCase			:			str       			=					yaml.safe_load(f.read()			)

  # Get to the API doc
  __lowerCamelCase			:			Union[str, Any]       			=					0
  while content[api_idx]["title"] != "API":
    api_idx += 1
  __lowerCamelCase			:			List[str]       			=					content[api_idx]['sections']

  # Then to the model doc
  __lowerCamelCase			:			List[Any]       			=					0
  while api_doc[scheduler_idx]["title"] != "Schedulers":
    scheduler_idx += 1

  __lowerCamelCase			:			List[Any]       			=					api_doc[scheduler_idx]['sections']
  __lowerCamelCase			:			Optional[Any]       			=					clean_doc_toc(_lowerCAmelCase			)

  __lowerCamelCase			:			str       			=					False
  if new_scheduler_doc != scheduler_doc:
    __lowerCamelCase			:			int       			=					True
    if overwrite:
      __lowerCamelCase			:			Any       			=					new_scheduler_doc

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

def 			a_	(				_lowerCAmelCase=False			) -> List[Any]:
  with open(_lowerCAmelCase			,encoding='utf-8'			) as f:
    __lowerCamelCase			:			List[str]       			=					yaml.safe_load(f.read()			)

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

  # Then to the model doc
  __lowerCamelCase			:			List[Any]       			=					0
  while api_doc[pipeline_idx]["title"] != "Pipelines":
    pipeline_idx += 1

  __lowerCamelCase			:			int       			=					False
  __lowerCamelCase			:			str       			=					api_doc[pipeline_idx]['sections']
  __lowerCamelCase			:			Optional[Any]       			=					[]

  # sort sub pipeline docs
  for pipeline_doc in pipeline_docs:
    if "section" in pipeline_doc:
      __lowerCamelCase			:			str       			=					pipeline_doc['section']
      __lowerCamelCase			:			Optional[Any]       			=					clean_doc_toc(_lowerCAmelCase			)
      if overwrite:
        __lowerCamelCase			:			Union[str, Any]       			=					new_sub_pipeline_doc
    new_pipeline_docs.append(_lowerCAmelCase			)

  # sort overall pipeline doc
  __lowerCamelCase			:			int       			=					clean_doc_toc(_lowerCAmelCase			)

  if new_pipeline_docs != pipeline_docs:
    __lowerCamelCase			:			Tuple       			=					True
    if overwrite:
      __lowerCamelCase			:			int       			=					new_pipeline_docs

  if diff:
    if overwrite:
      __lowerCamelCase			:			Tuple       			=					api_doc
      with open(_lowerCAmelCase			,'w'			,encoding='utf-8'			) as f:
        f.write(yaml.dump(_lowerCAmelCase			,allow_unicode=_lowerCAmelCase			)			)
    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__":
   _UpperCamelCase											=		argparse.ArgumentParser()
   parser.add_argument('--fix_and_overwrite', action='store_true', help='Whether to fix inconsistencies.')
   _UpperCamelCase											=		parser.parse_args()

   check_scheduler_doc(args.fix_and_overwrite)
   check_pipeline_doc(args.fix_and_overwrite)






