






from typing import List, Optional, Union

import torch
from transformers import (
    XLMRobertaTokenizer,
)

from ...models import UNetaDConditionModel, VQModel
from ...pipelines import DiffusionPipeline
from ...pipelines.pipeline_utils import ImagePipelineOutput
from ...schedulers import DDIMScheduler, DDPMScheduler
from ...utils import (
    is_accelerate_available,
    is_accelerate_version,
    logging,
    randn_tensor,
    replace_example_docstring,
)
from .text_encoder import MultilingualCLIP


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

__A					=	"\n    Examples:\n        ```py\n        >>> from diffusers import KandinskyPipeline, KandinskyPriorPipeline\n        >>> import torch\n\n        >>> pipe_prior = KandinskyPriorPipeline.from_pretrained(\"kandinsky-community/Kandinsky-2-1-prior\")\n        >>> pipe_prior.to(\"cuda\")\n\n        >>> prompt = \"red cat, 4k photo\"\n        >>> out = pipe_prior(prompt)\n        >>> image_emb = out.image_embeds\n        >>> negative_image_emb = out.negative_image_embeds\n\n        >>> pipe = KandinskyPipeline.from_pretrained(\"kandinsky-community/kandinsky-2-1\")\n        >>> pipe.to(\"cuda\")\n\n        >>> image = pipe(\n        ...     prompt,\n        ...     image_embeds=image_emb,\n        ...     negative_image_embeds=negative_image_emb,\n        ...     height=768,\n        ...     width=768,\n        ...     num_inference_steps=100,\n        ... ).images\n\n        >>> image[0].save(\"cat.png\")\n        ```\n"







def       lowercase__	(					A_:						Tuple  ,      A_:						List[str]  ,      A_:						str=8		)	->				Union[str, Any]:







       """simple docstring"""






       __UpperCAmelCase     				=h // scale_factor**2
       if h % scale_factor**2 != 0:
              new_h += 1
       __UpperCAmelCase     				=w // scale_factor**2
       if w % scale_factor**2 != 0:
              new_w += 1
       return new_h * scale_factor, new_w * scale_factor





