137 lines
5.5 KiB
Python
137 lines
5.5 KiB
Python
"""Qwen3-VL text conditioning for Krea 2.
|
|
|
|
Vendored / adapted from the reference ``encoder.py``. Krea 2 conditions on a
|
|
*stack* of hidden states pulled from several layers of Qwen3-VL-4B-Instruct
|
|
(``SELECT_LAYERS``), wrapped in a fixed instruction template. The MMDiT's
|
|
``TextFusionTransformer`` later collapses that layer axis down to one.
|
|
|
|
The reference encodes a whole batch padded to ``max_length``; here we encode one
|
|
prompt at a time at its natural length (the ai-toolkit pattern -- caches stay
|
|
small, any prompts can share a batch, and per-sample padding is deferred to the
|
|
model call). The fixed instruction prefix is fed through the model as context but
|
|
its hidden states are sliced off the returned features, exactly like the
|
|
reference.
|
|
"""
|
|
|
|
from typing import List, Optional
|
|
|
|
import torch
|
|
from torch import Tensor
|
|
|
|
|
|
# Layers of Qwen3-VL whose hidden states are stacked and fed to the MMDiT (12).
|
|
SELECT_LAYERS: tuple[int, ...] = (2, 5, 8, 11, 14, 17, 20, 23, 26, 29, 32, 35)
|
|
|
|
# Fixed instruction template wrapped around every prompt. The prefix is fed
|
|
# through the model as context but its hidden states are dropped from the output
|
|
# (the assistant only ever sees the prompt + suffix tokens as conditioning).
|
|
PROMPT_TEMPLATE_ENCODE_PREFIX = (
|
|
"<|im_start|>system\nDescribe the image by detailing the color, shape, size, "
|
|
"texture, quantity, text, spatial relationships of the objects and "
|
|
"background:<|im_end|>\n<|im_start|>user\n"
|
|
)
|
|
PROMPT_TEMPLATE_ENCODE_SUFFIX = "<|im_end|>\n<|im_start|>assistant\n"
|
|
|
|
# Number of leading tokens (the system prefix) sliced off the encoded features.
|
|
PROMPT_TEMPLATE_ENCODE_START_IDX = 34
|
|
|
|
|
|
@torch.no_grad()
|
|
def encode_krea_prompt(
|
|
qwen,
|
|
tokenizer,
|
|
processor,
|
|
prompt: str,
|
|
max_length: int = 512,
|
|
select_layers: tuple[int, ...] = SELECT_LAYERS,
|
|
prefix_idx: int = PROMPT_TEMPLATE_ENCODE_START_IDX,
|
|
images: Optional[List[Tensor]] = None,
|
|
vl_processor=None,
|
|
dtype: Optional[torch.dtype] = None,
|
|
) -> Tensor:
|
|
"""Encode a single prompt into stacked Qwen3-VL hidden states.
|
|
|
|
Returns a ``(L, num_select_layers, hidden)`` float tensor (in the encoder's
|
|
dtype) holding the prompt + suffix token features -- the system prefix has
|
|
been sliced off. ``L`` is the natural (unpadded) length so the caller stores
|
|
one tensor per prompt and pads to the batch max at the model call.
|
|
|
|
``images`` (optional) are reference images -- ``(C, H, W)`` tensors in
|
|
``[0, 1]`` -- embedded in the user message ahead of the prompt via named
|
|
vision placeholders (``Picture 1: <|vision_start|><|image_pad|><|vision_end|>``,
|
|
the ComfyUI ``TextEncodeQwenImageEditPlus`` layout). The ``vl_processor``
|
|
(Qwen3-VL AutoProcessor) expands each ``<|image_pad|>`` to the image's token
|
|
grid, so the returned features carry the vision tokens as extra conditioning.
|
|
The system prefix is unchanged, so ``prefix_idx`` slicing stays valid and the
|
|
image + prompt tokens all survive the slice.
|
|
"""
|
|
device = qwen.device
|
|
|
|
# The suffix ("...assistant\n") is tokenized without the BOS/template extras
|
|
# the main tokenizer adds, matching the reference's separate processor pass.
|
|
suffix_inputs = processor(
|
|
text=[PROMPT_TEMPLATE_ENCODE_SUFFIX], return_tensors="pt"
|
|
).to(device, non_blocking=True)
|
|
suffix_ids = suffix_inputs["input_ids"]
|
|
suffix_mask = suffix_inputs["attention_mask"].bool()
|
|
|
|
extra_inputs = {}
|
|
if images is not None and len(images) > 0:
|
|
image_prompt = "".join(
|
|
f"Picture {i + 1}: <|vision_start|><|image_pad|><|vision_end|>"
|
|
for i in range(len(images))
|
|
)
|
|
text = PROMPT_TEMPLATE_ENCODE_PREFIX + image_prompt + prompt
|
|
# No truncation here: the expanded image-pad runs must stay intact.
|
|
inputs = vl_processor(
|
|
text=[text],
|
|
images=list(images),
|
|
return_tensors="pt",
|
|
do_rescale=False,
|
|
).to(device)
|
|
for k, v in inputs.items():
|
|
if k in ("input_ids", "attention_mask"):
|
|
continue
|
|
if (
|
|
isinstance(v, torch.Tensor)
|
|
and v.is_floating_point()
|
|
and dtype is not None
|
|
):
|
|
v = v.to(dtype)
|
|
extra_inputs[k] = v
|
|
else:
|
|
# Prefix + prompt at natural length (no padding); truncate very long prompts.
|
|
text = PROMPT_TEMPLATE_ENCODE_PREFIX + prompt
|
|
inputs = tokenizer(
|
|
[text],
|
|
truncation=True,
|
|
return_length=False,
|
|
return_overflowing_tokens=False,
|
|
max_length=max_length + prefix_idx,
|
|
return_tensors="pt",
|
|
).to(device, non_blocking=True)
|
|
|
|
input_ids = torch.cat([inputs["input_ids"], suffix_ids], dim=1)
|
|
mask = torch.cat([inputs["attention_mask"].bool(), suffix_mask], dim=1)
|
|
|
|
# mm_token_type_ids (used for M-RoPE) must cover the appended suffix tokens
|
|
# too; they are plain text -> type 0.
|
|
if "mm_token_type_ids" in extra_inputs:
|
|
tt = extra_inputs["mm_token_type_ids"]
|
|
extra_inputs["mm_token_type_ids"] = torch.cat(
|
|
[tt, torch.zeros_like(suffix_ids, dtype=tt.dtype)], dim=1
|
|
)
|
|
|
|
states = qwen(
|
|
input_ids=input_ids,
|
|
attention_mask=mask,
|
|
output_hidden_states=True,
|
|
**extra_inputs,
|
|
)
|
|
|
|
# (1, L, num_layers, hidden)
|
|
hiddens = torch.stack([states.hidden_states[i] for i in select_layers], dim=2)
|
|
# Drop the system-prefix tokens; what remains is prompt + suffix conditioning.
|
|
hiddens = hiddens[:, prefix_idx:]
|
|
return hiddens[0]
|