






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

UpperCAmelCase_   : Optional[Any]   			= """
    Examples:
        ```py
        >>> from diffusers import KandinskyV22Pipeline, KandinskyV22PriorPipeline
        >>> import torch

        >>> pipe_prior = KandinskyV22PriorPipeline.from_pretrained(\"kandinsky-community/kandinsky-2-2-prior\")
        >>> pipe_prior.to(\"cuda\")
        >>> prompt = \"red cat, 4k photo\"
        >>> out = pipe_prior(prompt)
        >>> image_emb = out.image_embeds
        >>> zero_image_emb = out.negative_image_embeds
        >>> pipe = KandinskyV22Pipeline.from_pretrained(\"kandinsky-community/kandinsky-2-2-decoder\")
        >>> pipe.to(\"cuda\")
        >>> image = pipe(
        ...     image_embeds=image_emb,
        ...     negative_image_embeds=zero_image_emb,
        ...     height=768,
        ...     width=768,
        ...     num_inference_steps=50,
        ... ).images
        >>> image[0].save(\"cat.png\")
        ```
"""



def 					_A			(__a    ,				__a    ,				__a=8							) ->       List[str]:
    """simple docstring"""






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





class 							lowerCAmelCase__				(			UpperCAmelCase__ ):






      '''simple docstring'''

      def __init__(						self     :  Optional[int]						,	lowercase_     :  UNetaDConditionModel						,	lowercase_     :  DDPMScheduler						,	lowercase_     :  VQModel						,	):




          '''simple docstring'''


          super().__init__()

          self.register_modules(
              unet=lowercase_						,	scheduler=lowercase_						,	movq=lowercase_						,	)
          SCREAMING_SNAKE_CASE_      :  Tuple    						=		2 ** (len(self.movq.config.block_out_channels) - 1)
      def _SCREAMING_SNAKE_CASE (						self     :  List[Any]						,	lowercase_     :  Optional[Any]						,	lowercase_     :  str						,	lowercase_     :  Union[str, Any]						,	lowercase_     :  Union[str, Any]						,	lowercase_     :  List[str]						,	lowercase_     :  Tuple):




          '''simple docstring'''


          if latents is None:
              SCREAMING_SNAKE_CASE_      :  Optional[Any]    						=		randn_tensor(lowercase_						,	generator=lowercase_						,	device=lowercase_						,	dtype=lowercase_)
          else:
              if latents.shape != shape:
                  raise ValueError(F'Unexpected latents shape, got {latents.shape}, expected {shape}')
              SCREAMING_SNAKE_CASE_      :  Optional[int]    						=		latents.to(lowercase_)

          SCREAMING_SNAKE_CASE_      :  Any    						=		latents * scheduler.init_noise_sigma
          return latents
      def _SCREAMING_SNAKE_CASE (						self     :  Dict						,	lowercase_     :  int=0):




          '''simple docstring'''


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

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

          SCREAMING_SNAKE_CASE_      :  Optional[Any]    						=		[
              self.unet,
              self.movq,
          ]
          for cpu_offloaded_model in models:
              if cpu_offloaded_model is not None:
                  cpu_offload(lowercase_						,	lowercase_)
      def _SCREAMING_SNAKE_CASE (						self     :  Any						,	lowercase_     :  Dict=0):




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

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

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

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

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




          '''simple docstring'''


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




          '''simple docstring'''


          SCREAMING_SNAKE_CASE_      :  Dict    						=		self._execution_device

          SCREAMING_SNAKE_CASE_      :  Optional[int]    						=		guidance_scale > 1.0

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

          if do_classifier_free_guidance:
              SCREAMING_SNAKE_CASE_      :  Optional[int]    						=		image_embeds.repeat_interleave(lowercase_						,	dim=0)
              SCREAMING_SNAKE_CASE_      :  Optional[Any]    						=		negative_image_embeds.repeat_interleave(lowercase_						,	dim=0)

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

          self.scheduler.set_timesteps(lowercase_						,	device=lowercase_)
          SCREAMING_SNAKE_CASE_      :  str    						=		self.scheduler.timesteps

          SCREAMING_SNAKE_CASE_      :  Any    						=		self.unet.config.in_channels

          SCREAMING_SNAKE_CASE_						,			SCREAMING_SNAKE_CASE_      :  Optional[int]    						=		downscale_height_and_width(lowercase_						,	lowercase_						,	self.movq_scale_factor)

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

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

              SCREAMING_SNAKE_CASE_      :  str    						=		{'''image_embeds''': image_embeds}
              SCREAMING_SNAKE_CASE_      :  Optional[Any]    						=		self.unet(
                  sample=lowercase_						,	timestep=lowercase_						,	encoder_hidden_states=lowercase_						,	added_cond_kwargs=lowercase_						,	return_dict=lowercase_						,	)[0]

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

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

          if output_type == "pil":
              SCREAMING_SNAKE_CASE_      :  int    						=		self.numpy_to_pil(lowercase_)

          if not return_dict:
              return (image,)

          return ImagePipelineOutput(images=lowercase_)










