'''simple docstring'''

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: List[str]          =   logging.get_logger(__name__)  # pylint: disable=invalid-name

__a: List[Any]          =   """
    Examples:
        ```py
        >>> from diffusers import KandinskyPipeline, KandinskyPriorPipeline
        >>> import torch

        >>> pipe_prior = KandinskyPriorPipeline.from_pretrained(\"kandinsky-community/Kandinsky-2-1-prior\")
        >>> pipe_prior.to(\"cuda\")

        >>> prompt = \"red cat, 4k photo\"
        >>> out = pipe_prior(prompt)
        >>> image_emb = out.image_embeds
        >>> negative_image_emb = out.negative_image_embeds

        >>> pipe = KandinskyPipeline.from_pretrained(\"kandinsky-community/kandinsky-2-1\")
        >>> pipe.to(\"cuda\")

        >>> image = pipe(
        ...     prompt,
        ...     image_embeds=image_emb,
        ...     negative_image_embeds=negative_image_emb,
        ...     height=768,
        ...     width=768,
        ...     num_inference_steps=100,
        ... ).images

        >>> image[0].save(\"cat.png\")
        ```
"""
def 			__UpperCamelCase       (   UpperCAmelCase		,     UpperCAmelCase		,     UpperCAmelCase=8     ):
 lowercase__				:    str  =					h // scale_factor**2
 if h % scale_factor**2 != 0:
  new_h += 1
 lowercase__				:    Tuple  =					w // scale_factor**2
 if w % scale_factor**2 != 0:
  new_w += 1
 return new_h * scale_factor, new_w * scale_factor





