




"""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__						=			'''
    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		(		snake_case			:					List[str]      ,       snake_case			:					List[str]      ,       snake_case			:					Optional[Any]=8    ):
					_lowerCAmelCase:int									=  height // scale_factor**2
					if height % scale_factor**2 != 0:
										new_height += 1
					_lowerCAmelCase: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   a__			(       UpperCamelCase_  ):





						def __init__(							self      :	Union[str, Any]     ,a__      :	UNetaDConditionModel     ,a__      :	DDPMScheduler     ,a__      :	VQModel     ,)						->       Any:




											"""simple docstring"""







											super().__init__()

											self.register_modules(
											    unet=a__     ,scheduler=a__     ,movq=a__     ,)
											_lowerCAmelCase:List[str]									=  2 ** (len(self.movq.config.block_out_channels) - 1)





						def       __UpperCamelCase     (							self      :	Any     ,a__      :	Union[str, Any]     ,a__      :	Dict     ,a__      :	List[str]     ,a__      :	str     ,a__      :	Any     ,a__      :	Union[str, Any])						->       Dict:




											"""simple docstring"""







											if latents is None:
																_lowerCAmelCase:str									=  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:Optional[Any]									=  latents.to(a__)

											_lowerCAmelCase:Optional[Any]									=  latents * scheduler.init_noise_sigma
											return latents





						def       __UpperCamelCase     (							self      :	str     ,a__      :	Union[str, Any]=0)						->       int:




											"""simple docstring"""







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

											_lowerCAmelCase:Any									=  torch.device(F'cuda:{gpu_id}')

											_lowerCAmelCase:List[Any]									=  [
											    self.unet,
											    self.movq,
											]
											for cpu_offloaded_model in models:
																if cpu_offloaded_model is not None:
																					cpu_offload(a__     ,a__)





						def       __UpperCamelCase     (							self      :	int     ,a__      :	str=0)						->       List[Any]:




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

											_lowerCAmelCase:List[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:List[Any]									=  None
											for cpu_offloaded_model in [self.unet, self.movq]:
																_lowerCAmelCase							,   _lowerCAmelCase:Tuple									=  cpu_offload_with_hook(a__     ,a__     ,prev_module_hook=a__)

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





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




											"""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      :	Optional[int]     ,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     ,)						->       Optional[int]:




											"""simple docstring"""







											_lowerCAmelCase:str									=  self._execution_device

											_lowerCAmelCase:Tuple									=  guidance_scale > 1.0

											if isinstance(a__     ,a__):
																_lowerCAmelCase:Optional[int]									=  torch.cat(a__     ,dim=0)
											_lowerCAmelCase:Optional[int]									=  image_embeds.shape[0] * num_images_per_prompt
											if isinstance(a__     ,a__):
																_lowerCAmelCase:List[Any]									=  torch.cat(a__     ,dim=0)

											if do_classifier_free_guidance:
																_lowerCAmelCase:Tuple									=  image_embeds.repeat_interleave(a__     ,dim=0)
																_lowerCAmelCase:List[Any]									=  negative_image_embeds.repeat_interleave(a__     ,dim=0)

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

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

											_lowerCAmelCase:Tuple									=  self.unet.config.in_channels

											_lowerCAmelCase							,   _lowerCAmelCase:Tuple									=  downscale_height_and_width(a__     ,a__     ,self.movq_scale_factor)

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

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

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

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

											if not return_dict:
																return (image,)

											return ImagePipelineOutput(images=a__)












"""simple docstring"""






import warnings

from ...utils import logging
from .image_processing_poolformer import PoolFormerImageProcessor


UpperCamelCase__						=			logging.get_logger(__name__)




class   a__			(       UpperCamelCase_  ):





						def __init__(							self      :	Dict     ,*a__      :	List[str]     ,**a__      :	str)						->       None:




											"""simple docstring"""







											warnings.warn(
											    '''The class PoolFormerFeatureExtractor is deprecated and will be removed in version 5 of Transformers.'''
											    ''' Please use PoolFormerImageProcessor instead.'''     ,a__     ,)
											super().__init__(*a__     ,**a__)






