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


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

__lowerCAmelCase         =    """
    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 					UpperCAmelCase_		(__a      :    Any ,	__a      :    str ,	__a      :    Optional[Any]=8 ):






   """simple docstring"""

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







class UpperCAmelCase__	( lowercase__	):





  """simple docstring"""

  def __init__(		self      :     Tuple							,_a      :     UNetaDConditionModel							,_a      :     DDPMScheduler							,_a      :     VQModel							,):


     '''simple docstring'''






     super().__init__()

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


  def   __lowercase		(		self      :     Optional[int]							,_a      :     Optional[int]							,_a      :     Optional[Any]							,_a      :     int							,_a      :     Any							,_a      :     List[str]							,_a      :     Optional[Any]       ):


     '''simple docstring'''






     if latents is None:
        _a				:	Optional[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}"""       )
        _a				:	Tuple			=   latents.to(_a       )

     _a				:	str			=   latents * scheduler.init_noise_sigma
     return latents


  def   __lowercase		(		self      :     str							,_a      :     Dict=0       ):


     '''simple docstring'''






     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				:	Union[str, Any]			=   [
         self.unet,
         self.movq,
     ]
     for cpu_offloaded_model in models:
        if cpu_offloaded_model is not None:
           cpu_offload(_a							,_a       )


  def   __lowercase		(		self      :     Tuple							,_a      :     Optional[Any]=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.'       )

     _a				:	Union[str, 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)

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

     # We'll offload the last model manually.
     _a				:	Dict			=   hook


  @property
  # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline._execution_device
  def   __lowercase		(		self      :     Union[str, Any]       ):


     '''simple docstring'''






     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      :     Union[str, Any]							,_a      :     Union[torch.FloatTensor, List[torch.FloatTensor]]							,_a      :     Union[torch.FloatTensor, List[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							,):


     '''simple docstring'''






     _a				:	List[Any]			=   self._execution_device

     _a				:	List[Any]			=   guidance_scale > 1.0

     if isinstance(_a							,_a       ):
        _a				:	str			=   torch.cat(_a							,dim=0       )
     _a				:	Optional[Any]			=   image_embeds.shape[0] * num_images_per_prompt
     if isinstance(_a							,_a       ):
        _a				:	Optional[Any]			=   torch.cat(_a							,dim=0       )

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

     _a				:	Tuple			=   torch.cat([negative_image_embeds, image_embeds]							,dim=0       ).to(dtype=self.unet.dtype							,device=_a       )

     self.scheduler.set_timesteps(_a							,device=_a       )
     _a				:	Any			=   self.scheduler.timesteps

     _a				:	List[Any]			=   self.unet.config.in_channels

     _a,  _a				:	Any			=   downscale_height_and_width(_a							,_a							,self.movq_scale_factor       )

     # create initial latent
     _a				:	List[str]			=   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
        _a				:	List[str]			=   torch.cat([latents] * 2       ) if do_classifier_free_guidance else latents

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

        if do_classifier_free_guidance:
           _a,  _a				:	str			=   noise_pred.split(latents.shape[1]							,dim=1       )
           _a,  _a				:	str			=   noise_pred.chunk(2       )
           _a,  _a				:	Any			=   variance_pred.chunk(2       )
           _a				:	Union[str, Any]			=   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				:	Dict			=   noise_pred.split(latents.shape[1]							,dim=1       )

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

     if output_type == "pil":
        _a				:	Tuple			=   self.numpy_to_pil(_a       )

     if not return_dict:
        return (image,)

     return ImagePipelineOutput(images=_a       )





'''simple docstring'''







from typing import List, Optional, Tuple, Union

import torch

from ...models import UNetaDModel
from ...schedulers import KarrasVeScheduler
from ...utils import randn_tensor
from ..pipeline_utils import DiffusionPipeline, ImagePipelineOutput







class UpperCAmelCase__	( lowercase__	):





  """simple docstring"""

  __UpperCAmelCase				:				UNetaDModel
  __UpperCAmelCase				:				KarrasVeScheduler


  def __init__(		self      :     Union[str, Any]							,_a      :     UNetaDModel							,_a      :     KarrasVeScheduler       ):


     '''simple docstring'''






     super().__init__()
     self.register_modules(unet=_a							,scheduler=_a       )


  @torch.no_grad()
  def __call__(		self      :     List[Any]							,_a      :     int = 1							,_a      :     int = 50							,_a      :     Optional[Union[torch.Generator, List[torch.Generator]]] = None							,_a      :     Optional[str] = "pil"							,_a      :     bool = True							,**_a      :     List[Any]							,):


     '''simple docstring'''






     _a				:	Any			=   self.unet.config.sample_size
     _a				:	Optional[int]			=   (batch_size, 3, img_size, img_size)

     _a				:	Dict			=   self.unet

     # sample x_0 ~ N(0, sigma_0^2 * I)
     _a				:	Dict			=   randn_tensor(_a							,generator=_a							,device=self.device       ) * self.scheduler.init_noise_sigma

     self.scheduler.set_timesteps(_a       )

     for t in self.progress_bar(self.scheduler.timesteps       ):
        # here sigma_t == t_i from the paper
        _a				:	Optional[int]			=   self.scheduler.schedule[t]
        _a				:	List[str]			=   self.scheduler.schedule[t - 1] if t > 0 else 0

        # 1. Select temporarily increased noise level sigma_hat
        # 2. Add new noise to move from sample_i to sample_hat
        _a,  _a				:	List[Any]			=   self.scheduler.add_noise_to_input(_a							,_a							,generator=_a       )

        # 3. Predict the noise residual given the noise magnitude `sigma_hat`
        # The model inputs and output are adjusted by following eq. (213) in [1].
        _a				:	Optional[int]			=   (sigma_hat / 2) * model((sample_hat + 1) / 2							,sigma_hat / 2       ).sample

        # 4. Evaluate dx/dt at sigma_hat
        # 5. Take Euler step from sigma to sigma_prev
        _a				:	Tuple			=   self.scheduler.step(_a							,_a							,_a							,_a       )

        if sigma_prev != 0:
           # 6. Apply 2nd order correction
           # The model inputs and output are adjusted by following eq. (213) in [1].
           _a				:	Optional[int]			=   (sigma_prev / 2) * model((step_output.prev_sample + 1) / 2							,sigma_prev / 2       ).sample
           _a				:	Optional[Any]			=   self.scheduler.step_correct(
               _a							,_a							,_a							,_a							,step_output.prev_sample							,step_output['derivative']							,)
        _a				:	Dict			=   step_output.prev_sample

     _a				:	Tuple			=   (sample / 2 + 0.5).clamp(0							,1       )
     _a				:	Optional[Any]			=   sample.cpu().permute(0							,2							,3							,1       ).numpy()
     if output_type == "pil":
        _a				:	List[str]			=   self.numpy_to_pil(_a       )

     if not return_dict:
        return (image,)

     return ImagePipelineOutput(images=_a       )



