



from typing import List, Optional, Union

import numpy as np
import PIL
import torch
from PIL import Image

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			  =			'''
    Examples:
        ```py
        >>> from diffusers import KandinskyV22Img2ImgPipeline, KandinskyV22PriorPipeline
        >>> from diffusers.utils import load_image
        >>> import torch

        >>> pipe_prior = KandinskyV22PriorPipeline.from_pretrained(
        ...     "kandinsky-community/kandinsky-2-2-prior", torch_dtype=torch.float16
        ... )
        >>> pipe_prior.to("cuda")

        >>> prompt = "A red cartoon frog, 4k"
        >>> image_emb, zero_image_emb = pipe_prior(prompt, return_dict=False)

        >>> pipe = KandinskyV22Img2ImgPipeline.from_pretrained(
        ...     "kandinsky-community/kandinsky-2-2-decoder", torch_dtype=torch.float16
        ... )
        >>> pipe.to("cuda")

        >>> init_image = load_image(
        ...     "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main"
        ...     "/kandinsky/frog.png"
        ... )

        >>> image = pipe(
        ...     image=init_image,
        ...     image_embeds=image_emb,
        ...     negative_image_embeds=zero_image_emb,
        ...     height=768,
        ...     width=768,
        ...     num_inference_steps=100,
        ...     strength=0.2,
        ... ).images

        >>> image[0].save("red_frog.png")
        ```
'''
def 						UpperCamelCase      (  snake_case__			:       Union[str, Any]			,	snake_case__			:       Any			,	snake_case__			:       Union[str, Any]=8      )      ->	Optional[int]:
   UpperCamelCase       :    Dict		      =      height // scale_factor**2
   if height % scale_factor**2 != 0:
      new_height += 1
   UpperCamelCase       :    Tuple		      =      width // scale_factor**2
   if width % scale_factor**2 != 0:
      new_width += 1
   return new_height * scale_factor, new_width * scale_factor
def 						UpperCamelCase      (  snake_case__			:       Any			,	snake_case__			:       int=512			,	snake_case__			:       List[str]=512      )      ->	Dict:
   UpperCamelCase       :    int		      =      pil_image.resize((w, h)			,	resample=Image.BICUBIC			,	reducing_gap=1      )
   UpperCamelCase       :    str		      =      np.array(pil_image.convert('RGB'      )      )
   UpperCamelCase       :    Dict		      =      arr.astype(np.floataa      ) / 127.5 - 1
   UpperCamelCase       :    int		      =      np.transpose(snake_case__			,	[2, 0, 1]      )
   UpperCamelCase       :    Optional[Any]		      =      torch.from_numpy(snake_case__      ).unsqueeze(0      )
   return image



