Function get_weighted_text_embeddings

modules/prompt_parser_diffusers.py:604–754

(pipe, prompt: str = "", neg_prompt: str = "", clip_skip: int | None = None, prompt_mean_norm=None, diffusers_zeros_prompt_pad=None, te_pooled_embeds=None)

def get_weighted_text_embeddings(pipe, prompt: str = "", neg_prompt: str = "", clip_skip: int | None = None, prompt_mean_norm=None, diffusers_zeros_prompt_pad=None, te_pooled_embeds=None):
    device = devices.device
    if prompt is None:
        prompt = ''
    if neg_prompt is None:
        neg_prompt = ''
    SD3 = bool(hasattr(pipe, 'text_encoder_3') and not hasattr(pipe, 'text_encoder_4'))
    prompt, prompt_2, prompt_3, prompt_4 = split_prompts(pipe, prompt, SD3)
    neg_prompt, neg_prompt_2, neg_prompt_3, neg_prompt_4 = split_prompts(pipe, neg_prompt, SD3)

    if "Flux" in pipe.__class__.__name__: # clip is only used for the pooled embeds
        prompt_embeds, pooled_prompt_embeds, _ = pipe.encode_prompt(prompt=prompt, prompt_2=prompt_2, device=device, num_images_per_prompt=1)
        return prompt_embeds, pooled_prompt_embeds, None, None, None, None # no negative support

    if "Chroma" in pipe.__class__.__name__: # does not use clip and has no pooled embeds
        prompt_embeds, _, prompt_attention_mask, negative_prompt_embeds, _, negative_prompt_attention_mask = pipe.encode_prompt(prompt=prompt, negative_prompt=neg_prompt, device=device, num_images_per_prompt=1)
        return prompt_embeds, None, prompt_attention_mask, negative_prompt_embeds, None, negative_prompt_attention_mask

    if "HiDreamImage" in pipe.__class__.__name__: # clip is only used for the pooled embeds
        prompt_embeds_t5, negative_prompt_embeds_t5, prompt_embeds_llama3, negative_prompt_embeds_llama3, pooled_prompt_embeds, negative_pooled_prompt_embeds = pipe.encode_prompt(
            prompt=prompt, prompt_2=prompt_2, prompt_3=prompt_3, prompt_4=prompt_4,
            negative_prompt=neg_prompt, negative_prompt_2=neg_prompt_2, negative_prompt_3=neg_prompt_3, negative_prompt_4=neg_prompt_4,
            device=device, num_images_per_prompt=1,
        )
        prompt_embeds = [prompt_embeds_t5, prompt_embeds_llama3]
        negative_prompt_embeds = [negative_prompt_embeds_t5, negative_prompt_embeds_llama3]
        return prompt_embeds, pooled_prompt_embeds, None, negative_prompt_embeds, negative_pooled_prompt_embeds, None

    _zeros_pad = diffusers_zeros_prompt_pad if diffusers_zeros_prompt_pad is not None else shared.opts.diffusers_zeros_prompt_pad
    _te_pooled = te_pooled_embeds if te_pooled_embeds is not None else shared.opts.te_pooled_embeds
    if prompt != prompt_2:
        ps = [get_prompts_with_weights(pipe, p, prompt_mean_norm=prompt_mean_norm) for p in [prompt, prompt_2]]
        ns = [get_prompts_with_weights(pipe, p, prompt_mean_norm=prompt_mean_norm) for p in [neg_prompt, neg_prompt_2]]
    else:
        ps = 2 * [get_prompts_with_weights(pipe, prompt, prompt_mean_norm=prompt_mean_norm)]
        ns = 2 * [get_prompts_with_weights(pipe, neg_prompt, prompt_mean_norm=prompt_mean_norm)]

    positives, positive_weights = zip(*ps, strict=False)
    negatives, negative_weights = zip(*ns, strict=False)
    if hasattr(pipe, "tokenizer_2") and not hasattr(pipe, "tokenizer"):
        positives.pop(0)
        positive_weights.pop(0)
        negatives.pop(0)
        negative_weights.pop(0)

    embedding_providers = prepare_embedding_providers(pipe, clip_skip)
    if len(embedding_providers) == 0:
        log.error("Prompt encode: cannot find text encoder in model")
        return None, None, None, None, None, None
    empty_embedding_providers = None
    if 'StableCascade' in pipe.__class__.__name__:
        empty_embedding_providers = [embedding_providers[1]]
        embedding_providers = [embedding_providers[0]]

    prompt_embeds = []
    negative_prompt_embeds = []
    pooled_prompt_embeds = []
    negative_pooled_prompt_embeds = []