class 			UpperCAmelCase  (						a__					):





       '''simple docstring'''



       def __init__(						self	,   __lowerCAmelCase	,   __lowerCAmelCase	,   __lowerCAmelCase	,   __lowerCAmelCase	,   __lowerCAmelCase	,   )  ->  Tuple:
        super().__init__()

        self.register_modules(
            text_encoder=__lowerCAmelCase	,   tokenizer=__lowerCAmelCase	,   unet=__lowerCAmelCase	,   scheduler=__lowerCAmelCase	,   movq=__lowerCAmelCase	,   )
        lowercase__				:    Any  =					2 ** (len(self.movq.config.block_out_channels ) - 1)

       def _lowerCAmelCase(						self	,   __lowerCAmelCase	,   __lowerCAmelCase	,   __lowerCAmelCase	,   __lowerCAmelCase	,   __lowerCAmelCase	,   __lowerCAmelCase )  ->  str:
        if latents is None:
         lowercase__				:    Tuple  =					randn_tensor(__lowerCAmelCase	,   generator=__lowerCAmelCase	,   device=__lowerCAmelCase	,   dtype=__lowerCAmelCase )
        else:
         if latents.shape != shape:
          raise ValueError(F"""Unexpected latents shape, got {latents.shape}, expected {shape}""" )
         lowercase__				:    Union[str, Any]  =					latents.to(__lowerCAmelCase )

        lowercase__				:    int  =					latents * scheduler.init_noise_sigma
        return latents

       def _lowerCAmelCase(						self	,   __lowerCAmelCase	,   __lowerCAmelCase	,   __lowerCAmelCase	,   __lowerCAmelCase	,   __lowerCAmelCase=None	,   )  ->  Union[str, Any]:
        lowercase__				:    Union[str, Any]  =					len(__lowerCAmelCase ) if isinstance(__lowerCAmelCase	,   __lowerCAmelCase ) else 1
        # get prompt text embeddings
        lowercase__				:    Optional[int]  =					self.tokenizer(
            __lowerCAmelCase	,   padding='''max_length'''	,   truncation=__lowerCAmelCase	,   max_length=77	,   return_attention_mask=__lowerCAmelCase	,   add_special_tokens=__lowerCAmelCase	,   return_tensors='''pt'''	,   )

        lowercase__				:    Optional[int]  =					text_inputs.input_ids
        lowercase__				:    Optional[int]  =					self.tokenizer(__lowerCAmelCase	,   padding='''longest'''	,   return_tensors='''pt''' ).input_ids

        if untruncated_ids.shape[-1] >= text_input_ids.shape[-1] and not torch.equal(__lowerCAmelCase	,   __lowerCAmelCase ):
         lowercase__				:    str  =					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}""" )

        lowercase__				:    Tuple  =					text_input_ids.to(__lowerCAmelCase )
        lowercase__				:    Any  =					text_inputs.attention_mask.to(__lowerCAmelCase )

        lowercase__ ,							lowercase__				:    Dict  =					self.text_encoder(
            input_ids=__lowerCAmelCase	,   attention_mask=__lowerCAmelCase )

        lowercase__				:    int  =					prompt_embeds.repeat_interleave(__lowerCAmelCase	,   dim=0 )
        lowercase__				:    Tuple  =					text_encoder_hidden_states.repeat_interleave(__lowerCAmelCase	,   dim=0 )
        lowercase__				:    Union[str, Any]  =					text_mask.repeat_interleave(__lowerCAmelCase	,   dim=0 )

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

         lowercase__				:    Dict  =					self.tokenizer(
             __lowerCAmelCase	,   padding='''max_length'''	,   max_length=77	,   truncation=__lowerCAmelCase	,   return_attention_mask=__lowerCAmelCase	,   add_special_tokens=__lowerCAmelCase	,   return_tensors='''pt'''	,   )
         lowercase__				:    Optional[Any]  =					uncond_input.input_ids.to(__lowerCAmelCase )
         lowercase__				:    Optional[int]  =					uncond_input.attention_mask.to(__lowerCAmelCase )

         lowercase__ ,							lowercase__				:    Dict  =					self.text_encoder(
             input_ids=__lowerCAmelCase	,   attention_mask=__lowerCAmelCase )

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

         lowercase__				:    Tuple  =					negative_prompt_embeds.shape[1]
         lowercase__				:    int  =					negative_prompt_embeds.repeat(1	,   __lowerCAmelCase )
         lowercase__				:    str  =					negative_prompt_embeds.view(batch_size * num_images_per_prompt	,   __lowerCAmelCase )

         lowercase__				:    Any  =					uncond_text_encoder_hidden_states.shape[1]
         lowercase__				:    str  =					uncond_text_encoder_hidden_states.repeat(1	,   __lowerCAmelCase	,   1 )
         lowercase__				:    List[Any]  =					uncond_text_encoder_hidden_states.view(
             batch_size * num_images_per_prompt	,   __lowerCAmelCase	,   -1 )
         lowercase__				:    List[Any]  =					uncond_text_mask.repeat_interleave(__lowerCAmelCase	,   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
         lowercase__				:    Optional[int]  =					torch.cat([negative_prompt_embeds, prompt_embeds] )
         lowercase__				:    Dict  =					torch.cat([uncond_text_encoder_hidden_states, text_encoder_hidden_states] )

         lowercase__				:    str  =					torch.cat([uncond_text_mask, text_mask] )

        return prompt_embeds, text_encoder_hidden_states, text_mask

       def _lowerCAmelCase(						self	,   __lowerCAmelCase=0 )  ->  Tuple:
        if is_accelerate_available():
         from accelerate import cpu_offload
        else:
         raise ImportError('''Please install accelerate via `pip install accelerate`''' )

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

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

       def _lowerCAmelCase(						self	,   __lowerCAmelCase=0 )  ->  List[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.''' )

        lowercase__				:    Union[str, Any]  =					torch.device(F"""cuda:{gpu_id}""" )

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

        lowercase__				:    int  =					None
        for cpu_offloaded_model in [self.text_encoder, self.unet, self.movq]:
         lowercase__ ,							lowercase__				:    Any  =					cpu_offload_with_hook(__lowerCAmelCase	,   __lowerCAmelCase	,   prev_module_hook=__lowerCAmelCase )

        if self.safety_checker is not None:
         lowercase__ ,							lowercase__				:    int  =					cpu_offload_with_hook(self.safety_checker	,   __lowerCAmelCase	,   prev_module_hook=__lowerCAmelCase )

        # We'll offload the last model manually.
        lowercase__				:    Tuple  =					hook

       @property
       # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline._execution_device
       def _lowerCAmelCase(						self )  ->  Tuple:
        if not hasattr(self.unet	,   '''_hf_hook''' ):
         return self.device
        for module in self.unet.modules():
         if (
             hasattr(__lowerCAmelCase	,   '''_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(__lowerCAmelCase )
       def __call__(						self	,   __lowerCAmelCase	,   __lowerCAmelCase	,   __lowerCAmelCase	,   __lowerCAmelCase = None	,   __lowerCAmelCase = 512	,   __lowerCAmelCase = 512	,   __lowerCAmelCase = 100	,   __lowerCAmelCase = 4.0	,   __lowerCAmelCase = 1	,   __lowerCAmelCase = None	,   __lowerCAmelCase = None	,   __lowerCAmelCase = "pil"	,   __lowerCAmelCase = True	,   )  ->  Union[str, Any]:

        if isinstance(__lowerCAmelCase	,   __lowerCAmelCase ):
         lowercase__				:    str  =					1
        elif isinstance(__lowerCAmelCase	,   __lowerCAmelCase ):
         lowercase__				:    Tuple  =					len(__lowerCAmelCase )
        else:
         raise ValueError(F"""`prompt` has to be of type `str` or `list` but is {type(__lowerCAmelCase )}""" )

        lowercase__				:    List[Any]  =					self._execution_device

        lowercase__				:    Optional[Any]  =					batch_size * num_images_per_prompt
        lowercase__				:    Optional[Any]  =					guidance_scale > 1.0

        lowercase__ ,							lowercase__ ,							lowercase__				:    List[Any]  =					self._encode_prompt(
            __lowerCAmelCase	,   __lowerCAmelCase	,   __lowerCAmelCase	,   __lowerCAmelCase	,   __lowerCAmelCase )

        if isinstance(__lowerCAmelCase	,   __lowerCAmelCase ):
         lowercase__				:    Union[str, Any]  =					torch.cat(__lowerCAmelCase	,   dim=0 )
        if isinstance(__lowerCAmelCase	,   __lowerCAmelCase ):
         lowercase__				:    Any  =					torch.cat(__lowerCAmelCase	,   dim=0 )

        if do_classifier_free_guidance:
         lowercase__				:    Optional[Any]  =					image_embeds.repeat_interleave(__lowerCAmelCase	,   dim=0 )
         lowercase__				:    Tuple  =					negative_image_embeds.repeat_interleave(__lowerCAmelCase	,   dim=0 )

        lowercase__				:    List[str]  =					torch.cat([negative_image_embeds, image_embeds]	,   dim=0 ).to(
            dtype=prompt_embeds.dtype	,   device=__lowerCAmelCase )

        self.scheduler.set_timesteps(__lowerCAmelCase	,   device=__lowerCAmelCase )
        lowercase__				:    str  =					self.scheduler.timesteps

        lowercase__				:    Dict  =					self.unet.config.in_channels

        lowercase__ ,							lowercase__				:    int  =					get_new_h_w(__lowerCAmelCase	,   __lowerCAmelCase	,   self.movq_scale_factor )

        # create initial latent
        lowercase__				:    Optional[Any]  =					self.prepare_latents(
            (batch_size, num_channels_latents, height, width)	,   text_encoder_hidden_states.dtype	,   __lowerCAmelCase	,   __lowerCAmelCase	,   __lowerCAmelCase	,   self.scheduler	,   )

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

         lowercase__				:    Any  =					{'''text_embeds''': prompt_embeds, '''image_embeds''': image_embeds}
         lowercase__				:    Any  =					self.unet(
             sample=__lowerCAmelCase	,   timestep=__lowerCAmelCase	,   encoder_hidden_states=__lowerCAmelCase	,   added_cond_kwargs=__lowerCAmelCase	,   return_dict=__lowerCAmelCase	,   )[0]

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

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

        if output_type == "pil":
         lowercase__				:    Optional[Any]  =					self.numpy_to_pil(__lowerCAmelCase )

        if not return_dict:
         return (image,)

        return ImagePipelineOutput(images=__lowerCAmelCase )


'''simple docstring'''

from math import pow
def 			__UpperCamelCase       (   UpperCAmelCase		,     UpperCAmelCase		,     UpperCAmelCase		,     UpperCAmelCase		,     UpperCAmelCase		,     ):
 if current_sum == needed_sum:
  # If the sum of the powers is equal to needed_sum, then we have a solution.
  solutions_count += 1
  return current_sum, solutions_count

 lowercase__				:    Optional[Any]  =					int(pow(UpperCAmelCase		,     UpperCAmelCase     )     )
 if current_sum + i_to_n <= needed_sum:
  # If the sum of the powers is less than needed_sum, then continue adding powers.
  current_sum += i_to_n
  lowercase__ ,							lowercase__				:    Dict  =					backtrack(
      UpperCAmelCase		,     UpperCAmelCase		,     current_number + 1		,     UpperCAmelCase		,     UpperCAmelCase     )
  current_sum -= i_to_n
 if i_to_n < needed_sum:
  # If the power of i is less than needed_sum, then try with the next power.
  lowercase__ ,							lowercase__				:    str  =					backtrack(
      UpperCAmelCase		,     UpperCAmelCase		,     current_number + 1		,     UpperCAmelCase		,     UpperCAmelCase     )
 return current_sum, solutions_count
def 			__UpperCamelCase       (   UpperCAmelCase		,     UpperCAmelCase     ):
 if not (1 <= needed_sum <= 1000 and 2 <= power <= 10):
  raise ValueError(
      '''Invalid input\n'''
      '''needed_sum must be between 1 and 1000, power between 2 and 10.'''     )

 return backtrack(UpperCAmelCase		,     UpperCAmelCase		,     1		,     0		,     0     )[1]  # Return the solutions_count


if __name__ == "__main__":
  import doctest

  doctest.testmod()