class 							lowerCAmelCase_						(    a__				):
       def __init__(     self,     SCREAMING_SNAKE_CASE_,     SCREAMING_SNAKE_CASE_,     SCREAMING_SNAKE_CASE_,     )	->						Optional[Any]:
          super().__init__()

          self.register_modules(
              unet=SCREAMING_SNAKE_CASE_,     scheduler=SCREAMING_SNAKE_CASE_,     movq=SCREAMING_SNAKE_CASE_,     )
          UpperCamelCase       :    List[Any]		      =      2 ** (len(self.movq.config.block_out_channels			) - 1)
       def    snake_case_			(     self,     SCREAMING_SNAKE_CASE_,     SCREAMING_SNAKE_CASE_,     SCREAMING_SNAKE_CASE_			)	->						Union[str, Any]:
          # get the original timestep using init_timestep
          UpperCamelCase       :    str		      =      min(int(num_inference_steps * strength			),     SCREAMING_SNAKE_CASE_			)

          UpperCamelCase       :    Optional[Any]		      =      max(num_inference_steps - init_timestep,     0			)
          UpperCamelCase       :    Dict		      =      self.scheduler.timesteps[t_start:]

          return timesteps, num_inference_steps - t_start
       def    snake_case_			(     self,     SCREAMING_SNAKE_CASE_,     SCREAMING_SNAKE_CASE_,     SCREAMING_SNAKE_CASE_,     SCREAMING_SNAKE_CASE_,     SCREAMING_SNAKE_CASE_,     SCREAMING_SNAKE_CASE_,     SCREAMING_SNAKE_CASE_=None			)	->						List[Any]:
          if not isinstance(SCREAMING_SNAKE_CASE_,     (torch.Tensor, PIL.Image.Image, list)			):
             raise ValueError(
                 F"""`image` has to be of type `torch.Tensor`, `PIL.Image.Image` or list but is {type(SCREAMING_SNAKE_CASE_			)}"""			)

          UpperCamelCase       :    Tuple		      =      image.to(device=SCREAMING_SNAKE_CASE_,     dtype=SCREAMING_SNAKE_CASE_			)

          UpperCamelCase       :    Any		      =      batch_size * num_images_per_prompt

          if image.shape[1] == 4:
             UpperCamelCase       :    int		      =      image

          else:
             if isinstance(SCREAMING_SNAKE_CASE_,     SCREAMING_SNAKE_CASE_			) and len(SCREAMING_SNAKE_CASE_			) != batch_size:
                raise ValueError(
                    F"""You have passed a list of generators of length {len(SCREAMING_SNAKE_CASE_			)}, but requested an effective batch"""
                    F""" size of {batch_size}. Make sure the batch size matches the length of the generators."""			)

             elif isinstance(SCREAMING_SNAKE_CASE_,     SCREAMING_SNAKE_CASE_			):
                UpperCamelCase       :    List[Any]		      =      [
                    self.movq.encode(image[i : i + 1]			).latent_dist.sample(generator[i]			) for i in range(SCREAMING_SNAKE_CASE_			)
                ]
                UpperCamelCase       :    Union[str, Any]		      =      torch.cat(SCREAMING_SNAKE_CASE_,     dim=0			)
             else:
                UpperCamelCase       :    Optional[Any]		      =      self.movq.encode(SCREAMING_SNAKE_CASE_			).latent_dist.sample(SCREAMING_SNAKE_CASE_			)

             UpperCamelCase       :    Dict		      =      self.movq.config.scaling_factor * init_latents

          UpperCamelCase       :    Optional[int]		      =      torch.cat([init_latents],     dim=0			)

          UpperCamelCase       :    int		      =      init_latents.shape
          UpperCamelCase       :    Optional[int]		      =      randn_tensor(SCREAMING_SNAKE_CASE_,     generator=SCREAMING_SNAKE_CASE_,     device=SCREAMING_SNAKE_CASE_,     dtype=SCREAMING_SNAKE_CASE_			)

          # get latents
          UpperCamelCase       :    Optional[int]		      =      self.scheduler.add_noise(SCREAMING_SNAKE_CASE_,     SCREAMING_SNAKE_CASE_,     SCREAMING_SNAKE_CASE_			)

          UpperCamelCase       :    List[Any]		      =      init_latents

          return latents
       def    snake_case_			(     self,     SCREAMING_SNAKE_CASE_=0			)	->						str:
          if is_accelerate_available():
             from accelerate import cpu_offload
          else:
             raise ImportError('Please install accelerate via `pip install accelerate`'			)

          UpperCamelCase       :    Dict		      =      torch.device(F"""cuda:{gpu_id}"""			)

          UpperCamelCase       :    Optional[Any]		      =      [
              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    snake_case_			(     self,     SCREAMING_SNAKE_CASE_=0			)	->						Any:
          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.'			)

          UpperCamelCase       :    Optional[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)

          UpperCamelCase       :    int		      =      None
          for cpu_offloaded_model in [self.unet, self.movq]:
             UpperCamelCase      ,	UpperCamelCase       :    Any		      =      cpu_offload_with_hook(SCREAMING_SNAKE_CASE_,     SCREAMING_SNAKE_CASE_,     prev_module_hook=SCREAMING_SNAKE_CASE_			)

          # We'll offload the last model manually.
          UpperCamelCase       :    List[str]		      =      hook
       @property
       # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline._execution_device
       def    snake_case_			(     self			)	->						Optional[Any]:
          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,     SCREAMING_SNAKE_CASE_,     SCREAMING_SNAKE_CASE_,     SCREAMING_SNAKE_CASE_,     SCREAMING_SNAKE_CASE_ = 512,     SCREAMING_SNAKE_CASE_ = 512,     SCREAMING_SNAKE_CASE_ = 100,     SCREAMING_SNAKE_CASE_ = 4.0,     SCREAMING_SNAKE_CASE_ = 0.3,     SCREAMING_SNAKE_CASE_ = 1,     SCREAMING_SNAKE_CASE_ = None,     SCREAMING_SNAKE_CASE_ = "pil",     SCREAMING_SNAKE_CASE_ = True,     )	->						Optional[Any]:
          UpperCamelCase       :    str		      =      self._execution_device

          UpperCamelCase       :    Union[str, Any]		      =      guidance_scale > 1.0

          if isinstance(SCREAMING_SNAKE_CASE_,     SCREAMING_SNAKE_CASE_			):
             UpperCamelCase       :    Union[str, Any]		      =      torch.cat(SCREAMING_SNAKE_CASE_,     dim=0			)
          UpperCamelCase       :    int		      =      image_embeds.shape[0]
          if isinstance(SCREAMING_SNAKE_CASE_,     SCREAMING_SNAKE_CASE_			):
             UpperCamelCase       :    Union[str, Any]		      =      torch.cat(SCREAMING_SNAKE_CASE_,     dim=0			)

          if do_classifier_free_guidance:
             UpperCamelCase       :    List[str]		      =      image_embeds.repeat_interleave(SCREAMING_SNAKE_CASE_,     dim=0			)
             UpperCamelCase       :    Optional[Any]		      =      negative_image_embeds.repeat_interleave(SCREAMING_SNAKE_CASE_,     dim=0			)

          UpperCamelCase       :    Union[str, Any]		      =      torch.cat([negative_image_embeds, image_embeds],     dim=0			).to(dtype=self.unet.dtype,     device=SCREAMING_SNAKE_CASE_			)

          if not isinstance(SCREAMING_SNAKE_CASE_,     SCREAMING_SNAKE_CASE_			):
             UpperCamelCase       :    List[str]		      =      [image]
          if not all(isinstance(SCREAMING_SNAKE_CASE_,     (PIL.Image.Image, torch.Tensor)			) for i in image			):
             raise ValueError(
                 F"""Input is in incorrect format: {[type(SCREAMING_SNAKE_CASE_			) for i in image]}. Currently, we only support  PIL image and pytorch tensor"""			)

          UpperCamelCase       :    Any		      =      torch.cat([prepare_image(SCREAMING_SNAKE_CASE_,     SCREAMING_SNAKE_CASE_,     SCREAMING_SNAKE_CASE_			) for i in image],     dim=0			)
          UpperCamelCase       :    List[str]		      =      image.to(dtype=image_embeds.dtype,     device=SCREAMING_SNAKE_CASE_			)

          UpperCamelCase       :    Any		      =      self.movq.encode(SCREAMING_SNAKE_CASE_			)['latents']
          UpperCamelCase       :    List[Any]		      =      latents.repeat_interleave(SCREAMING_SNAKE_CASE_,     dim=0			)
          self.scheduler.set_timesteps(SCREAMING_SNAKE_CASE_,     device=SCREAMING_SNAKE_CASE_			)
          UpperCamelCase      ,	UpperCamelCase       :    str		      =      self.get_timesteps(SCREAMING_SNAKE_CASE_,     SCREAMING_SNAKE_CASE_,     SCREAMING_SNAKE_CASE_			)
          UpperCamelCase       :    Optional[int]		      =      timesteps[:1].repeat(batch_size * num_images_per_prompt			)
          UpperCamelCase      ,	UpperCamelCase       :    Union[str, Any]		      =      downscale_height_and_width(SCREAMING_SNAKE_CASE_,     SCREAMING_SNAKE_CASE_,     self.movq_scale_factor			)
          UpperCamelCase       :    Optional[int]		      =      self.prepare_latents(
              SCREAMING_SNAKE_CASE_,     SCREAMING_SNAKE_CASE_,     SCREAMING_SNAKE_CASE_,     SCREAMING_SNAKE_CASE_,     image_embeds.dtype,     SCREAMING_SNAKE_CASE_,     SCREAMING_SNAKE_CASE_			)
          for i, t in enumerate(self.progress_bar(SCREAMING_SNAKE_CASE_			)			):
             # expand the latents if we are doing classifier free guidance
             UpperCamelCase       :    Optional[Any]		      =      torch.cat([latents] * 2			) if do_classifier_free_guidance else latents

             UpperCamelCase       :    List[Any]		      =      {'image_embeds': image_embeds}
             UpperCamelCase       :    Optional[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:
                UpperCamelCase      ,	UpperCamelCase       :    Tuple		      =      noise_pred.split(latents.shape[1],     dim=1			)
                UpperCamelCase      ,	UpperCamelCase       :    str		      =      noise_pred.chunk(2			)
                UpperCamelCase      ,	UpperCamelCase       :    Any		      =      variance_pred.chunk(2			)
                UpperCamelCase       :    Dict		      =      noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond)
                UpperCamelCase       :    Union[str, 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"]
             ):
                UpperCamelCase      ,	UpperCamelCase       :    Any		      =      noise_pred.split(latents.shape[1],     dim=1			)

             # compute the previous noisy sample x_t -> x_t-1
             UpperCamelCase       :    Dict		      =      self.scheduler.step(
                 SCREAMING_SNAKE_CASE_,     SCREAMING_SNAKE_CASE_,     SCREAMING_SNAKE_CASE_,     generator=SCREAMING_SNAKE_CASE_,     )[0]

          # post-processing
          UpperCamelCase       :    int		      =      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"]:
             UpperCamelCase       :    Dict		      =      image * 0.5 + 0.5
             UpperCamelCase       :    Optional[Any]		      =      image.clamp(0,     1			)
             UpperCamelCase       :    str		      =      image.cpu().permute(0,     2,     3,     1			).float().numpy()

          if output_type == "pil":
             UpperCamelCase       :    Optional[Any]		      =      self.numpy_to_pil(SCREAMING_SNAKE_CASE_			)

          if not return_dict:
             return (image,)

          return ImagePipelineOutput(images=SCREAMING_SNAKE_CASE_			)