class 				_A			(     UpperCamelCase				):
    """simple docstring"""

    def __init__(				self      :     List[str] ,     __SCREAMING_SNAKE_CASE      :     MultilingualCLIP ,     __SCREAMING_SNAKE_CASE      :     XLMRobertaTokenizer ,     __SCREAMING_SNAKE_CASE      :     UNetaDConditionModel ,     __SCREAMING_SNAKE_CASE      :     Union[DDIMScheduler, DDPMScheduler] ,     __SCREAMING_SNAKE_CASE      :     VQModel ,     )      ->							Optional[Any]:
           super().__init__()

           self.register_modules(
               text_encoder=__SCREAMING_SNAKE_CASE ,     tokenizer=__SCREAMING_SNAKE_CASE ,     unet=__SCREAMING_SNAKE_CASE ,     scheduler=__SCREAMING_SNAKE_CASE ,     movq=__SCREAMING_SNAKE_CASE ,     )
           __UpperCAmelCase     				=2 ** (len(self.movq.config.block_out_channels      ) - 1)



    def _a							(				self      :     int ,     __SCREAMING_SNAKE_CASE      :     Tuple ,     __SCREAMING_SNAKE_CASE      :     Optional[Any] ,     __SCREAMING_SNAKE_CASE      :     Union[str, Any] ,     __SCREAMING_SNAKE_CASE      :     Tuple ,     __SCREAMING_SNAKE_CASE      :     Union[str, Any] ,     __SCREAMING_SNAKE_CASE      :     Optional[Any]      )      ->							List[str]:
           if latents is None:
                  __UpperCAmelCase     				=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}'''      )
                  __UpperCAmelCase     				=latents.to(__SCREAMING_SNAKE_CASE      )

           __UpperCAmelCase     				=latents * scheduler.init_noise_sigma
           return latents



    def _a							(				self      :     Tuple ,     __SCREAMING_SNAKE_CASE      :     List[str] ,     __SCREAMING_SNAKE_CASE      :     Tuple ,     __SCREAMING_SNAKE_CASE      :     List[str] ,     __SCREAMING_SNAKE_CASE      :     List[Any] ,     __SCREAMING_SNAKE_CASE      :     int=None ,     )      ->							Any:
           __UpperCAmelCase     				=len(__SCREAMING_SNAKE_CASE      ) if isinstance(__SCREAMING_SNAKE_CASE ,     __SCREAMING_SNAKE_CASE      ) else 1
           # get prompt text embeddings
           __UpperCAmelCase     				=self.tokenizer(
               __SCREAMING_SNAKE_CASE ,     padding="""max_length""" ,     truncation=__SCREAMING_SNAKE_CASE ,     max_length=77 ,     return_attention_mask=__SCREAMING_SNAKE_CASE ,     add_special_tokens=__SCREAMING_SNAKE_CASE ,     return_tensors="""pt""" ,     )

           __UpperCAmelCase     				=text_inputs.input_ids
           __UpperCAmelCase     				=self.tokenizer(__SCREAMING_SNAKE_CASE ,     padding="""longest""" ,     return_tensors="""pt"""      ).input_ids

           if untruncated_ids.shape[-1] >= text_input_ids.shape[-1] and not torch.equal(__SCREAMING_SNAKE_CASE ,     __SCREAMING_SNAKE_CASE      ):
                  __UpperCAmelCase     				=self.tokenizer.batch_decode(untruncated_ids[:, self.tokenizer.model_max_length - 1 : -1]      )
                  logger.warning(
                      """The following part of your input was truncated because CLIP can only handle sequences up to"""
                      f''' {self.tokenizer.model_max_length} tokens: {removed_text}'''      )

           __UpperCAmelCase     				=text_input_ids.to(__SCREAMING_SNAKE_CASE      )
           __UpperCAmelCase     				=text_inputs.attention_mask.to(__SCREAMING_SNAKE_CASE      )

           __UpperCAmelCase				, __UpperCAmelCase     				=self.text_encoder(
               input_ids=__SCREAMING_SNAKE_CASE ,     attention_mask=__SCREAMING_SNAKE_CASE      )

           __UpperCAmelCase     				=prompt_embeds.repeat_interleave(__SCREAMING_SNAKE_CASE ,     dim=0      )
           __UpperCAmelCase     				=text_encoder_hidden_states.repeat_interleave(__SCREAMING_SNAKE_CASE ,     dim=0      )
           __UpperCAmelCase     				=text_mask.repeat_interleave(__SCREAMING_SNAKE_CASE ,     dim=0      )

           if do_classifier_free_guidance:
                  __UpperCAmelCase     				=42
                  if negative_prompt is None:
                         __UpperCAmelCase     				=[""""""] * batch_size
                  elif type(__SCREAMING_SNAKE_CASE      ) is not type(__SCREAMING_SNAKE_CASE      ):
                         raise TypeError(
                             f'''`negative_prompt` should be the same type to `prompt`, but got {type(__SCREAMING_SNAKE_CASE      )} !='''
                             f''' {type(__SCREAMING_SNAKE_CASE      )}.'''      )
                  elif isinstance(__SCREAMING_SNAKE_CASE ,     __SCREAMING_SNAKE_CASE      ):
                         __UpperCAmelCase     				=[negative_prompt]
                  elif batch_size != len(__SCREAMING_SNAKE_CASE      ):
                         raise ValueError(
                             f'''`negative_prompt`: {negative_prompt} has batch size {len(__SCREAMING_SNAKE_CASE      )}, but `prompt`:'''
                             f''' {prompt} has batch size {batch_size}. Please make sure that passed `negative_prompt` matches'''
                             """ the batch size of `prompt`."""      )
                  else:
                         __UpperCAmelCase     				=negative_prompt

                  __UpperCAmelCase     				=self.tokenizer(
                      __SCREAMING_SNAKE_CASE ,     padding="""max_length""" ,     max_length=77 ,     truncation=__SCREAMING_SNAKE_CASE ,     return_attention_mask=__SCREAMING_SNAKE_CASE ,     add_special_tokens=__SCREAMING_SNAKE_CASE ,     return_tensors="""pt""" ,     )
                  __UpperCAmelCase     				=uncond_input.input_ids.to(__SCREAMING_SNAKE_CASE      )
                  __UpperCAmelCase     				=uncond_input.attention_mask.to(__SCREAMING_SNAKE_CASE      )

                  __UpperCAmelCase				, __UpperCAmelCase     				=self.text_encoder(
                      input_ids=__SCREAMING_SNAKE_CASE ,     attention_mask=__SCREAMING_SNAKE_CASE      )

                  # duplicate unconditional embeddings for each generation per prompt, using mps friendly method

                  __UpperCAmelCase     				=negative_prompt_embeds.shape[1]
                  __UpperCAmelCase     				=negative_prompt_embeds.repeat(1 ,     __SCREAMING_SNAKE_CASE      )
                  __UpperCAmelCase     				=negative_prompt_embeds.view(batch_size * num_images_per_prompt ,     __SCREAMING_SNAKE_CASE      )

                  __UpperCAmelCase     				=uncond_text_encoder_hidden_states.shape[1]
                  __UpperCAmelCase     				=uncond_text_encoder_hidden_states.repeat(1 ,     __SCREAMING_SNAKE_CASE ,     1      )
                  __UpperCAmelCase     				=uncond_text_encoder_hidden_states.view(
                      batch_size * num_images_per_prompt ,     __SCREAMING_SNAKE_CASE ,     -1      )
                  __UpperCAmelCase     				=uncond_text_mask.repeat_interleave(__SCREAMING_SNAKE_CASE ,     dim=0      )

                  # done duplicates

                  # For classifier free guidance, we need to do two forward passes.
                  # Here we concatenate the unconditional and text embeddings into a single batch
                  # to avoid doing two forward passes
                  __UpperCAmelCase     				=torch.cat([negative_prompt_embeds, prompt_embeds]      )
                  __UpperCAmelCase     				=torch.cat([uncond_text_encoder_hidden_states, text_encoder_hidden_states]      )

                  __UpperCAmelCase     				=torch.cat([uncond_text_mask, text_mask]      )

           return prompt_embeds, text_encoder_hidden_states, text_mask



    def _a							(				self      :     Dict ,     __SCREAMING_SNAKE_CASE      :     List[str]=0      )      ->							Dict:
           if is_accelerate_available():
                  from accelerate import cpu_offload
           else:
                  raise ImportError("""Please install accelerate via `pip install accelerate`"""      )

           __UpperCAmelCase     				=torch.device(f'''cuda:{gpu_id}'''      )

           __UpperCAmelCase     				=[
               self.unet,
               self.text_encoder,
               self.movq,
           ]
           for cpu_offloaded_model in models:
                  if cpu_offloaded_model is not None:
                         cpu_offload(__SCREAMING_SNAKE_CASE ,     __SCREAMING_SNAKE_CASE      )



    def _a							(				self      :     Optional[Any] ,     __SCREAMING_SNAKE_CASE      :     Dict=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     				=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     				=None
           for cpu_offloaded_model in [self.text_encoder, self.unet, self.movq]:
                  __UpperCAmelCase				, __UpperCAmelCase     				=cpu_offload_with_hook(__SCREAMING_SNAKE_CASE ,     __SCREAMING_SNAKE_CASE ,     prev_module_hook=__SCREAMING_SNAKE_CASE      )

           if self.safety_checker is not None:
                  __UpperCAmelCase				, __UpperCAmelCase     				=cpu_offload_with_hook(self.safety_checker ,     __SCREAMING_SNAKE_CASE ,     prev_module_hook=__SCREAMING_SNAKE_CASE      )

           # We'll offload the last model manually.
           __UpperCAmelCase     				=hook



    @property
    # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline._execution_device
    def _a							(				self      :     Union[str, Any]      )      ->							Dict:
           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      :     Any ,     __SCREAMING_SNAKE_CASE      :     Union[str, List[str]] ,     __SCREAMING_SNAKE_CASE      :     Union[torch.FloatTensor, List[torch.FloatTensor]] ,     __SCREAMING_SNAKE_CASE      :     Union[torch.FloatTensor, List[torch.FloatTensor]] ,     __SCREAMING_SNAKE_CASE      :     Optional[Union[str, List[str]]] = None ,     __SCREAMING_SNAKE_CASE      :     int = 512 ,     __SCREAMING_SNAKE_CASE      :     int = 512 ,     __SCREAMING_SNAKE_CASE      :     int = 100 ,     __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]:

           if isinstance(__SCREAMING_SNAKE_CASE ,     __SCREAMING_SNAKE_CASE      ):
                  __UpperCAmelCase     				=1
           elif isinstance(__SCREAMING_SNAKE_CASE ,     __SCREAMING_SNAKE_CASE      ):
                  __UpperCAmelCase     				=len(__SCREAMING_SNAKE_CASE      )
           else:
                  raise ValueError(f'''`prompt` has to be of type `str` or `list` but is {type(__SCREAMING_SNAKE_CASE      )}'''      )

           __UpperCAmelCase     				=self._execution_device

           __UpperCAmelCase     				=batch_size * num_images_per_prompt
           __UpperCAmelCase     				=guidance_scale > 1.0

           __UpperCAmelCase				, __UpperCAmelCase				, __UpperCAmelCase     				=self._encode_prompt(
               __SCREAMING_SNAKE_CASE ,     __SCREAMING_SNAKE_CASE ,     __SCREAMING_SNAKE_CASE ,     __SCREAMING_SNAKE_CASE ,     __SCREAMING_SNAKE_CASE      )

           if isinstance(__SCREAMING_SNAKE_CASE ,     __SCREAMING_SNAKE_CASE      ):
                  __UpperCAmelCase     				=torch.cat(__SCREAMING_SNAKE_CASE ,     dim=0      )
           if isinstance(__SCREAMING_SNAKE_CASE ,     __SCREAMING_SNAKE_CASE      ):
                  __UpperCAmelCase     				=torch.cat(__SCREAMING_SNAKE_CASE ,     dim=0      )

           if do_classifier_free_guidance:
                  __UpperCAmelCase     				=image_embeds.repeat_interleave(__SCREAMING_SNAKE_CASE ,     dim=0      )
                  __UpperCAmelCase     				=negative_image_embeds.repeat_interleave(__SCREAMING_SNAKE_CASE ,     dim=0      )

           __UpperCAmelCase     				=torch.cat([negative_image_embeds, image_embeds] ,     dim=0      ).to(
               dtype=prompt_embeds.dtype ,     device=__SCREAMING_SNAKE_CASE      )

           self.scheduler.set_timesteps(__SCREAMING_SNAKE_CASE ,     device=__SCREAMING_SNAKE_CASE      )
           __UpperCAmelCase     				=self.scheduler.timesteps

           __UpperCAmelCase     				=self.unet.config.in_channels

           __UpperCAmelCase				, __UpperCAmelCase     				=get_new_h_w(__SCREAMING_SNAKE_CASE ,     __SCREAMING_SNAKE_CASE ,     self.movq_scale_factor      )

           # create initial latent
           __UpperCAmelCase     				=self.prepare_latents(
               (batch_size, num_channels_latents, height, width) ,     text_encoder_hidden_states.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
                  __UpperCAmelCase     				=torch.cat([latents] * 2      ) if do_classifier_free_guidance else latents

                  __UpperCAmelCase     				={"""text_embeds""": prompt_embeds, """image_embeds""": image_embeds}
                  __UpperCAmelCase     				=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     				=noise_pred.split(latents.shape[1] ,     dim=1      )
                         __UpperCAmelCase				, __UpperCAmelCase     				=noise_pred.chunk(2      )
                         __UpperCAmelCase				, __UpperCAmelCase     				=variance_pred.chunk(2      )
                         __UpperCAmelCase     				=noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond)
                         __UpperCAmelCase     				=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     				=noise_pred.split(latents.shape[1] ,     dim=1      )

                  # compute the previous noisy sample x_t -> x_t-1
                  __UpperCAmelCase     				=self.scheduler.step(
                      __SCREAMING_SNAKE_CASE ,     __SCREAMING_SNAKE_CASE ,     __SCREAMING_SNAKE_CASE ,     generator=__SCREAMING_SNAKE_CASE ,     ).prev_sample
           # post-processing
           __UpperCAmelCase     				=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     				=image * 0.5 + 0.5
                  __UpperCAmelCase     				=image.clamp(0 ,     1      )
                  __UpperCAmelCase     				=image.cpu().permute(0 ,     2 ,     3 ,     1      ).float().numpy()

           if output_type == "pil":
                  __UpperCAmelCase     				=self.numpy_to_pil(__SCREAMING_SNAKE_CASE      )

           if not return_dict:
                  return (image,)

           return ImagePipelineOutput(images=__SCREAMING_SNAKE_CASE      )








import copy
from dataclasses import dataclass, field
from typing import ClassVar, Dict

from ..features import ClassLabel, Features, Image
from .base import TaskTemplate


@dataclass(frozen=A     )
class 		_SCREAMING_SNAKE_CASE     (						A     ):
     __SCREAMING_SNAKE_CASE     					=    field(default='''image-classification'''							,	metadata={'''include_in_asdict_even_if_is_default''': True}     )
     __SCREAMING_SNAKE_CASE     					=    Features({'''image''': Image()}     )
     __SCREAMING_SNAKE_CASE     					=    Features({'''labels''': ClassLabel}     )
     __SCREAMING_SNAKE_CASE     					=    "image"
     __SCREAMING_SNAKE_CASE     					=    "labels"


     def        __snake_case(  self							,					A_    ):
       if self.label_column not in features:
         raise ValueError(F'''Column {self.label_column} is not present in features.'''    )
       if not isinstance(features[self.label_column]							,					A_    ):
         raise ValueError(F'''Column {self.label_column} is not a ClassLabel.'''    )
       _UpperCAmelCase		:							Tuple							=     copy.deepcopy(self    )
       _UpperCAmelCase		:							str							=     self.label_schema.copy()
       _UpperCAmelCase		:							Optional[Any]							=     features[self.label_column]
       _UpperCAmelCase		:							int							=     label_schema
       return task_template




     @property
     def        __snake_case(  self    ):
       return {
           self.image_column: "image",
           self.label_column: "labels",
       }





