


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


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

__UpperCamelCase    :						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      :			Tuple						,	UpperCAmelCase      :			Optional[int]						,	UpperCAmelCase      :			Any=8      ):
      """simple docstring"""




      __lowerCamelCase     :      Union[str, Any]       						=						h // scale_factor**2
      if h % scale_factor**2 != 0:
            new_h += 1
      __lowerCamelCase     :      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 (						__snake_case						):

 '''simple docstring'''







 def __init__(  self  : Dict       ,  _lowerCamelCase  : str       ,  _lowerCamelCase  : int       ,  _lowerCamelCase  : str       ,  _lowerCamelCase  : List[str]       ,  _lowerCamelCase  : Tuple       ,  ):






       '''simple docstring'''



       super().__init__()

       self.register_modules(
           text_encoder=A_       ,  tokenizer=A_       ,  unet=A_       ,  scheduler=A_       ,  movq=A_       ,  )
       __lowerCamelCase     :      List[Any]       						=						2 ** (len(self.movq.config.block_out_channels							) - 1)

 def   _snake_case    (  self  : Optional[Any]       ,  _lowerCamelCase  : List[Any]       ,  _lowerCamelCase  : List[str]       ,  _lowerCamelCase  : Any       ,  _lowerCamelCase  : Optional[int]       ,  _lowerCamelCase  : str       ,  _lowerCamelCase  : List[str]							):






       '''simple docstring'''



       if latents is None:
             __lowerCamelCase     :      int       						=						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     :      Dict       						=						latents * scheduler.init_noise_sigma
       return latents

 def   _snake_case    (  self  : List[str]       ,  _lowerCamelCase  : int       ,  _lowerCamelCase  : Dict       ,  _lowerCamelCase  : Optional[Any]       ,  _lowerCamelCase  : Union[str, Any]       ,  _lowerCamelCase  : Union[str, Any]=None       ,  ):






       '''simple docstring'''



       __lowerCamelCase     :      int       						=						len(A_							) if isinstance(A_       ,  A_							) else 1
       # get prompt text embeddings
       __lowerCamelCase     :      Any       						=						self.tokenizer(
           A_       ,  padding="""max_length"""       ,  truncation=A_       ,  max_length=7_7       ,  return_attention_mask=A_       ,  add_special_tokens=A_       ,  return_tensors="""pt"""       ,  )

       __lowerCamelCase     :      List[Any]       						=						text_inputs.input_ids
       __lowerCamelCase     :      List[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_							):
             __lowerCamelCase     :      Union[str, Any]       						=						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}"""							)

       __lowerCamelCase     :      str       						=						text_input_ids.to(A_							)
       __lowerCamelCase     :      int       						=						text_inputs.attention_mask.to(A_							)

       __lowerCamelCase     :      Optional[int]       						=						self.text_encoder(
           input_ids=A_       ,  attention_mask=A_							)

       __lowerCamelCase     :      Optional[Any]       						=						prompt_embeds.repeat_interleave(A_       ,  dim=0							)
       __lowerCamelCase     :      Any       						=						text_encoder_hidden_states.repeat_interleave(A_       ,  dim=0							)
       __lowerCamelCase     :      int       						=						text_mask.repeat_interleave(A_       ,  dim=0							)

       if do_classifier_free_guidance:
             __lowerCamelCase     :      List[str]
             if negative_prompt is None:
                   __lowerCamelCase     :      Optional[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_							):
                   __lowerCamelCase     :      Union[str, Any]       						=						[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:
                   __lowerCamelCase     :      Optional[int]       						=						negative_prompt

             __lowerCamelCase     :      Tuple       						=						self.tokenizer(
                 A_       ,  padding="""max_length"""       ,  max_length=7_7       ,  truncation=A_       ,  return_attention_mask=A_       ,  add_special_tokens=A_       ,  return_tensors="""pt"""       ,  )
             __lowerCamelCase     :      Optional[Any]       						=						uncond_input.input_ids.to(A_							)
             __lowerCamelCase     :      Tuple       						=						uncond_input.attention_mask.to(A_							)

             __lowerCamelCase     :      Optional[int]       						=						self.text_encoder(
                 input_ids=A_       ,  attention_mask=A_							)

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

             __lowerCamelCase     :      int       						=						negative_prompt_embeds.shape[1]
             __lowerCamelCase     :      Any       						=						negative_prompt_embeds.repeat(1       ,  A_							)
             __lowerCamelCase     :      List[Any]       						=						negative_prompt_embeds.view(batch_size * num_images_per_prompt       ,  A_							)

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

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

       return prompt_embeds, text_encoder_hidden_states, text_mask

 def   _snake_case    (  self  : Optional[int]       ,  _lowerCamelCase  : Tuple=0							):






       '''simple docstring'''



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

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

       __lowerCamelCase     :      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   _snake_case    (  self  : Tuple       ,  _lowerCamelCase  : List[str]=0							):






       '''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     :      Union[str, 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     :      Tuple       						=						None
       for cpu_offloaded_model in [self.text_encoder, self.unet, self.movq]:
             __lowerCamelCase     :      Optional[int]       						=						cpu_offload_with_hook(A_       ,  A_       ,  prev_module_hook=A_							)

       if self.safety_checker is not None:
             __lowerCamelCase     :      Any       						=						cpu_offload_with_hook(self.safety_checker       ,  A_       ,  prev_module_hook=A_							)

       # We'll offload the last model manually.
       __lowerCamelCase     :      str       						=						hook

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






       '''simple docstring'''



       if isinstance(A_       ,  A_							):
             __lowerCamelCase     :      Optional[Any]       						=						1
       elif isinstance(A_       ,  A_							):
             __lowerCamelCase     :      List[Any]       						=						len(A_							)
       else:
             raise ValueError(F"""`prompt` has to be of type `str` or `list` but is {type(A_							)}"""							)

       __lowerCamelCase     :      str       						=						self._execution_device

       __lowerCamelCase     :      Dict       						=						batch_size * num_images_per_prompt
       __lowerCamelCase     :      Optional[int]       						=						guidance_scale > 1.0

       __lowerCamelCase     :      Union[str, Any]       						=						self._encode_prompt(
           A_       ,  A_       ,  A_       ,  A_       ,  A_							)

       if isinstance(A_       ,  A_							):
             __lowerCamelCase     :      Any       						=						torch.cat(A_       ,  dim=0							)
       if isinstance(A_       ,  A_							):
             __lowerCamelCase     :      int       						=						torch.cat(A_       ,  dim=0							)

       if do_classifier_free_guidance:
             __lowerCamelCase     :      int       						=						image_embeds.repeat_interleave(A_       ,  dim=0							)
             __lowerCamelCase     :      Optional[Any]       						=						negative_image_embeds.repeat_interleave(A_       ,  dim=0							)

       __lowerCamelCase     :      List[Any]       						=						torch.cat([negative_image_embeds, image_embeds]       ,  dim=0							).to(
           dtype=prompt_embeds.dtype       ,  device=A_							)

       self.scheduler.set_timesteps(A_       ,  device=A_							)
       __lowerCamelCase     :      Any       						=						self.scheduler.timesteps

       __lowerCamelCase     :      Tuple       						=						self.unet.config.in_channels

       __lowerCamelCase     :      Union[str, Any]       						=						get_new_h_w(A_       ,  A_       ,  self.movq_scale_factor							)

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

             __lowerCamelCase     :      Dict       						=						{"text_embeds": prompt_embeds, "image_embeds": image_embeds}
             __lowerCamelCase     :      int       						=						self.unet(
                 sample=A_       ,  timestep=A_       ,  encoder_hidden_states=A_       ,  added_cond_kwargs=A_       ,  return_dict=A_       ,  )[0]

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

             # compute the previous noisy sample x_t -> x_t-1
             __lowerCamelCase     :      Optional[int]       						=						self.scheduler.step(
                 A_       ,  A_       ,  A_       ,  generator=A_       ,  ).prev_sample
       # post-processing
       __lowerCamelCase     :      Tuple       						=						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     :      Tuple       						=						image * 0.5 + 0.5
             __lowerCamelCase     :      Union[str, Any]       						=						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_							)







# Copyright 2023 The HuggingFace Team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
from typing import TYPE_CHECKING

from ...utils import OptionalDependencyNotAvailable, _LazyModule, is_torch_available


__UpperCamelCase    :						Any     =							{
    'configuration_mgp_str': ['MGP_STR_PRETRAINED_CONFIG_ARCHIVE_MAP', 'MgpstrConfig'],
    'processing_mgp_str': ['MgpstrProcessor'],
    'tokenization_mgp_str': ['MgpstrTokenizer'],
}

try:
    if not is_torch_available():
        raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
    pass
else:
    __UpperCamelCase    :						Tuple     =							[
        'MGP_STR_PRETRAINED_MODEL_ARCHIVE_LIST',
        'MgpstrModel',
        'MgpstrPreTrainedModel',
        'MgpstrForSceneTextRecognition',
    ]

if TYPE_CHECKING:
    from .configuration_mgp_str import MGP_STR_PRETRAINED_CONFIG_ARCHIVE_MAP, MgpstrConfig
    from .processing_mgp_str import MgpstrProcessor
    from .tokenization_mgp_str import MgpstrTokenizer

    try:
        if not is_torch_available():
            raise OptionalDependencyNotAvailable()
    except OptionalDependencyNotAvailable:
        pass
    else:
        from .modeling_mgp_str import (
            MGP_STR_PRETRAINED_MODEL_ARCHIVE_LIST,
            MgpstrForSceneTextRecognition,
            MgpstrModel,
            MgpstrPreTrainedModel,
        )
else:
    import sys

    __UpperCamelCase    :						Optional[Any]     =							_LazyModule(__name__, globals()['__file__'], _import_structure, module_spec=__spec__)



