



"""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


_lowerCAmelCase    :Optional[int]     =					logging.get_logger(__name__)  # pylint: disable=invalid-name

_lowerCAmelCase    :Union[str, Any]     =					'\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  lowerCamelCase_						(UpperCamelCase__		:		Dict     ,       UpperCamelCase__		:		List[Any]     ,       UpperCamelCase__		:		Union[str, Any]=8							):
    _UpperCAmelCase							:  Optional[int]											=    h // scale_factor**2
    if h % scale_factor**2 != 0:
        new_h += 1
    _UpperCAmelCase							:  Optional[Any]											=    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       ,	A       ,	A       ,	A       ,	A       ,	A       ,	)							->					Optional[Any]:
         super().__init__()

         self.register_modules(
             text_encoder=A       ,	tokenizer=A       ,	unet=A       ,	scheduler=A       ,	movq=A       ,	)
         _UpperCAmelCase							:  Tuple											=    2 ** (len(self.movq.config.block_out_channels		) - 1)





     def  __lowerCAmelCase	(					self       ,	A       ,	A       ,	A       ,	A       ,	A       ,	A		)							->					Dict:
         if latents is None:
             _UpperCAmelCase							:  Dict											=    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}'		)
             _UpperCAmelCase							:  Dict											=    latents.to(A		)

         _UpperCAmelCase							:  Union[str, Any]											=    latents * scheduler.init_noise_sigma
         return latents





     def  __lowerCAmelCase	(					self       ,	A       ,	A       ,	A       ,	A       ,	A=None       ,	)							->					List[Any]:
         _UpperCAmelCase							:  Optional[Any]											=    len(A		) if isinstance(A       ,	A		) else 1
         # get prompt text embeddings
         _UpperCAmelCase							:  List[str]											=    self.tokenizer(
             A       ,	padding='''max_length'''       ,	truncation=A       ,	max_length=7_7       ,	return_attention_mask=A       ,	add_special_tokens=A       ,	return_tensors='''pt'''       ,	)

         _UpperCAmelCase							:  str											=    text_inputs.input_ids
         _UpperCAmelCase							:  Any											=    self.tokenizer(A       ,	padding='''longest'''       ,	return_tensors='''pt'''		).input_ids

         if untruncated_ids.shape[-1] >= text_input_ids.shape[-1] and not torch.equal(A       ,	A		):
             _UpperCAmelCase							:  Tuple											=    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							:  List[str]											=    text_input_ids.to(A		)
         _UpperCAmelCase							:  int											=    text_inputs.attention_mask.to(A		)

         _UpperCAmelCase							,				_UpperCAmelCase							:  Optional[Any]											=    self.text_encoder(
             input_ids=A       ,	attention_mask=A		)

         _UpperCAmelCase							:  Tuple											=    prompt_embeds.repeat_interleave(A       ,	dim=0		)
         _UpperCAmelCase							:  Union[str, Any]											=    text_encoder_hidden_states.repeat_interleave(A       ,	dim=0		)
         _UpperCAmelCase							:  Any											=    text_mask.repeat_interleave(A       ,	dim=0		)

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

             _UpperCAmelCase							:  List[Any]											=    self.tokenizer(
                 A       ,	padding='''max_length'''       ,	max_length=7_7       ,	truncation=A       ,	return_attention_mask=A       ,	add_special_tokens=A       ,	return_tensors='''pt'''       ,	)
             _UpperCAmelCase							:  Tuple											=    uncond_input.input_ids.to(A		)
             _UpperCAmelCase							:  int											=    uncond_input.attention_mask.to(A		)

             _UpperCAmelCase							,				_UpperCAmelCase							:  List[Any]											=    self.text_encoder(
                 input_ids=A       ,	attention_mask=A		)

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

             _UpperCAmelCase							:  Any											=    negative_prompt_embeds.shape[1]
             _UpperCAmelCase							:  Optional[int]											=    negative_prompt_embeds.repeat(1       ,	A		)
             _UpperCAmelCase							:  Any											=    negative_prompt_embeds.view(batch_size * num_images_per_prompt       ,	A		)

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

             _UpperCAmelCase							:  Tuple											=    torch.cat([uncond_text_mask, text_mask]		)

         return prompt_embeds, text_encoder_hidden_states, text_mask





     def  __lowerCAmelCase	(					self       ,	A=0		)							->					List[Any]:
         if is_accelerate_available():
             from accelerate import cpu_offload
         else:
             raise ImportError('''Please install accelerate via `pip install accelerate`'''		)

         _UpperCAmelCase							:  List[str]											=    torch.device(f'cuda:{gpu_id}'		)

         _UpperCAmelCase							:  Optional[int]											=    [
             self.unet,
             self.text_encoder,
             self.movq,
         ]
         for cpu_offloaded_model in models:
             if cpu_offloaded_model is not None:
                 cpu_offload(A       ,	A		)





     def  __lowerCAmelCase	(					self       ,	A=0		)							->					Tuple:
         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							:  Dict											=    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)

         _UpperCAmelCase							:  Optional[int]											=    None
         for cpu_offloaded_model in [self.text_encoder, self.unet, self.movq]:
             _UpperCAmelCase							,				_UpperCAmelCase							:  Dict											=    cpu_offload_with_hook(A       ,	A       ,	prev_module_hook=A		)

         if self.safety_checker is not None:
             _UpperCAmelCase							,				_UpperCAmelCase							:  int											=    cpu_offload_with_hook(self.safety_checker       ,	A       ,	prev_module_hook=A		)

         # We'll offload the last model manually.
         _UpperCAmelCase							:  int											=    hook





     @property
     # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline._execution_device
     def  __lowerCAmelCase	(					self		)							->					List[Any]:
         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       ,	A       ,	A       ,	A       ,	A = None       ,	A = 5_1_2       ,	A = 5_1_2       ,	A = 1_0_0       ,	A = 4.0       ,	A = 1       ,	A = None       ,	A = None       ,	A = "pil"       ,	A = True       ,	)							->					List[str]:

         if isinstance(A       ,	A		):
             _UpperCAmelCase							:  Optional[int]											=    1
         elif isinstance(A       ,	A		):
             _UpperCAmelCase							:  Tuple											=    len(A		)
         else:
             raise ValueError(f'`prompt` has to be of type `str` or `list` but is {type(A		)}'		)

         _UpperCAmelCase							:  str											=    self._execution_device

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

         _UpperCAmelCase							,				_UpperCAmelCase							,				_UpperCAmelCase							:  int											=    self._encode_prompt(
             A       ,	A       ,	A       ,	A       ,	A		)

         if isinstance(A       ,	A		):
             _UpperCAmelCase							:  Any											=    torch.cat(A       ,	dim=0		)
         if isinstance(A       ,	A		):
             _UpperCAmelCase							:  Optional[int]											=    torch.cat(A       ,	dim=0		)

         if do_classifier_free_guidance:
             _UpperCAmelCase							:  int											=    image_embeds.repeat_interleave(A       ,	dim=0		)
             _UpperCAmelCase							:  List[str]											=    negative_image_embeds.repeat_interleave(A       ,	dim=0		)

         _UpperCAmelCase							:  str											=    torch.cat([negative_image_embeds, image_embeds]       ,	dim=0		).to(
             dtype=prompt_embeds.dtype       ,	device=A		)

         self.scheduler.set_timesteps(A       ,	device=A		)
         _UpperCAmelCase							:  int											=    self.scheduler.timesteps

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

         _UpperCAmelCase							,				_UpperCAmelCase							:  List[Any]											=    get_new_h_w(A       ,	A       ,	self.movq_scale_factor		)

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

             _UpperCAmelCase							:  List[str]											=    {'''text_embeds''': prompt_embeds, '''image_embeds''': image_embeds}
             _UpperCAmelCase							:  List[str]											=    self.unet(
                 sample=A       ,	timestep=A       ,	encoder_hidden_states=A       ,	added_cond_kwargs=A       ,	return_dict=A       ,	)[0]

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

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

         if output_type == "pil":
             _UpperCAmelCase							:  List[str]											=    self.numpy_to_pil(A		)

         if not return_dict:
             return (image,)

         return ImagePipelineOutput(images=A		)







"""simple docstring"""
def  lowerCamelCase_						(UpperCamelCase__		:		int							):
    if not isinstance(UpperCamelCase__     ,       UpperCamelCase__							) or number < 0:
        raise ValueError('''Input must be a non-negative integer'''							)

    _UpperCAmelCase							:  str											=    0
    while number:
        # This way we arrive at next set bit (next 1) instead of looping
        # through each bit and checking for 1s hence the
        # loop won't run 32 times it will only run the number of `1` times
        number &= number - 1
        count += 1
    return count


if __name__ == "__main__":
  import doctest

  doctest.testmod()