"""simple docstring"""






from typing import Dict, List, Optional, Union

import numpy as np

from transformers.utils import is_vision_available
from transformers.utils.generic import TensorType

from ...image_processing_utils import BaseImageProcessor, BatchFeature, get_size_dict
from ...image_transforms import (
    center_crop,
    get_resize_output_image_size,
    normalize,
    rescale,
    resize,
    to_channel_dimension_format,
)
from ...image_utils import (
    IMAGENET_STANDARD_MEAN,
    IMAGENET_STANDARD_STD,
    ChannelDimension,
    ImageInput,
    PILImageResampling,
    is_valid_image,
    to_numpy_array,
    valid_images,
)
from ...utils import logging


if is_vision_available():
  import PIL

UpperCAmelCase_   : int   			= logging.get_logger(__name__)



def 					_A			(__a							) ->       List[List[ImageInput]]:
    """simple docstring"""






    if isinstance(__a    ,				(list, tuple)							) and isinstance(videos[0]    ,				(list, tuple)							) and is_valid_image(videos[0][0]							):
        return videos

    elif isinstance(__a    ,				(list, tuple)							) and is_valid_image(videos[0]							):
        return [videos]

    elif is_valid_image(__a							):
        return [[videos]]

    raise ValueError(f'Could not make batched video from {videos}'							)





