from typing import List, Optional

import torch

from diffusers.loaders import TextualInversionLoaderMixin, LoraLoaderMixin

from . import USE_PEFT_BACKEND, deprecate, logging, scale_lora_layers, unscale_lora_layers

from ..models.lora import adjust_lora_scale_text_encoder

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

# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.encode_prompt

def encode_prompt(

pipeline: 'DiffusionPipeline',

prompt,

device,

num_images_per_prompt,

do_classifier_free_guidance,

negative_prompt=None,

prompt_embeds: Optional[torch.FloatTensor] = None,

negative_prompt_embeds: Optional[torch.FloatTensor] = None,

lora_scale: Optional[float] = None,

clip_skip: Optional[int] = None,

):

r"""

Encodes the prompt into text encoder hidden states.

Args:

prompt (`str` or `List[str]`, *optional*):

prompt to be encoded

device: (`torch.device`):

torch device

num_images_per_prompt (`int`):

number of images that should be generated per prompt

do_classifier_free_guidance (`bool`):

whether to use classifier free guidance or not

negative_prompt (`str` or `List[str]`, *optional*):

The prompt or prompts not to guide the image generation. If not defined, one has to pass

`negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is

less than `1`).

prompt_embeds (`torch.FloatTensor`, *optional*):

Pre-generated text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. If not

provided, text embeddings will be generated from `prompt` input argument.

negative_prompt_embeds (`torch.FloatTensor`, *optional*):

Pre-generated negative text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt

weighting. If not provided, negative_prompt_embeds will be generated from `negative_prompt` input

argument.

lora_scale (`float`, *optional*):

A LoRA scale that will be applied to all LoRA layers of the text encoder if LoRA layers are loaded.

clip_skip (`int`, *optional*):

Number of layers to be skipped from CLIP while computing the prompt embeddings. A value of 1 means that

the output of the pre-final layer will be used for computing the prompt embeddings.

"""

# set lora scale so that monkey patched LoRA

# function of text encoder can correctly access it

if lora_scale is not None and isinstance(pipeline, LoraLoaderMixin):

pipeline._lora_scale = lora_scale

# dynamically adjust the LoRA scale

if not USE_PEFT_BACKEND:

adjust_lora_scale_text_encoder(pipeline.text_encoder, lora_scale)

else:

scale_lora_layers(pipeline.text_encoder, lora_scale)

if prompt is not None and isinstance(prompt, str):

batch_size = 1

elif prompt is not None and isinstance(prompt, list):

batch_size = len(prompt)

else:

batch_size = prompt_embeds.shape[0]

if prompt_embeds is None:

# textual inversion: process multi-vector tokens if necessary

if isinstance(pipeline, TextualInversionLoaderMixin):

prompt = pipeline.maybe_convert_prompt(prompt, pipeline.tokenizer)

text_inputs = pipeline.tokenizer(

prompt,

padding="max_length",

max_length=pipeline.tokenizer.model_max_length,

truncation=True,

return_tensors="pt",

)

text_input_ids = text_inputs.input_ids

untruncated_ids = pipeline.tokenizer(prompt, padding="longest", return_tensors="pt").input_ids

if untruncated_ids.shape[-1] >= text_input_ids.shape[-1] and not torch.equal(

text_input_ids, untruncated_ids

):

removed_text = pipeline.tokenizer.batch_decode(

untruncated_ids[:, pipeline.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" {pipeline.tokenizer.model_max_length} tokens: {removed_text}"

)

if hasattr(pipeline.text_encoder.config, "use_attention_mask") and pipeline.text_encoder.config.use_attention_mask:

attention_mask = text_inputs.attention_mask.to(device)

else:

attention_mask = None

if clip_skip is None:

prompt_embeds = pipeline.text_encoder(text_input_ids.to(device), attention_mask=attention_mask)

prompt_embeds = prompt_embeds[0]

else:

prompt_embeds = pipeline.text_encoder(

text_input_ids.to(device), attention_mask=attention_mask, output_hidden_states=True

)

# Access the `hidden_states` first, that contains a tuple of

# all the hidden states from the encoder layers. Then index into

# the tuple to access the hidden states from the desired layer.

prompt_embeds = prompt_embeds[-1][-(clip_skip + 1)]

# We also need to apply the final LayerNorm here to not mess with the

# representations. The `last_hidden_states` that we typically use for

# obtaining the final prompt representations passes through the LayerNorm

# layer.

prompt_embeds = pipeline.text_encoder.text_model.final_layer_norm(prompt_embeds)

if pipeline.text_encoder is not None:

prompt_embeds_dtype = pipeline.text_encoder.dtype

elif pipeline.unet is not None:

prompt_embeds_dtype = pipeline.unet.dtype

else:

prompt_embeds_dtype = prompt_embeds.dtype

prompt_embeds = prompt_embeds.to(dtype=prompt_embeds_dtype, device=device)

bs_embed, seq_len, _ = prompt_embeds.shape

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

prompt_embeds = prompt_embeds.repeat(1, num_images_per_prompt, 1)

prompt_embeds = prompt_embeds.view(bs_embed * num_images_per_prompt, seq_len, -1)

# get unconditional embeddings for classifier free guidance

if do_classifier_free_guidance and negative_prompt_embeds is None:

uncond_tokens: List[str]

if negative_prompt is None:

uncond_tokens = [""] * batch_size

elif prompt is not None and type(prompt) is not type(negative_prompt):

raise TypeError(

f"`negative_prompt` should be the same type to `prompt`, but got {type(negative_prompt)} !="

f" {type(prompt)}."

)

elif isinstance(negative_prompt, str):

uncond_tokens = [negative_prompt]

elif batch_size != len(negative_prompt):

raise ValueError(

f"`negative_prompt`: {negative_prompt} has batch size {len(negative_prompt)}, but `prompt`:"

f" {prompt} has batch size {batch_size}. Please make sure that passed `negative_prompt` matches"

" the batch size of `prompt`."

)

else:

uncond_tokens = negative_prompt

# textual inversion: process multi-vector tokens if necessary

if isinstance(pipeline, TextualInversionLoaderMixin):

uncond_tokens = pipeline.maybe_convert_prompt(uncond_tokens, pipeline.tokenizer)

max_length = prompt_embeds.shape[1]

uncond_input = pipeline.tokenizer(

uncond_tokens,

padding="max_length",

max_length=max_length,

truncation=True,

return_tensors="pt",

)

if hasattr(pipeline.text_encoder.config, "use_attention_mask") and pipeline.text_encoder.config.use_attention_mask:

attention_mask = uncond_input.attention_mask.to(device)

else:

attention_mask = None

negative_prompt_embeds = pipeline.text_encoder(

uncond_input.input_ids.to(device),

attention_mask=attention_mask,

)

negative_prompt_embeds = negative_prompt_embeds[0]

if do_classifier_free_guidance:

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

seq_len = negative_prompt_embeds.shape[1]

negative_prompt_embeds = negative_prompt_embeds.to(dtype=prompt_embeds_dtype, device=device)

negative_prompt_embeds = negative_prompt_embeds.repeat(1, num_images_per_prompt, 1)

negative_prompt_embeds = negative_prompt_embeds.view(batch_size * num_images_per_prompt, seq_len, -1)

if isinstance(pipeline, LoraLoaderMixin) and USE_PEFT_BACKEND:

# Retrieve the original scale by scaling back the LoRA layers

unscale_lora_layers(pipeline.text_encoder, lora_scale)

return prompt_embeds, negative_prompt_embeds