



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


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

lowercase							       =			'''
    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_       (		lowercase__       ,	lowercase__       ,	lowercase__=8       ):


     '''simple docstring'''







     a_       		=h // scale_factor**2
     if h % scale_factor**2 != 0:
          new_h += 1
     a_       		=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_							,  ) ->      Dict:



      """simple docstring"""







      super().__init__()

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




 def       lowercase_  (   self							,  lowerCAmelCase_							,  lowerCAmelCase_							,  lowerCAmelCase_							,  lowerCAmelCase_							,  lowerCAmelCase_							,  lowerCAmelCase_) ->      Dict:



      """simple docstring"""







      if latents is None:
           a_       		=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}""")
           a_       		=latents.to(lowerCAmelCase_)

      a_       		=latents * scheduler.init_noise_sigma
      return latents




 def       lowercase_  (   self							,  lowerCAmelCase_							,  lowerCAmelCase_							,  lowerCAmelCase_							,  lowerCAmelCase_							,  lowerCAmelCase_=None							,  ) ->      List[str]:



      """simple docstring"""







      a_       		=len(lowerCAmelCase_) if isinstance(lowerCAmelCase_							,  lowerCAmelCase_) else 1
      # get prompt text embeddings
      a_       		=self.tokenizer(
          lowerCAmelCase_							,  padding="max_length"							,  truncation=lowerCAmelCase_							,  max_length=7_7							,  return_attention_mask=lowerCAmelCase_							,  add_special_tokens=lowerCAmelCase_							,  return_tensors="pt"							,  )

      a_       		=text_inputs.input_ids
      a_       		=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_):
           a_       		=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}""")

      a_       		=text_input_ids.to(lowerCAmelCase_)
      a_       		=text_inputs.attention_mask.to(lowerCAmelCase_)

      a_		,						a_       		=self.text_encoder(
          input_ids=lowerCAmelCase_							,  attention_mask=lowerCAmelCase_)

      a_       		=prompt_embeds.repeat_interleave(lowerCAmelCase_							,  dim=0)
      a_       		=text_encoder_hidden_states.repeat_interleave(lowerCAmelCase_							,  dim=0)
      a_       		=text_mask.repeat_interleave(lowerCAmelCase_							,  dim=0)

      if do_classifier_free_guidance:
           a_       		=42
           if negative_prompt is None:
                a_       		=[""] * 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_):
                a_       		=[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:
                a_       		=negative_prompt

           a_       		=self.tokenizer(
               lowerCAmelCase_							,  padding="max_length"							,  max_length=7_7							,  truncation=lowerCAmelCase_							,  return_attention_mask=lowerCAmelCase_							,  add_special_tokens=lowerCAmelCase_							,  return_tensors="pt"							,  )
           a_       		=uncond_input.input_ids.to(lowerCAmelCase_)
           a_       		=uncond_input.attention_mask.to(lowerCAmelCase_)

           a_		,						a_       		=self.text_encoder(
               input_ids=lowerCAmelCase_							,  attention_mask=lowerCAmelCase_)

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

           a_       		=negative_prompt_embeds.shape[1]
           a_       		=negative_prompt_embeds.repeat(1							,  lowerCAmelCase_)
           a_       		=negative_prompt_embeds.view(batch_size * num_images_per_prompt							,  lowerCAmelCase_)

           a_       		=uncond_text_encoder_hidden_states.shape[1]
           a_       		=uncond_text_encoder_hidden_states.repeat(1							,  lowerCAmelCase_							,  1)
           a_       		=uncond_text_encoder_hidden_states.view(
               batch_size * num_images_per_prompt							,  lowerCAmelCase_							,  -1)
           a_       		=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
           a_       		=torch.cat([negative_prompt_embeds, prompt_embeds])
           a_       		=torch.cat([uncond_text_encoder_hidden_states, text_encoder_hidden_states])

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

      return prompt_embeds, text_encoder_hidden_states, text_mask




 def       lowercase_  (   self							,  lowerCAmelCase_=0) ->      List[Any]:



      """simple docstring"""







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

      a_       		=torch.device(f"""cuda:{gpu_id}""")

      a_       		=[
          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       lowercase_  (   self							,  lowerCAmelCase_=0) ->      Union[str, 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.")

      a_       		=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)

      a_       		=None
      for cpu_offloaded_model in [self.text_encoder, self.unet, self.movq]:
           a_		,						a_       		=cpu_offload_with_hook(lowerCAmelCase_							,  lowerCAmelCase_							,  prev_module_hook=lowerCAmelCase_)

      if self.safety_checker is not None:
           a_		,						a_       		=cpu_offload_with_hook(self.safety_checker							,  lowerCAmelCase_							,  prev_module_hook=lowerCAmelCase_)

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




 @property
 # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline._execution_device
 def       lowercase_  (   self) ->      Optional[Any]:



      """simple docstring"""







      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_ = 5_1_2							,  lowerCAmelCase_ = 5_1_2							,  lowerCAmelCase_ = 1_0_0							,  lowerCAmelCase_ = 4.0							,  lowerCAmelCase_ = 1							,  lowerCAmelCase_ = None							,  lowerCAmelCase_ = None							,  lowerCAmelCase_ = "pil"							,  lowerCAmelCase_ = True							,  ) ->      List[str]:



      """simple docstring"""







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

      a_       		=self._execution_device

      a_       		=batch_size * num_images_per_prompt
      a_       		=guidance_scale > 1.0

      a_		,						a_		,						a_       		=self._encode_prompt(
          lowerCAmelCase_							,  lowerCAmelCase_							,  lowerCAmelCase_							,  lowerCAmelCase_							,  lowerCAmelCase_)

      if isinstance(lowerCAmelCase_							,  lowerCAmelCase_):
           a_       		=torch.cat(lowerCAmelCase_							,  dim=0)
      if isinstance(lowerCAmelCase_							,  lowerCAmelCase_):
           a_       		=torch.cat(lowerCAmelCase_							,  dim=0)

      if do_classifier_free_guidance:
           a_       		=image_embeds.repeat_interleave(lowerCAmelCase_							,  dim=0)
           a_       		=negative_image_embeds.repeat_interleave(lowerCAmelCase_							,  dim=0)

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

      self.scheduler.set_timesteps(lowerCAmelCase_							,  device=lowerCAmelCase_)
      a_       		=self.scheduler.timesteps

      a_       		=self.unet.config.in_channels

      a_		,						a_       		=get_new_h_w(lowerCAmelCase_							,  lowerCAmelCase_							,  self.movq_scale_factor)

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

           a_       		={"text_embeds": prompt_embeds, "image_embeds": image_embeds}
           a_       		=self.unet(
               sample=lowerCAmelCase_							,  timestep=lowerCAmelCase_							,  encoder_hidden_states=lowerCAmelCase_							,  added_cond_kwargs=lowerCAmelCase_							,  return_dict=lowerCAmelCase_							,  )[0]

           if do_classifier_free_guidance:
                a_		,						a_       		=noise_pred.split(latents.shape[1]							,  dim=1)
                a_		,						a_       		=noise_pred.chunk(2)
                a_		,						a_       		=variance_pred.chunk(2)
                a_       		=noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond)
                a_       		=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_       		=noise_pred.split(latents.shape[1]							,  dim=1)

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

      if output_type == "pil":
           a_       		=self.numpy_to_pil(lowerCAmelCase_)

      if not return_dict:
           return (image,)

      return ImagePipelineOutput(images=lowerCAmelCase_)







'''simple docstring'''


import torch

from diffusers import StableDiffusionPipeline


lowercase							       =			'''path-to-your-trained-model'''
lowercase							       =			StableDiffusionPipeline.from_pretrained(model_id, torch_dtype=torch.floataa).to('''cuda''')

lowercase							       =			'''A photo of sks dog in a bucket'''
lowercase							       =			pipe(prompt, num_inference_steps=50, guidance_scale=7.5).images[0]

image.save('''dog-bucket.png''')