'''simple docstring'''
import functools
import operator

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


__SCREAMING_SNAKE_CASE					:				Dict									=		logging.get_logger(__name__)

__SCREAMING_SNAKE_CASE					:				int									=		{
    '''facebook/wav2vec2-base-960h''': '''https://huggingface.co/facebook/wav2vec2-base-960h/resolve/main/config.json''',
    # See all Wav2Vec2 models at https://huggingface.co/models?filter=wav2vec2
}






class 					__lowerCAmelCase   (   lowercase						):

						"""simple docstring"""






						_UpperCAmelCase		:       Dict   						="wav2vec2"






						def __init__(  self							:      str ,  lowerCAmelCase							:      Dict=32 ,  lowerCAmelCase							:      List[Any]=7_68 ,  lowerCAmelCase							:      List[str]=12 ,  lowerCAmelCase							:      List[str]=12 ,  lowerCAmelCase							:      List[Any]=30_72 ,  lowerCAmelCase							:      Optional[Any]="gelu" ,  lowerCAmelCase							:      Optional[int]=0.1 ,  lowerCAmelCase							:      Optional[int]=0.1 ,  lowerCAmelCase							:      List[Any]=0.1 ,  lowerCAmelCase							:      Dict=0.0 ,  lowerCAmelCase							:      Union[str, Any]=0.0 ,  lowerCAmelCase							:      str=0.1 ,  lowerCAmelCase							:      List[str]=0.1 ,  lowerCAmelCase							:      Dict=0.0_2 ,  lowerCAmelCase							:      List[str]=1e-5 ,  lowerCAmelCase							:      Optional[Any]="group" ,  lowerCAmelCase							:      List[Any]="gelu" ,  lowerCAmelCase							:      int=(5_12, 5_12, 5_12, 5_12, 5_12, 5_12, 5_12) ,  lowerCAmelCase							:      Union[str, Any]=(5, 2, 2, 2, 2, 2, 2) ,  lowerCAmelCase							:      Any=(10, 3, 3, 3, 3, 2, 2) ,  lowerCAmelCase							:      str=False ,  lowerCAmelCase							:      Tuple=1_28 ,  lowerCAmelCase							:      Union[str, Any]=16 ,  lowerCAmelCase							:      List[str]=False ,  lowerCAmelCase							:      Union[str, Any]=True ,  lowerCAmelCase							:      Any=0.0_5 ,  lowerCAmelCase							:      Any=10 ,  lowerCAmelCase							:      Dict=2 ,  lowerCAmelCase							:      List[Any]=0.0 ,  lowerCAmelCase							:      int=10 ,  lowerCAmelCase							:      List[str]=0 ,  lowerCAmelCase							:      Tuple=3_20 ,  lowerCAmelCase							:      List[str]=2 ,  lowerCAmelCase							:      int=0.1 ,  lowerCAmelCase							:      List[Any]=1_00 ,  lowerCAmelCase							:      List[str]=2_56 ,  lowerCAmelCase							:      int=2_56 ,  lowerCAmelCase							:      Optional[Any]=0.1 ,  lowerCAmelCase							:      str="sum" ,  lowerCAmelCase							:      List[Any]=False ,  lowerCAmelCase							:      List[str]=False ,  lowerCAmelCase							:      Tuple=2_56 ,  lowerCAmelCase							:      str=(5_12, 5_12, 5_12, 5_12, 15_00) ,  lowerCAmelCase							:      List[str]=(5, 3, 3, 1, 1) ,  lowerCAmelCase							:      Tuple=(1, 2, 3, 1, 1) ,  lowerCAmelCase							:      Optional[Any]=5_12 ,  lowerCAmelCase							:      Any=0 ,  lowerCAmelCase							:      str=1 ,  lowerCAmelCase							:      Any=2 ,  lowerCAmelCase							:      Any=False ,  lowerCAmelCase							:      List[Any]=3 ,  lowerCAmelCase							:      Any=2 ,  lowerCAmelCase							:      List[str]=3 ,  lowerCAmelCase							:      int=None ,  lowerCAmelCase							:      Optional[int]=None ,  **lowerCAmelCase							:      Union[str, Any] ,  ):
													super().__init__(**lowerCAmelCase ,  pad_token_id=lowerCAmelCase ,  bos_token_id=lowerCAmelCase ,  eos_token_id=lowerCAmelCase	)
													A_			    =		hidden_size
													A_			    =		feat_extract_norm
													A_			    =		feat_extract_activation
													A_			    =		list(lowerCAmelCase	)
													A_			    =		list(lowerCAmelCase	)
													A_			    =		list(lowerCAmelCase	)
													A_			    =		conv_bias
													A_			    =		num_conv_pos_embeddings
													A_			    =		num_conv_pos_embedding_groups
													A_			    =		len(self.conv_dim	)
													A_			    =		num_hidden_layers
													A_			    =		intermediate_size
													A_			    =		hidden_act
													A_			    =		num_attention_heads
													A_			    =		hidden_dropout
													A_			    =		attention_dropout
													A_			    =		activation_dropout
													A_			    =		feat_proj_dropout
													A_			    =		final_dropout
													A_			    =		layerdrop
													A_			    =		layer_norm_eps
													A_			    =		initializer_range
													A_			    =		vocab_size
													A_			    =		do_stable_layer_norm
													A_			    =		use_weighted_layer_sum

													if (
													    (len(self.conv_stride	) != self.num_feat_extract_layers)
													    or (len(self.conv_kernel	) != self.num_feat_extract_layers)
													    or (len(self.conv_dim	) != self.num_feat_extract_layers)
													):
																				raise ValueError(
																				    "Configuration for convolutional layers is incorrect. It is required that `len(config.conv_dim)` =="
																				    " `len(config.conv_stride)` == `len(config.conv_kernel)`, but is `len(config.conv_dim) ="
																				    F" {len(self.conv_dim	)}`, `len(config.conv_stride) = {len(self.conv_stride	)}`,"
																				    F" `len(config.conv_kernel) = {len(self.conv_kernel	)}`."	)

													# fine-tuning config parameters for SpecAugment: https://arxiv.org/abs/1904.08779
													A_			    =		apply_spec_augment
													A_			    =		mask_time_prob
													A_			    =		mask_time_length
													A_			    =		mask_time_min_masks
													A_			    =		mask_feature_prob
													A_			    =		mask_feature_length
													A_			    =		mask_feature_min_masks

													# parameters for pretraining with codevector quantized representations
													A_			    =		num_codevectors_per_group
													A_			    =		num_codevector_groups
													A_			    =		contrastive_logits_temperature
													A_			    =		feat_quantizer_dropout
													A_			    =		num_negatives
													A_			    =		codevector_dim
													A_			    =		proj_codevector_dim
													A_			    =		diversity_loss_weight

													# ctc loss
													A_			    =		ctc_loss_reduction
													A_			    =		ctc_zero_infinity

													# adapter
													A_			    =		add_adapter
													A_			    =		adapter_kernel_size
													A_			    =		adapter_stride
													A_			    =		num_adapter_layers
													A_			    =		output_hidden_size or hidden_size
													A_			    =		adapter_attn_dim

													# SequenceClassification-specific parameter. Feel free to ignore for other classes.
													A_			    =		classifier_proj_size

													# XVector-specific parameters. Feel free to ignore for other classes.
													A_			    =		list(lowerCAmelCase	)
													A_			    =		list(lowerCAmelCase	)
													A_			    =		list(lowerCAmelCase	)
													A_			    =		xvector_output_dim






						@property
						def 				_UpperCAmelCase       (  self							:      Tuple	):
													return functools.reduce(operator.mul ,  self.conv_stride ,  1	)





