


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_      : Any  					=       logging.get_logger(__name__)  # pylint: disable=invalid-name

UpperCAmelCase_      : Optional[Any]  					=       '\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        SCREAMING_SNAKE_CASE_						(			__A :			str  ,    __A :			List[Any]  ,    __A :			int=8				)	->		Optional[Any]:

       """simple docstring"""

       a_					:		int							=  height // scale_factor**2
       if height % scale_factor**2 != 0:
              new_height += 1
       a_					:		Optional[int]							=  width // scale_factor**2
       if width % scale_factor**2 != 0:
              new_width += 1
       return new_height * scale_factor, new_width * scale_factor
class  SCREAMING_SNAKE_CASE__  (				lowercase__			):
      def __init__(       self				:       List[Any]     ,							SCREAMING_SNAKE_CASE__				:       UNetaDConditionModel     ,							SCREAMING_SNAKE_CASE__				:       DDPMScheduler     ,							SCREAMING_SNAKE_CASE__				:       VQModel     ,							)		->			Optional[int]:
             super().__init__()

             self.register_modules(
                 unet=SCREAMING_SNAKE_CASE__     ,							scheduler=SCREAMING_SNAKE_CASE__     ,							movq=SCREAMING_SNAKE_CASE__     ,							)
             a_					:		Optional[int]							=  2 ** (len(self.movq.config.block_out_channels     ) - 1)
      def        SCREAMING_SNAKE_CASE							(       self				:       Any     ,							SCREAMING_SNAKE_CASE__				:       int     ,							SCREAMING_SNAKE_CASE__				:       int     ,							SCREAMING_SNAKE_CASE__				:       int     ,							SCREAMING_SNAKE_CASE__				:       str     ,							SCREAMING_SNAKE_CASE__				:       int     ,							SCREAMING_SNAKE_CASE__				:       Optional[int]     )		->			Optional[Any]:
             if latents is None:
                    a_					:		Tuple							=  randn_tensor(SCREAMING_SNAKE_CASE__     ,							generator=SCREAMING_SNAKE_CASE__     ,							device=SCREAMING_SNAKE_CASE__     ,							dtype=SCREAMING_SNAKE_CASE__     )
             else:
                    if latents.shape != shape:
                           raise ValueError(F"""Unexpected latents shape, got {latents.shape}, expected {shape}"""     )
                    a_					:		Tuple							=  latents.to(SCREAMING_SNAKE_CASE__     )

             a_					:		int							=  latents * scheduler.init_noise_sigma
             return latents
      def        SCREAMING_SNAKE_CASE							(       self				:       str     ,							SCREAMING_SNAKE_CASE__				:       Optional[Any]=0     )		->			Any:
             if is_accelerate_available():
                    from accelerate import cpu_offload
             else:
                    raise ImportError('Please install accelerate via `pip install accelerate`'     )

             a_					:		int							=  torch.device(F"""cuda:{gpu_id}"""     )

             a_					:		Tuple							=  [
                 self.unet,
                 self.movq,
             ]
             for cpu_offloaded_model in models:
                    if cpu_offloaded_model is not None:
                           cpu_offload(SCREAMING_SNAKE_CASE__     ,							SCREAMING_SNAKE_CASE__     )
      def        SCREAMING_SNAKE_CASE							(       self				:       Tuple     ,							SCREAMING_SNAKE_CASE__				:       Union[str, Any]=0     )		->			List[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.'     )

             a_					:		Any							=  torch.device(F"""cuda:{gpu_id}"""     )

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

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

             # We'll offload the last model manually.
             a_					:		Union[str, Any]							=  hook
      @property
      # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline._execution_device
      def        SCREAMING_SNAKE_CASE							(       self				:       Union[str, Any]     )		->			List[str]:
             if not hasattr(self.unet     ,							'_hf_hook'     ):
                    return self.device
             for module in self.unet.modules():
                    if (
                        hasattr(SCREAMING_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(SCREAMING_SNAKE_CASE__     )
      def __call__(       self				:       List[Any]     ,							SCREAMING_SNAKE_CASE__				:       Union[torch.FloatTensor, List[torch.FloatTensor]]     ,							SCREAMING_SNAKE_CASE__				:       Union[torch.FloatTensor, List[torch.FloatTensor]]     ,							SCREAMING_SNAKE_CASE__				:       torch.FloatTensor     ,							SCREAMING_SNAKE_CASE__				:       int = 5_1_2     ,							SCREAMING_SNAKE_CASE__				:       int = 5_1_2     ,							SCREAMING_SNAKE_CASE__				:       int = 1_0_0     ,							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     ,							)		->			Optional[int]:
             a_					:		Union[str, Any]							=  self._execution_device

             a_					:		int							=  guidance_scale > 1.0

             if isinstance(SCREAMING_SNAKE_CASE__     ,							SCREAMING_SNAKE_CASE__     ):
                    a_					:		List[Any]							=  torch.cat(SCREAMING_SNAKE_CASE__     ,							dim=0     )
             if isinstance(SCREAMING_SNAKE_CASE__     ,							SCREAMING_SNAKE_CASE__     ):
                    a_					:		str							=  torch.cat(SCREAMING_SNAKE_CASE__     ,							dim=0     )
             if isinstance(SCREAMING_SNAKE_CASE__     ,							SCREAMING_SNAKE_CASE__     ):
                    a_					:		List[str]							=  torch.cat(SCREAMING_SNAKE_CASE__     ,							dim=0     )

             a_					:		Union[str, Any]							=  image_embeds.shape[0] * num_images_per_prompt

             if do_classifier_free_guidance:
                    a_					:		int							=  image_embeds.repeat_interleave(SCREAMING_SNAKE_CASE__     ,							dim=0     )
                    a_					:		Optional[Any]							=  negative_image_embeds.repeat_interleave(SCREAMING_SNAKE_CASE__     ,							dim=0     )
                    a_					:		Optional[Any]							=  hint.repeat_interleave(SCREAMING_SNAKE_CASE__     ,							dim=0     )

             a_					:		Union[str, Any]							=  torch.cat([negative_image_embeds, image_embeds]     ,							dim=0     ).to(dtype=self.unet.dtype     ,							device=SCREAMING_SNAKE_CASE__     )
             a_					:		Tuple							=  torch.cat([hint, hint]     ,							dim=0     ).to(dtype=self.unet.dtype     ,							device=SCREAMING_SNAKE_CASE__     )
             self.scheduler.set_timesteps(SCREAMING_SNAKE_CASE__     ,							device=SCREAMING_SNAKE_CASE__     )
             a_					:		List[str]							=  self.scheduler.timesteps

             a_					:		str							=  self.movq.config.latent_channels

             a_							,		a_					:		Optional[Any]							=  downscale_height_and_width(SCREAMING_SNAKE_CASE__     ,							SCREAMING_SNAKE_CASE__     ,							self.movq_scale_factor     )

             # create initial latent
             a_					:		Tuple							=  self.prepare_latents(
                 (batch_size, num_channels_latents, height, width)     ,							image_embeds.dtype     ,							SCREAMING_SNAKE_CASE__     ,							SCREAMING_SNAKE_CASE__     ,							SCREAMING_SNAKE_CASE__     ,							self.scheduler     ,							)

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

                    a_					:		Dict							=  {'image_embeds': image_embeds, 'hint': hint}
                    a_					:		List[Any]							=  self.unet(
                        sample=SCREAMING_SNAKE_CASE__     ,							timestep=SCREAMING_SNAKE_CASE__     ,							encoder_hidden_states=SCREAMING_SNAKE_CASE__     ,							added_cond_kwargs=SCREAMING_SNAKE_CASE__     ,							return_dict=SCREAMING_SNAKE_CASE__     ,							)[0]

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

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

             if output_type == "pil":
                    a_					:		List[Any]							=  self.numpy_to_pil(SCREAMING_SNAKE_CASE__     )

             if not return_dict:
                    return (image,)

             return ImagePipelineOutput(images=SCREAMING_SNAKE_CASE__     )









from collections import OrderedDict
from typing import TYPE_CHECKING, Any, Mapping, Optional, Union

from ...configuration_utils import PretrainedConfig
from ...onnx import OnnxConfig, OnnxSeqaSeqConfigWithPast
from ...utils import logging


if TYPE_CHECKING:
    from ...feature_extraction_utils import FeatureExtractionMixin
    from ...tokenization_utils_base import PreTrainedTokenizerBase
    from ...utils import TensorType

lowercase      :       List[str]									=  logging.get_logger(__name__)

lowercase      :       Any									=  {
    """openai/whisper-base""": """https://huggingface.co/openai/whisper-base/resolve/main/config.json""",
}

# fmt: off
lowercase      :       List[str]									=  [
    1, 2, 7, 8, 9, 1_0, 1_4, 2_5,
    2_6, 2_7, 2_8, 2_9, 3_1, 5_8, 5_9, 6_0, 6_1, 6_2,
    6_3, 9_0, 9_1, 9_2, 9_3, 3_5_7, 3_6_6, 4_3_8, 5_3_2, 6_8_5,
    7_0_5, 7_9_6, 9_3_0, 1_0_5_8, 1_2_2_0, 1_2_6_7, 1_2_7_9, 1_3_0_3, 1_3_4_3, 1_3_7_7,
    1_3_9_1, 1_6_3_5, 1_7_8_2, 1_8_7_5, 2_1_6_2, 2_3_6_1, 2_4_8_8, 3_4_6_7, 4_0_0_8, 4_2_1_1,
    4_6_0_0, 4_8_0_8, 5_2_9_9, 5_8_5_5, 6_3_2_9, 7_2_0_3, 9_6_0_9, 9_9_5_9, 1_0_5_6_3, 1_0_7_8_6,
    1_1_4_2_0, 1_1_7_0_9, 1_1_9_0_7, 1_3_1_6_3, 1_3_6_9_7, 1_3_7_0_0, 1_4_8_0_8, 1_5_3_0_6, 1_6_4_1_0, 1_6_7_9_1,
    1_7_9_9_2, 1_9_2_0_3, 1_9_5_1_0, 2_0_7_2_4, 2_2_3_0_5, 2_2_9_3_5, 2_7_0_0_7, 3_0_1_0_9, 3_0_4_2_0, 3_3_4_0_9,
    3_4_9_4_9, 4_0_2_8_3, 4_0_4_9_3, 4_0_5_4_9, 4_7_2_8_2, 4_9_1_4_6, 5_0_2_5_7, 5_0_3_5_9, 5_0_3_6_0, 5_0_3_6_1
]
lowercase      :       List[Any]									=  [
    1, 2, 7, 8, 9, 1_0, 1_4, 2_5,
    2_6, 2_7, 2_8, 2_9, 3_1, 5_8, 5_9, 6_0, 6_1, 6_2,
    6_3, 9_0, 9_1, 9_2, 9_3, 3_5_9, 5_0_3, 5_2_2, 5_4_2, 8_7_3,
    8_9_3, 9_0_2, 9_1_8, 9_2_2, 9_3_1, 1_3_5_0, 1_8_5_3, 1_9_8_2, 2_4_6_0, 2_6_2_7,
    3_2_4_6, 3_2_5_3, 3_2_6_8, 3_5_3_6, 3_8_4_6, 3_9_6_1, 4_1_8_3, 4_6_6_7, 6_5_8_5, 6_6_4_7,
    7_2_7_3, 9_0_6_1, 9_3_8_3, 1_0_4_2_8, 1_0_9_2_9, 1_1_9_3_8, 1_2_0_3_3, 1_2_3_3_1, 1_2_5_6_2, 1_3_7_9_3,
    1_4_1_5_7, 1_4_6_3_5, 1_5_2_6_5, 1_5_6_1_8, 1_6_5_5_3, 1_6_6_0_4, 1_8_3_6_2, 1_8_9_5_6, 2_0_0_7_5, 2_1_6_7_5,
    2_2_5_2_0, 2_6_1_3_0, 2_6_1_6_1, 2_6_4_3_5, 2_8_2_7_9, 2_9_4_6_4, 3_1_6_5_0, 3_2_3_0_2, 3_2_4_7_0, 3_6_8_6_5,
    4_2_8_6_3, 4_7_4_2_5, 4_9_8_7_0, 5_0_2_5_4, 5_0_2_5_8, 5_0_3_6_0, 5_0_3_6_1, 5_0_3_6_2
]





class  A__    (							__UpperCAmelCase	):




   """simple docstring"""



   __A   :				int         =  '''whisper'''
   __A   :				List[Any]         =  ['''past_key_values''']
   __A   :				Optional[int]         =  {'''num_attention_heads''': '''encoder_attention_heads''', '''hidden_size''': '''d_model'''}

   def __init__(					self					,  lowercase=5_1865					,  lowercase=80					,  lowercase=6					,  lowercase=4					,  lowercase=6					,  lowercase=4					,  lowercase=1536					,  lowercase=1536					,  lowercase=0.0					,  lowercase=0.0					,  lowercase=5_0257					,  lowercase=True					,  lowercase=True					,  lowercase="gelu"					,  lowercase=256					,  lowercase=0.0					,  lowercase=0.0					,  lowercase=0.0					,  lowercase=0.02					,  lowercase=False					,  lowercase=1500					,  lowercase=448					,  lowercase=5_0256					,  lowercase=5_0256					,  lowercase=5_0256					,  lowercase=None					,  lowercase=[220, 5_0256]					,  lowercase=False					,  lowercase=256					,  lowercase=False					,  lowercase=0.05					,  lowercase=10					,  lowercase=2					,  lowercase=0.0					,  lowercase=10					,  lowercase=0					,  lowercase=7					,  **lowercase					,  )			->						str:







          '''simple docstring'''

          a__						:	int =				vocab_size
          a__						:	int =				num_mel_bins
          a__						:	Optional[int] =				d_model
          a__						:	List[str] =				encoder_layers
          a__						:	Dict =				encoder_attention_heads
          a__						:	List[str] =				decoder_layers
          a__						:	Tuple =				decoder_attention_heads
          a__						:	List[str] =				decoder_ffn_dim
          a__						:	Optional[Any] =				encoder_ffn_dim
          a__						:	Tuple =				dropout
          a__						:	Optional[int] =				attention_dropout
          a__						:	Any =				activation_dropout
          a__						:	Any =				activation_function
          a__						:	List[Any] =				init_std
          a__						:	Optional[int] =				encoder_layerdrop
          a__						:	Union[str, Any] =				decoder_layerdrop
          a__						:	Tuple =				use_cache
          a__						:	List[str] =				encoder_layers
          a__						:	Dict =				scale_embedding  # scale factor will be sqrt(d_model) if True
          a__						:	Dict =				max_source_positions
          a__						:	Dict =				max_target_positions

          # Audio Classification-specific parameters. Feel free to ignore for other classes.
          a__						:	Optional[int] =				classifier_proj_size
          a__						:	List[Any] =				use_weighted_layer_sum

          # fine-tuning config parameters for SpecAugment: https://arxiv.org/abs/1904.08779
          a__						:	List[Any] =				apply_spec_augment
          a__						:	int =				mask_time_prob
          a__						:	int =				mask_time_length
          a__						:	List[Any] =				mask_time_min_masks
          a__						:	str =				mask_feature_prob
          a__						:	Optional[int] =				mask_feature_length
          a__						:	Union[str, Any] =				mask_feature_min_masks

          a__						:	Tuple =				median_filter_width

          super().__init__(
              pad_token_id=lowercase					,  bos_token_id=lowercase					,  eos_token_id=lowercase					,  is_encoder_decoder=lowercase					,  decoder_start_token_id=lowercase					,  suppress_tokens=lowercase					,  begin_suppress_tokens=lowercase					,  **lowercase					,  )





class  A__    (							__UpperCAmelCase	):




   """simple docstring"""



   @property
   def 							__lowercase		(					self)			->						Mapping[str, Mapping[int, str]]:







          '''simple docstring'''

          a__						:	List[str] =				OrderedDict(
              [
                  ('input_features', {0: 'batch', 1: 'feature_size', 2: 'encoder_sequence'}),
              ])
          if self.use_past:
                 a__						:	Optional[Any] =				{0: 'batch'}
          else:
                 a__						:	Optional[Any] =				{0: 'batch', 1: 'decoder_sequence'}

          if self.use_past:
                 self.fill_with_past_key_values_(lowercase					,  direction='inputs')

          return common_inputs

   def 							__lowercase		(					self					,  lowercase					,  lowercase = -1					,  lowercase = -1					,  lowercase = False					,  lowercase = None					,  lowercase = 2_2050					,  lowercase = 5.0					,  lowercase = 220					,  )			->						Mapping[str, Any]:







          '''simple docstring'''

          a__						:	Union[str, Any] =				OrderedDict()
          a__						:	int =				OnnxConfig.generate_dummy_inputs(
              self					,  preprocessor=preprocessor.feature_extractor					,  batch_size=lowercase					,  framework=lowercase					,  sampling_rate=lowercase					,  time_duration=lowercase					,  frequency=lowercase					,  )
          a__						:	List[Any] =				encoder_inputs['input_features'].shape[2]
          a__						:	Optional[int] =				encoder_sequence_length // 2 if self.use_past else seq_length

          a__						:	Any =				super().generate_dummy_inputs(
              preprocessor.tokenizer					,  lowercase					,  lowercase					,  lowercase					,  lowercase)

          a__						:	List[str] =				encoder_inputs.pop('input_features')
          a__						:	Optional[int] =				decoder_inputs.pop('decoder_input_ids')

          if "past_key_values" in decoder_inputs:
                 a__						:	List[str] =				decoder_inputs.pop('past_key_values')

          return dummy_inputs

   @property
   def 							__lowercase		(					self)			->						float:







          '''simple docstring'''

          return 1e-3