class 							lowerCAmelCase__				(			UpperCAmelCase__ ):






      '''simple docstring'''

      __UpperCamelCase	 =  ["pixel_values"]
      def __init__(						self     :  Dict						,	lowercase_     :  bool = True						,	lowercase_     :  Dict[str, int] = None						,	lowercase_     :  PILImageResampling = PILImageResampling.BILINEAR						,	lowercase_     :  bool = True						,	lowercase_     :  Dict[str, int] = None						,	lowercase_     :  bool = True						,	lowercase_     :  Union[int, float] = 1 / 255						,	lowercase_     :  bool = True						,	lowercase_     :  bool = True						,	lowercase_     :  Optional[Union[float, List[float]]] = None						,	lowercase_     :  Optional[Union[float, List[float]]] = None						,	**lowercase_     :  Dict						,	):




          '''simple docstring'''


          super().__init__(**lowercase_)
          SCREAMING_SNAKE_CASE_      :  str    						=		size if size is not None else {'''shortest_edge''': 256}
          SCREAMING_SNAKE_CASE_      :  Optional[int]    						=		get_size_dict(lowercase_						,	default_to_square=lowercase_)
          SCREAMING_SNAKE_CASE_      :  Union[str, Any]    						=		crop_size if crop_size is not None else {'''height''': 224, '''width''': 224}
          SCREAMING_SNAKE_CASE_      :  Dict    						=		get_size_dict(lowercase_						,	param_name='''crop_size''')

          SCREAMING_SNAKE_CASE_      :  Optional[int]    						=		do_resize
          SCREAMING_SNAKE_CASE_      :  List[Any]    						=		size
          SCREAMING_SNAKE_CASE_      :  Tuple    						=		do_center_crop
          SCREAMING_SNAKE_CASE_      :  Dict    						=		crop_size
          SCREAMING_SNAKE_CASE_      :  List[Any]    						=		resample
          SCREAMING_SNAKE_CASE_      :  List[str]    						=		do_rescale
          SCREAMING_SNAKE_CASE_      :  List[str]    						=		rescale_factor
          SCREAMING_SNAKE_CASE_      :  List[Any]    						=		offset
          SCREAMING_SNAKE_CASE_      :  List[Any]    						=		do_normalize
          SCREAMING_SNAKE_CASE_      :  Tuple    						=		image_mean if image_mean is not None else IMAGENET_STANDARD_MEAN
          SCREAMING_SNAKE_CASE_      :  Optional[int]    						=		image_std if image_std is not None else IMAGENET_STANDARD_STD
      def _SCREAMING_SNAKE_CASE (						self     :  List[Any]						,	lowercase_     :  np.ndarray						,	lowercase_     :  Dict[str, int]						,	lowercase_     :  PILImageResampling = PILImageResampling.BILINEAR						,	lowercase_     :  Optional[Union[str, ChannelDimension]] = None						,	**lowercase_     :  Any						,	):




          '''simple docstring'''


          SCREAMING_SNAKE_CASE_      :  Tuple    						=		get_size_dict(lowercase_						,	default_to_square=lowercase_)
          if "shortest_edge" in size:
              SCREAMING_SNAKE_CASE_      :  List[Any]    						=		get_resize_output_image_size(lowercase_						,	size['''shortest_edge''']						,	default_to_square=lowercase_)
          elif "height" in size and "width" in size:
              SCREAMING_SNAKE_CASE_      :  Union[str, Any]    						=		(size['''height'''], size['''width'''])
          else:
              raise ValueError(F'Size must have \'height\' and \'width\' or \'shortest_edge\' as keys. Got {size.keys()}')
          return resize(lowercase_						,	size=lowercase_						,	resample=lowercase_						,	data_format=lowercase_						,	**lowercase_)
      def _SCREAMING_SNAKE_CASE (						self     :  Any						,	lowercase_     :  np.ndarray						,	lowercase_     :  Dict[str, int]						,	lowercase_     :  Optional[Union[str, ChannelDimension]] = None						,	**lowercase_     :  List[Any]						,	):




          '''simple docstring'''


          SCREAMING_SNAKE_CASE_      :  List[str]    						=		get_size_dict(lowercase_)
          if "height" not in size or "width" not in size:
              raise ValueError(F'Size must have \'height\' and \'width\' as keys. Got {size.keys()}')
          return center_crop(lowercase_						,	size=(size['''height'''], size['''width'''])						,	data_format=lowercase_						,	**lowercase_)
      def _SCREAMING_SNAKE_CASE (						self     :  List[Any]						,	lowercase_     :  np.ndarray						,	lowercase_     :  Union[int, float]						,	lowercase_     :  bool = True						,	lowercase_     :  Optional[Union[str, ChannelDimension]] = None						,	**lowercase_     :  Optional[int]						,	):




          '''simple docstring'''


          SCREAMING_SNAKE_CASE_      :  List[Any]    						=		image.astype(np.floataa)
          if offset:
              SCREAMING_SNAKE_CASE_      :  Tuple    						=		image - (scale / 2)
          return rescale(lowercase_						,	scale=lowercase_						,	data_format=lowercase_						,	**lowercase_)
      def _SCREAMING_SNAKE_CASE (						self     :  List[str]						,	lowercase_     :  np.ndarray						,	lowercase_     :  Union[float, List[float]]						,	lowercase_     :  Union[float, List[float]]						,	lowercase_     :  Optional[Union[str, ChannelDimension]] = None						,	**lowercase_     :  List[str]						,	):




          '''simple docstring'''


          return normalize(lowercase_						,	mean=lowercase_						,	std=lowercase_						,	data_format=lowercase_						,	**lowercase_)
      def _SCREAMING_SNAKE_CASE (						self     :  List[str]						,	lowercase_     :  ImageInput						,	lowercase_     :  bool = None						,	lowercase_     :  Dict[str, int] = None						,	lowercase_     :  PILImageResampling = None						,	lowercase_     :  bool = None						,	lowercase_     :  Dict[str, int] = None						,	lowercase_     :  bool = None						,	lowercase_     :  float = None						,	lowercase_     :  bool = None						,	lowercase_     :  bool = None						,	lowercase_     :  Optional[Union[float, List[float]]] = None						,	lowercase_     :  Optional[Union[float, List[float]]] = None						,	lowercase_     :  Optional[ChannelDimension] = ChannelDimension.FIRST						,	):




          '''simple docstring'''


          if do_resize and size is None or resample is None:
              raise ValueError('''Size and resample must be specified if do_resize is True.''')

          if do_center_crop and crop_size is None:
              raise ValueError('''Crop size must be specified if do_center_crop is True.''')

          if do_rescale and rescale_factor is None:
              raise ValueError('''Rescale factor must be specified if do_rescale is True.''')

          if do_normalize and (image_mean is None or image_std is None):
              raise ValueError('''Image mean and std must be specified if do_normalize is True.''')

          if offset and not do_rescale:
              raise ValueError('''For offset, do_rescale must also be set to True.''')

          # All transformations expect numpy arrays.
          SCREAMING_SNAKE_CASE_      :  List[str]    						=		to_numpy_array(lowercase_)

          if do_resize:
              SCREAMING_SNAKE_CASE_      :  List[Any]    						=		self.resize(image=lowercase_						,	size=lowercase_						,	resample=lowercase_)

          if do_center_crop:
              SCREAMING_SNAKE_CASE_      :  Dict    						=		self.center_crop(lowercase_						,	size=lowercase_)

          if do_rescale:
              SCREAMING_SNAKE_CASE_      :  int    						=		self.rescale(image=lowercase_						,	scale=lowercase_						,	offset=lowercase_)

          if do_normalize:
              SCREAMING_SNAKE_CASE_      :  Dict    						=		self.normalize(image=lowercase_						,	mean=lowercase_						,	std=lowercase_)

          SCREAMING_SNAKE_CASE_      :  List[Any]    						=		to_channel_dimension_format(lowercase_						,	lowercase_)
          return image



      def _SCREAMING_SNAKE_CASE (						self     :  Dict						,	lowercase_     :  ImageInput						,	lowercase_     :  bool = None						,	lowercase_     :  Dict[str, int] = None						,	lowercase_     :  PILImageResampling = None						,	lowercase_     :  bool = None						,	lowercase_     :  Dict[str, int] = None						,	lowercase_     :  bool = None						,	lowercase_     :  float = None						,	lowercase_     :  bool = None						,	lowercase_     :  bool = None						,	lowercase_     :  Optional[Union[float, List[float]]] = None						,	lowercase_     :  Optional[Union[float, List[float]]] = None						,	lowercase_     :  Optional[Union[str, TensorType]] = None						,	lowercase_     :  ChannelDimension = ChannelDimension.FIRST						,	**lowercase_     :  Optional[Any]						,	):




          '''simple docstring'''


          SCREAMING_SNAKE_CASE_      :  Optional[Any]    						=		do_resize if do_resize is not None else self.do_resize
          SCREAMING_SNAKE_CASE_      :  int    						=		resample if resample is not None else self.resample
          SCREAMING_SNAKE_CASE_      :  Union[str, Any]    						=		do_center_crop if do_center_crop is not None else self.do_center_crop
          SCREAMING_SNAKE_CASE_      :  Dict    						=		do_rescale if do_rescale is not None else self.do_rescale
          SCREAMING_SNAKE_CASE_      :  Dict    						=		rescale_factor if rescale_factor is not None else self.rescale_factor
          SCREAMING_SNAKE_CASE_      :  Dict    						=		offset if offset is not None else self.offset
          SCREAMING_SNAKE_CASE_      :  str    						=		do_normalize if do_normalize is not None else self.do_normalize
          SCREAMING_SNAKE_CASE_      :  Dict    						=		image_mean if image_mean is not None else self.image_mean
          SCREAMING_SNAKE_CASE_      :  List[str]    						=		image_std if image_std is not None else self.image_std

          SCREAMING_SNAKE_CASE_      :  Union[str, Any]    						=		size if size is not None else self.size
          SCREAMING_SNAKE_CASE_      :  Tuple    						=		get_size_dict(lowercase_						,	default_to_square=lowercase_)
          SCREAMING_SNAKE_CASE_      :  Any    						=		crop_size if crop_size is not None else self.crop_size
          SCREAMING_SNAKE_CASE_      :  Dict    						=		get_size_dict(lowercase_						,	param_name='''crop_size''')

          if not valid_images(lowercase_):
              raise ValueError(
                  '''Invalid image type. Must be of type PIL.Image.Image, numpy.ndarray, '''
                  '''torch.Tensor, tf.Tensor or jax.ndarray.''')

          SCREAMING_SNAKE_CASE_      :  Tuple    						=		make_batched(lowercase_)

          SCREAMING_SNAKE_CASE_      :  Optional[Any]    						=		[
              [
                  self._preprocess_image(
                      image=lowercase_						,	do_resize=lowercase_						,	size=lowercase_						,	resample=lowercase_						,	do_center_crop=lowercase_						,	crop_size=lowercase_						,	do_rescale=lowercase_						,	rescale_factor=lowercase_						,	offset=lowercase_						,	do_normalize=lowercase_						,	image_mean=lowercase_						,	image_std=lowercase_						,	data_format=lowercase_						,	)
                  for img in video
              ]
              for video in videos
          ]

          SCREAMING_SNAKE_CASE_      :  int    						=		{'''pixel_values''': videos}
          return BatchFeature(data=lowercase_						,	tensor_type=lowercase_)


