"""Prompt templating and Qwen conditioning for Z-Image."""
from __future__ import annotations
from dataclasses import dataclass
from pathlib import Path
from diffusion_cli.checkpoint import loadStateDict
from diffusion_cli.config import ModelSource, TEXT_ENCODER_ROLE
QWEN3_4B_VOCAB_SIZE = 151936
QWEN3_4B_HIDDEN_SIZE = 2560
QWEN3_4B_INTERMEDIATE_SIZE = 9728
QWEN3_4B_NUM_HIDDEN_LAYERS = 36
QWEN3_4B_NUM_ATTENTION_HEADS = 32
QWEN3_4B_NUM_KEY_VALUE_HEADS = 8
QWEN3_4B_MAX_POSITION_EMBEDDINGS = 40960
QWEN3_4B_HEAD_DIM = 128
QWEN3_4B_ROPE_THETA = 1000000.0
QWEN3_4B_PAD_TOKEN_ID = 151643
QWEN3_4B_LAYER_INDEX = -2
@dataclass(frozen=True)
class TextConditioning:
"""Text embeddings and masks consumed by the diffusion model."""
hidden_states: object
attention_mask: object
token_count: int
@dataclass(frozen=True)
class PromptConditioning:
"""Positive and negative prompt conditioning for CFG sampling."""
positive: TextConditioning
negative: TextConditioning
def applyZImagePromptTemplate(prompt: str) -> str:
"""Wrap a prompt with the chat template used by ComfyUI Z-Image."""
return f"<|im_start|>user\n{prompt}<|im_end|>\n<|im_start|>assistant\n"
def buildQwen3_4BConfig():
"""Build the Qwen3-4B config used by ComfyUI's Z-Image text encoder."""
from transformers import Qwen3Config
return Qwen3Config(
vocab_size=QWEN3_4B_VOCAB_SIZE,
hidden_size=QWEN3_4B_HIDDEN_SIZE,
intermediate_size=QWEN3_4B_INTERMEDIATE_SIZE,
num_hidden_layers=QWEN3_4B_NUM_HIDDEN_LAYERS,
num_attention_heads=QWEN3_4B_NUM_ATTENTION_HEADS,
num_key_value_heads=QWEN3_4B_NUM_KEY_VALUE_HEADS,
max_position_embeddings=QWEN3_4B_MAX_POSITION_EMBEDDINGS,
rms_norm_eps=1e-6,
rope_theta=QWEN3_4B_ROPE_THETA,
head_dim=QWEN3_4B_HEAD_DIM,
attention_bias=False,
tie_word_embeddings=False,
pad_token_id=QWEN3_4B_PAD_TOKEN_ID,
bos_token_id=None,
eos_token_id=None,
use_cache=False,
)
def normalizeQwenStateDict(state_dict: dict[str, object]) -> dict[str, object]:
"""Convert ComfyUI Qwen checkpoint keys to Transformers model keys."""
normalized = {}
for key, value in state_dict.items():
if key.startswith("model."):
normalized[key.removeprefix("model.")] = value
else:
normalized[key] = value
return normalized
def initializeQwenRotaryEmbedding(model) -> None:
"""Initialize non-persistent Qwen RoPE buffers after meta allocation."""
import torch
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_4B_ROPE_THETA)
device = model.rotary_emb.inv_freq.device
inv_freq = 1.0 / (
rope_theta
** (
torch.arange(0, head_dim, 2, dtype=torch.float32, device=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 ZImageTextEncoder:
"""Load the local Qwen3-4B encoder and produce text conditioning."""
def __init__(
self,
model_path: ModelSource | Path | None,
tokenizer_path: Path,
device,
dtype,
*,
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 encode(self, prompt: str) -> TextConditioning:
"""Encode one prompt into a hidden-state tensor and attention mask."""
import torch
text = applyZImagePromptTemplate(prompt)
token_batch = self.tokenizer(
text,
return_tensors="pt",
padding=False,
truncation=False,
add_special_tokens=False,
)
input_ids = token_batch["input_ids"].to(self.device)
attention_mask = token_batch["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,
)
hidden_states = output.hidden_states[QWEN3_4B_LAYER_INDEX]
return TextConditioning(
hidden_states=hidden_states,
attention_mask=attention_mask,
token_count=int(attention_mask.sum().item()),
)
def encodePrompts(
self,
prompt: str,
negative_prompt: str,
) -> PromptConditioning:
"""Encode positive and negative prompts for guided sampling."""
return PromptConditioning(
positive=self.encode(prompt),
negative=self.encode(negative_prompt),
)
def toDevice(self, device, dtype) -> None:
"""Move the encoder model to the active generation device."""
self.device = device
self.dtype = dtype
self.model.to(device=device, dtype=dtype)
self.model.eval()
def toCpu(self) -> None:
"""Move the encoder model back to CPU memory."""
import torch
self.device = torch.device("cpu")
self.model.to(device=self.device)
self.model.eval()
def _loadTokenizer(self, tokenizer_path: Path):
from transformers import Qwen2Tokenizer
tokenizer = Qwen2Tokenizer.from_pretrained(
tokenizer_path,
local_files_only=True,
padding_side="right",
)
tokenizer.pad_token_id = QWEN3_4B_PAD_TOKEN_ID
return tokenizer
def _loadModel(self, model_path: ModelSource | Path | None):
import torch
from transformers import Qwen3Model
if model_path is None:
raise ValueError("model_path is required when model is absent")
if isinstance(model_path, ModelSource):
model_source = model_path
else:
model_source = ModelSource(model_path, TEXT_ENCODER_ROLE)
state_dict = normalizeQwenStateDict(loadStateDict(model_source))
with torch.device("meta"):
model = Qwen3Model(buildQwen3_4BConfig())
model.to_empty(device="cpu")
missing_keys, unexpected_keys = model.load_state_dict(
state_dict,
strict=False,
assign=True,
)
allowed_missing = {"rotary_emb.inv_freq"}
unexpected = [key for key in unexpected_keys if key]
missing = [key for key in missing_keys if key not in allowed_missing]
if missing or unexpected:
details = []
if missing:
details.append(f"missing keys: {missing[:8]}")
if unexpected:
details.append(f"unexpected keys: {unexpected[:8]}")
raise ValueError("Could not load Qwen text encoder: " + "; ".join(
details,
))
initializeQwenRotaryEmbedding(model)
return model.to(device=self.device, dtype=self.dtype)