Spaces:
Running on Zero
Running on Zero
File size: 5,803 Bytes
855e1f5 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 | """Qwen3-VL prompt encoding for Mage-Flow.
The reference implementation (mage_flow/models/modules/text_encoder.py) packs
many prompts into one varlen forward; per-sequence the math reduces to a plain
causal forward with positions ``arange(L)`` replicated on every mrope axis and
the leading system-prompt tokens dropped from the output. ai-toolkit encodes
prompts one at a time (embeds are cached per caption), so that per-sequence
form is implemented here directly on the stock HF Qwen3-VL module.
The DiT consumes the final hidden states ``[L - drop_idx, 2560]``; the pooled
"vec" of the reference is unused by the released checkpoints (vec_type null —
the transformer adds a zero vector), so it is not computed.
"""
from __future__ import annotations
import math
from typing import List, Optional
import torch
from toolkit.models.v2.text_encoders.qwen3_vl import patch_qwen_vl_patch_embed
# Prompt templates from the reference (mage_flow/models/utils.py). ``start_idx``
# is the number of leading (system prompt) tokens dropped from the encoder
# output before it conditions the DiT.
PROMPT_TEMPLATE = {
"mage-flow": {
"template": (
"<|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{}<|im_end|>\n<|im_start|>assistant\n"
),
"start_idx": 34,
},
"mage-flow-edit": {
"template": (
"<|im_start|>system\nDescribe the key features of the input image (color, shape, size, texture,"
" objects, background), then explain how the user's text instruction should alter or modify the image. "
"Generate a new image that meets the user's requirements while maintaining consistency with the original "
"input where appropriate.<|im_end|>\n<|im_start|>user\n{}<|im_end|>\n<|im_start|>assistant\n"
),
"start_idx": 64,
},
}
# Fixed image placeholder used at edit training time (one per reference image).
EDIT_IMAGE_PLACEHOLDER = "<|vision_start|><|image_pad|><|vision_end|>"
def edit_prompt_body(instruction: str, num_refs: int) -> str:
"""Training-time multi-reference prompt body:
``Image 1: <ph>Image 2: <ph>…{instruction}``."""
prefix = "".join(
f"Image {j}: {EDIT_IMAGE_PLACEHOLDER}" for j in range(1, num_refs + 1)
)
return prefix + instruction
def resize_vl_images(
images: List[torch.Tensor], max_long_edge: int = 384
) -> List["PIL.Image.Image"]:
"""Prepare reference images for the Qwen3-VL conditioning pass.
Matches the reference's ``_resize_long_edge`` (cap the long edge at
``max_long_edge``, preserving aspect ratio, BICUBIC, never upscale) — the
MLLM only needs a coarse view of the reference; full-resolution detail
flows through the VAE reference latents. Input tensors are ``(C, H, W)``
or ``(1, C, H, W)`` in [0, 1]; output is PIL for the Qwen processor.
"""
from torchvision.transforms.functional import to_pil_image
out = []
for img in images:
if img.dim() == 4:
img = img[0]
img = img.float().clamp(0, 1).cpu()
pil = to_pil_image(img)
if max_long_edge and max_long_edge > 0:
w, h = pil.size
long_edge = max(w, h)
if long_edge > max_long_edge:
scale = max_long_edge / long_edge
new_w = max(1, int(round(w * scale)))
new_h = max(1, int(round(h * scale)))
pil = pil.resize((new_w, new_h), resample=3) # PIL.Image.BICUBIC
out.append(pil)
return out
@torch.no_grad()
def encode_mageflow_prompt(
text_encoder, # Qwen3VLForConditionalGeneration
tokenizer,
prompt: str,
template_name: str = "mage-flow",
max_length: int = 2048,
images: Optional[list] = None, # list of PIL images (already VL-resized)
processor=None, # AutoProcessor, required when images are given
dtype: torch.dtype = torch.bfloat16,
) -> torch.Tensor:
"""Encode one prompt (optionally with reference images) to DiT conditioning.
Returns ``(L - start_idx, 2560)`` final hidden states with the system
prompt dropped, matching the reference ``TextEncoder.forward`` txt output.
"""
info = PROMPT_TEMPLATE[template_name]
template = info["template"]
drop_idx = int(info["start_idx"])
device = next(text_encoder.parameters()).device
if images:
if processor is None:
raise ValueError(
"encoding reference images requires the Qwen3-VL AutoProcessor"
)
formatted = template.format(edit_prompt_body(prompt, len(images)))
vl = processor(text=[formatted], images=list(images), return_tensors="pt")
input_ids = vl["input_ids"].to(device)
extra = {}
for key in ("pixel_values", "image_grid_thw"):
if vl.get(key) is not None:
extra[key] = vl[key].to(device)
else:
formatted = template.format(prompt)
input_ids = tokenizer(
formatted,
max_length=max_length + drop_idx,
truncation=True,
return_tensors="pt",
).input_ids.to(device)
extra = {}
# The reference packed encoder gives every sequence positions arange(L)
# replicated on all mrope axes (even for image tokens); a 2D position_ids
# is expanded to exactly that by the HF text model.
position_ids = torch.arange(input_ids.shape[1], device=device).unsqueeze(0)
out = text_encoder.model(
input_ids=input_ids,
position_ids=position_ids,
**extra,
)
hidden = out.last_hidden_state[0] # (L, D)
return hidden[drop_idx:].to(dtype)
|