27 lines
1.0 KiB
Python
27 lines
1.0 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
|
|
from __future__ import annotations
|
|
|
|
import torch
|
|
|
|
|
|
def txt_seq_lens_from_embeds(prompt_embeds: torch.Tensor | None) -> list[int] | None:
|
|
"""Return per-row RoPE text lengths from padded encoder hidden-state width.
|
|
|
|
Diffusers derives Qwen-Image text RoPE length from
|
|
``encoder_hidden_states.shape[1]``, not from the count of valid (non-padding)
|
|
tokens in the attention mask. Callers that right-pad prompt embeddings for
|
|
batching must pass the padded width here so ``QwenEmbedRope`` builds a table
|
|
long enough for every text position.
|
|
"""
|
|
if prompt_embeds is None:
|
|
return None
|
|
if prompt_embeds.ndim == 2:
|
|
prompt_embeds = prompt_embeds.unsqueeze(0)
|
|
if prompt_embeds.ndim != 3:
|
|
raise ValueError(f"prompt_embeds must be 2D or 3D, got shape={tuple(prompt_embeds.shape)}")
|
|
seq_len = int(prompt_embeds.shape[1])
|
|
batch_size = int(prompt_embeds.shape[0])
|
|
return [seq_len] * batch_size
|