"""Local Qwen3-VL conditioning for Krea 2."""
from __future__ import annotations
from dataclasses import dataclass
from pathlib import Path
import torch
from diffusion_cli.checkpoint import loadStateDict
from diffusion_cli.config import ModelSource
from diffusion_cli.errors import DiffusionCliError
from diffusion_cli.krea2_model import Krea2Conditioning
from diffusion_cli.quantization import (
QuantizationManifest,
ScaledFp8Linear,
loadQuantizationManifest,
moveModulePreservingQuantization,
validateQuantizedStateDict,
)
KREA2_TAP_LAYERS = (2, 5, 8, 11, 14, 17, 20, 23, 26, 29, 32, 35)
KREA2_PROMPT_PREFIX = (
"<|im_start|>system\n"
"Describe 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"
)
KREA2_PROMPT_SUFFIX = "<|im_end|>\n<|im_start|>assistant\n"
KREA2_MAX_TEXT_TOKENS = 512
KREA2_PAD_TOKEN_ID = 151643
QWEN3_VL_4B_VOCAB_SIZE = 151936
QWEN3_VL_4B_HIDDEN_SIZE = 2560
QWEN3_VL_4B_INTERMEDIATE_SIZE = 9728
QWEN3_VL_4B_LAYERS = 36
QWEN3_VL_4B_HEADS = 32
QWEN3_VL_4B_KV_HEADS = 8
QWEN3_VL_4B_HEAD_DIM = 128
QWEN3_VL_4B_MAX_POSITION = 262144
QWEN3_VL_4B_ROPE_THETA = 5000000.0
@dataclass(frozen=True)
class Krea2PromptConditioning:
"""Positive and optional negative Krea conditioning."""
positive: Krea2Conditioning
negative: Krea2Conditioning | None = None
def buildQwen3Vl4BConfig():
"""Build the text-only Qwen3-VL-4B configuration in code."""
from transformers import Qwen3VLTextConfig
return Qwen3VLTextConfig(
vocab_size=QWEN3_VL_4B_VOCAB_SIZE,
hidden_size=QWEN3_VL_4B_HIDDEN_SIZE,
intermediate_size=QWEN3_VL_4B_INTERMEDIATE_SIZE,
num_hidden_layers=QWEN3_VL_4B_LAYERS,
num_attention_heads=QWEN3_VL_4B_HEADS,
num_key_value_heads=QWEN3_VL_4B_KV_HEADS,
head_dim=QWEN3_VL_4B_HEAD_DIM,
max_position_embeddings=QWEN3_VL_4B_MAX_POSITION,
rope_parameters={
"rope_theta": QWEN3_VL_4B_ROPE_THETA,
"rope_type": "default",
},
rms_norm_eps=1e-6,
attention_bias=False,
tie_word_embeddings=True,
pad_token_id=KREA2_PAD_TOKEN_ID,
use_cache=False,
)
def normalizeKreaTextStateDict(
state_dict: dict[str, torch.Tensor],
) -> dict[str, torch.Tensor]:
"""Normalize ComfyUI Qwen keys and consume ``comfy_quant`` metadata."""
normalized = {}
for key, value in state_dict.items():
if key.endswith(".comfy_quant"):
continue
if key.startswith("model."):
key = key.removeprefix("model.")
if key.startswith("visual."):
# Qwen3-VL files may carry a vision tower; Krea text conditioning
# deliberately documents it as an ignored component.
continue
normalized[key] = value
return normalized
def _findTokenBoundary(
token_ids: list[int],
prefix_ids: list[int],
) -> int:
"""Find the prefix boundary using token IDs rather than character offsets."""
for start in range(len(token_ids) - len(prefix_ids) + 1):
if token_ids[start : start + len(prefix_ids)] == prefix_ids:
return start + len(prefix_ids)
raise DiffusionCliError(
"Krea prompt template prefix is missing its known token boundary"
)
def _asTensor(value, *, dtype=None) -> torch.Tensor:
"""Convert tokenizer output from lists or tensors to a tensor."""
if isinstance(value, torch.Tensor):
return value if dtype is None else value.to(dtype=dtype)
return torch.tensor(value, dtype=dtype)
def initializeKreaRotaryEmbedding(model) -> None:
"""Initialize the non-persistent Qwen RoPE buffer after meta loading."""
config = model.config
head_dim = getattr(config, "head_dim", None)
if head_dim is None:
head_dim = config.hidden_size // config.num_attention_heads
rope_parameters = getattr(config, "rope_parameters", {})
rope_theta = rope_parameters.get(
"rope_theta",
QWEN3_VL_4B_ROPE_THETA,
)
parameter = next(model.parameters())
inv_freq = 1.0 / (
rope_theta
** (
torch.arange(
0,
head_dim,
2,
dtype=torch.float32,
device=parameter.device,
)
/ head_dim
)
)
model.rotary_emb.inv_freq = inv_freq
if hasattr(model.rotary_emb, "original_inv_freq"):
model.rotary_emb.original_inv_freq = inv_freq.clone()
class Krea2TextEncoder:
"""Load a local Qwen3-VL text backbone and emit twelve hidden-state taps."""
def __init__(
self,
model_path: ModelSource | Path,
tokenizer_path: Path,
device,
dtype=torch.bfloat16,
*,
model=None,
tokenizer=None,
) -> None:
self.device = device
self.dtype = dtype
self.tokenizer = tokenizer or self._loadTokenizer(tokenizer_path)
self.model = model or self._loadModel(model_path)
self.model.eval()
def encodePrompt(self, prompt: str) -> Krea2Conditioning:
"""Encode one prompt with the exact Krea chat template."""
if not isinstance(prompt, str):
raise DiffusionCliError("Krea prompt must be a string")
prefix_batch = self.tokenizer(
KREA2_PROMPT_PREFIX,
return_tensors="pt",
padding=False,
truncation=False,
add_special_tokens=False,
)
prefix_ids = _asTensor(prefix_batch["input_ids"], dtype=torch.long)
prefix_ids = prefix_ids.reshape(-1).tolist()
suffix_batch = self.tokenizer(
KREA2_PROMPT_SUFFIX,
return_tensors="pt",
padding=False,
truncation=False,
add_special_tokens=False,
)
suffix_ids = _asTensor(
suffix_batch["input_ids"],
dtype=torch.long,
)
suffix_mask = _asTensor(
suffix_batch["attention_mask"],
dtype=torch.bool,
)
prompt_batch = self.tokenizer(
[KREA2_PROMPT_PREFIX + prompt],
return_tensors="pt",
padding="max_length",
truncation=True,
max_length=(
KREA2_MAX_TEXT_TOKENS
+ len(prefix_ids)
- suffix_ids.shape[1]
),
add_special_tokens=False,
)
input_ids = _asTensor(
prompt_batch["input_ids"],
dtype=torch.long,
)
attention_mask = _asTensor(
prompt_batch["attention_mask"],
dtype=torch.bool,
)
input_ids = torch.cat((input_ids, suffix_ids), dim=1)
attention_mask = torch.cat((attention_mask, suffix_mask), dim=1)
boundary = _findTokenBoundary(
input_ids[0].tolist(),
prefix_ids,
)
input_ids = input_ids.to(self.device)
attention_mask = attention_mask.to(self.device)
with torch.inference_mode():
output = self.model(
input_ids=input_ids,
attention_mask=attention_mask,
output_hidden_states=True,
use_cache=False,
)
try:
hidden = torch.stack(
[output.hidden_states[index] for index in KREA2_TAP_LAYERS],
dim=2,
)
except (AttributeError, IndexError) as exc:
raise DiffusionCliError(
"Text encoder checkpoint does not expose the 36 Qwen layers "
"required by Krea 2"
) from exc
hidden = hidden[:, boundary:]
attention_mask = attention_mask[:, boundary:]
if hidden.shape[2] != len(KREA2_TAP_LAYERS):
raise DiffusionCliError("Krea text encoder returned invalid tap layout")
if hidden.shape[-1] != QWEN3_VL_4B_HIDDEN_SIZE:
raise DiffusionCliError(
"Text encoder checkpoint is not Qwen3-VL-4B"
)
return Krea2Conditioning(hidden, attention_mask)
def encode(self, prompt: str) -> Krea2Conditioning:
"""Encode one positive prompt."""
return self.encodePrompt(prompt)
def encodePrompts(
self,
prompt: str,
negative_prompt: str = "",
*,
cfg: float = 0.0,
) -> Krea2PromptConditioning:
"""Encode positive and, only when requested, negative prompts."""
positive = self.encodePrompt(prompt)
negative = self.encodePrompt(negative_prompt) if cfg > 0 else None
return Krea2PromptConditioning(positive, negative)
def toDevice(self, device, dtype=None) -> None:
"""Move the Qwen backbone without expanding its FP8 weights."""
self.device = device
self.dtype = dtype or self.dtype
moveModulePreservingQuantization(
self.model,
device,
self.dtype,
)
self.model.eval()
def toCpu(self) -> None:
"""Move the Qwen backbone back to CPU."""
self.toDevice(torch.device("cpu"))
def _loadTokenizer(self, tokenizer_path: Path):
"""Load only tokenizer assets from an explicit local directory."""
from transformers import Qwen2Tokenizer
tokenizer = Qwen2Tokenizer.from_pretrained(
tokenizer_path,
local_files_only=True,
padding_side="right",
)
tokenizer.pad_token_id = KREA2_PAD_TOKEN_ID
tokenizer.padding_side = "right"
return tokenizer
def _loadModel(self, model_path: ModelSource | Path):
"""Construct the text-only model and strictly load local tensors."""
from transformers import Qwen3VLTextModel
source = model_path if isinstance(model_path, ModelSource) else ModelSource(
model_path,
"text_encoder",
)
raw_state = loadStateDict(source)
manifest = validateQuantizedStateDict(
normalizeKreaTextStateDict(raw_state),
loadQuantizationManifest(source),
)
normalized = normalizeKreaTextStateDict(raw_state)
with torch.device("meta"):
model = Qwen3VLTextModel(buildQwen3Vl4BConfig())
for name in manifest.layers:
module = model.get_submodule(name)
if not isinstance(module, torch.nn.Linear):
raise DiffusionCliError(
f"Quantized Qwen layer is not linear: {name}"
)
replacement = ScaledFp8Linear(
module.in_features,
module.out_features,
bias=module.bias is not None,
device="meta",
dtype=self.dtype,
)
parent_name, _, child_name = name.rpartition(".")
parent = model.get_submodule(parent_name) if parent_name else model
setattr(parent, child_name, replacement)
missing, unexpected = model.load_state_dict(
normalized,
strict=False,
assign=True,
)
allowed_missing = {"rotary_emb.inv_freq"}
missing = [key for key in missing if key not in allowed_missing]
if missing or unexpected:
raise DiffusionCliError(
"Krea text encoder checkpoint key mismatch: "
f"missing={missing[:3]}, unexpected={unexpected[:3]}"
)
initializeKreaRotaryEmbedding(model)
moveModulePreservingQuantization(model, self.device, self.dtype)
return model
loadKrea2TextEncoder = Krea2TextEncoder