diffsynth/models/ideogram4_text_encoder.py

import warnings

import torch
import torch.nn as nn
import torch.nn.functional as F

LLM_TOKEN_INDICATOR = 3
QWEN3_VL_ACTIVATION_LAYERS = (0, 3, 6, 9, 12, 15, 18, 21, 24, 27, 30, 33, 35)

FP8_E4M3_MAX = 448.0
FP8_WEIGHT_DTYPE = torch.float8_e4m3fn
FP8_SCALE_SUFFIX = ".weight_scale"

_BNB_SIBLING_SUFFIXES = (
    ".absmax",
    ".quant_map",
    ".nested_absmax",
    ".nested_quant_map",
)


class Fp8Linear(nn.Module):
    """Linear layer holding an e4m3 float8 weight + per-row float32 scale."""

    weight: torch.Tensor
    weight_scale: torch.Tensor
    bias: torch.Tensor | None

    def __init__(
        self,
        in_features: int,
        out_features: int,
        bias: bool,
        compute_dtype: torch.dtype,
    ) -> None:
        super().__init__()
        self.in_features = in_features
        self.out_features = out_features
        self.compute_dtype = compute_dtype
        self.register_buffer(
            "weight",
            torch.empty(out_features, in_features, dtype=FP8_WEIGHT_DTYPE),
        )
        self.register_buffer("weight_scale", torch.empty(out_features, dtype=torch.float32))
        if bias:
            self.register_buffer("bias", torch.empty(out_features, dtype=compute_dtype))
        else:
            self.bias = None

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        w = self.weight.to(x.dtype) * self.weight_scale.to(x.dtype).unsqueeze(1)
        bias = self.bias.to(x.dtype) if self.bias is not None else None
        return F.linear(x, w, bias)


def is_fp8_state_dict(state_dict: dict[str, torch.Tensor]) -> bool:
    return any(k.endswith(FP8_SCALE_SUFFIX) for k in state_dict) or any(
        v.dtype == FP8_WEIGHT_DTYPE for v in state_dict.values()
    )


def is_bnb4bit_state_dict(state_dict: dict[str, torch.Tensor]) -> bool:
    return any(".quant_state.bitsandbytes__" in k for k in state_dict)


def swap_linears_to_fp8(
    module: nn.Module,
    state_dict: dict[str, torch.Tensor],
    compute_dtype: torch.dtype,
    *,
    prefix: str = "",
) -> None:
    for name, child in list(module.named_children()):
        child_prefix = f"{prefix}{name}"
        if (
            isinstance(child, nn.Linear) and f"{child_prefix}{FP8_SCALE_SUFFIX}" in state_dict
        ):
            setattr(
                module,
                name,
                Fp8Linear(
                    child.in_features,
                    child.out_features,
                    bias=child.bias is not None,
                    compute_dtype=compute_dtype,
                ),
            )
        else:
            swap_linears_to_fp8(child, state_dict, compute_dtype, prefix=f"{child_prefix}.")


def load_fp8_state_dict(
    model: nn.Module,
    state_dict: dict[str, torch.Tensor],
    device: torch.device,
    dtype: torch.dtype,
    *,
    assign: bool = False,
    strict: bool = True,
) -> None:
    prepared: dict[str, torch.Tensor] = {}
    for k, v in state_dict.items():
        if v.dtype == FP8_WEIGHT_DTYPE:
            prepared[k] = v.to(device=device)
        elif k.endswith(FP8_SCALE_SUFFIX):
            prepared[k] = v.to(device=device, dtype=torch.float32)
        elif v.is_floating_point():
            prepared[k] = v.to(device=device, dtype=dtype)
        else:
            prepared[k] = v.to(device=device)

    missing, unexpected = model.load_state_dict(prepared, strict=False, assign=assign)
    if unexpected:
        raise RuntimeError(f"unexpected keys after fp8 load: {unexpected[:10]}")
    if missing:
        if strict:
            raise RuntimeError(f"missing keys after fp8 load: {missing[:10]}")
        warnings.warn(f"missing keys after fp8 load: {missing[:10]}", stacklevel=2)

    model.to(device)


def swap_linears_to_bnb4bit(
    module: nn.Module,
    compute_dtype: torch.dtype,
    *,
    quant_type: str = "nf4",
    compress_statistics: bool = False,
) -> None:
    import bitsandbytes as bnb
    for name, child in list(module.named_children()):
        if isinstance(child, nn.Linear):
            new_linear = bnb.nn.Linear4bit(
                child.in_features,
                child.out_features,
                bias=child.bias is not None,
                compute_dtype=compute_dtype,
                compress_statistics=compress_statistics,
                quant_type=quant_type,
            )
            setattr(module, name, new_linear)
        else:
            swap_linears_to_bnb4bit(
                child,
                compute_dtype,
                quant_type=quant_type,
                compress_statistics=compress_statistics,
            )


def load_bnb4bit_state_dict(
    model: nn.Module,
    state_dict: dict[str, torch.Tensor],
    device: torch.device,
    dtype: torch.dtype,
) -> None:
    import bitsandbytes as bnb
    consumed: set[str] = set()
    for full_name, tensor in state_dict.items():
        if ".quant_state." in full_name or full_name.endswith(_BNB_SIBLING_SUFFIXES):
            continue
        parent_path, _, param_name = full_name.rpartition(".")
        parent = model.get_submodule(parent_path) if parent_path else model
        current = parent._parameters.get(param_name)
        if not isinstance(current, bnb.nn.Params4bit):
            continue
        prefix = full_name + "."
        quantized_stats = {k: v for k, v in state_dict.items() if k.startswith(prefix)}
        consumed.add(full_name)
        consumed.update(quantized_stats.keys())
        parent._parameters[param_name] = bnb.nn.Params4bit.from_prequantized(
            data=tensor,
            quantized_stats=quantized_stats,
            requires_grad=False,
            device=device,
        )

    remaining = {k: v for k, v in state_dict.items() if k not in consumed}
    for k in list(remaining):
        if remaining[k].is_floating_point():
            remaining[k] = remaining[k].to(device=device, dtype=dtype)
        else:
            remaining[k] = remaining[k].to(device=device)

    missing, unexpected = model.load_state_dict(remaining, strict=False)
    real_missing = [m for m in missing if m not in consumed]
    if real_missing:
        raise RuntimeError(f"missing keys after quantized load: {real_missing[:10]}")
    if unexpected:
        raise RuntimeError(f"unexpected keys after quantized load: {unexpected[:10]}")

    for p in model.parameters():
        if isinstance(p, bnb.nn.Params4bit):
            continue
        if p.is_floating_point() and p.dtype != dtype:
            p.data = p.data.to(dtype=dtype)
        if p.device != device:
            p.data = p.data.to(device=device)
    for name, b in list(model.named_buffers()):
        if b.is_floating_point() and b.dtype != dtype:
            parent_path, _, leaf = name.rpartition(".")
            parent = model.get_submodule(parent_path) if parent_path else model
            parent.register_buffer(
                leaf,
                b.to(device=device, dtype=dtype),
                persistent=leaf not in parent._non_persistent_buffers_set,
            )
        elif b.device != device:
            parent_path, _, leaf = name.rpartition(".")
            parent = model.get_submodule(parent_path) if parent_path else model
            parent.register_buffer(
                leaf,
                b.to(device=device),
                persistent=leaf not in parent._non_persistent_buffers_set,
            )


_DEFAULT_TEXT_ENCODER_CONFIG = {
    "architectures": ["Qwen3VLModel"],
    "dtype": "bfloat16",
    "image_token_id": 151655,
    "model_type": "qwen3_vl",
    "text_config": {
        "attention_bias": False,
        "attention_dropout": 0.0,
        "bos_token_id": 151643,
        "dtype": "bfloat16",
        "eos_token_id": 151645,
        "head_dim": 128,
        "hidden_act": "silu",
        "hidden_size": 4096,
        "initializer_range": 0.02,
        "intermediate_size": 12288,
        "max_position_embeddings": 262144,
        "model_type": "qwen3_vl_text",
        "num_attention_heads": 32,
        "num_hidden_layers": 36,
        "num_key_value_heads": 8,
        "pad_token_id": None,
        "rms_norm_eps": 1e-06,
        "rope_parameters": {
            "mrope_interleaved": True,
            "mrope_section": [24, 20, 20],
            "rope_theta": 5000000,
            "rope_type": "default",
        },
        "use_cache": True,
        "vocab_size": 151936,
    },
    "tie_word_embeddings": False,
    "transformers_version": "5.8.0",
    "video_token_id": 151656,
    "vision_config": {
        "deepstack_visual_indexes": [8, 16, 24],
        "depth": 27,
        "dtype": "bfloat16",
        "hidden_act": "gelu_pytorch_tanh",
        "hidden_size": 1152,
        "in_channels": 3,
        "initializer_range": 0.02,
        "intermediate_size": 4304,
        "model_type": "qwen3_vl_vision",
        "num_heads": 16,
        "num_position_embeddings": 2304,
        "out_hidden_size": 4096,
        "patch_size": 16,
        "spatial_merge_size": 2,
        "temporal_patch_size": 2,
    },
    "vision_end_token_id": 151653,
    "vision_start_token_id": 151652,
    "ideogram_fp8_weight_only": True,
}


class Ideogram4TextEncoder(nn.Module):
    """Qwen3-VL-8B-Instruct wrapper that extracts hidden states from specific layers."""

    def __init__(self, config_path: str = None, **kwargs) -> None:
        super().__init__()
        from transformers import AutoConfig, AutoModel
        if config_path is None:
            config_kwargs = {k: v for k, v in _DEFAULT_TEXT_ENCODER_CONFIG.items() if k != "model_type"}
            config = AutoConfig.for_model("qwen3_vl", **config_kwargs)
        else:
            config = AutoConfig.from_pretrained(config_path, trust_remote_code=True)
        self.model = AutoModel.from_config(config, trust_remote_code=True)
        self.config = config

    def load_state_dict(self, state_dict, strict=True, assign=False):
        if is_fp8_state_dict(state_dict):
            swap_linears_to_fp8(self.model, state_dict, torch.bfloat16)
            return self.model.load_state_dict(state_dict, strict=False, assign=assign)
        return self.model.load_state_dict(state_dict, strict=strict, assign=assign)

    def forward(
        self,
        token_ids: torch.Tensor,
        attention_mask: torch.Tensor,
        text_position_ids: torch.Tensor,
    ) -> torch.Tensor:
        """Extract hidden states from specific layers of Qwen3-VL.

        Args:
            token_ids: (B, L) token ids
            attention_mask: (B, L) attention mask
            text_position_ids: (B, L) position ids for text tokens

        Returns:
            (B, L, hidden_size * num_activation_layers) concatenated hidden states
        """
        from transformers.masking_utils import create_causal_mask

        language_model = self.model.language_model

        inputs_embeds = language_model.embed_tokens(token_ids)

        position_ids_4d = text_position_ids[None, ...].expand(4, text_position_ids.shape[0], -1)
        text_position_ids_4d = position_ids_4d[0]
        mrope_position_ids = position_ids_4d[1:]

        causal_mask = create_causal_mask(
            config=language_model.config,
            inputs_embeds=inputs_embeds,
            attention_mask=attention_mask,
            past_key_values=None,
            position_ids=text_position_ids_4d,
        )
        position_embeddings = language_model.rotary_emb(inputs_embeds, mrope_position_ids)

        tap_set = set(QWEN3_VL_ACTIVATION_LAYERS)
        captured: dict[int, torch.Tensor] = {}
        hidden_states = inputs_embeds
        for layer_idx, decoder_layer in enumerate(language_model.layers):
            hidden_states = decoder_layer(
                hidden_states,
                attention_mask=causal_mask,
                position_ids=text_position_ids_4d,
                past_key_values=None,
                position_embeddings=position_embeddings,
            )
            if layer_idx in tap_set:
                captured[layer_idx] = hidden_states

        selected = [captured[i] for i in QWEN3_VL_ACTIVATION_LAYERS]
        stacked = torch.stack(selected, dim=0)
        stacked = torch.permute(stacked, (1, 2, 3, 0))
        batch_size, seq_len, hidden_size = stacked.shape[:3]
        stacked = stacked.reshape(batch_size, seq_len, -1)

        text_mask = attention_mask.to(stacked.dtype).unsqueeze(-1)
        stacked = stacked * text_mask
        return stacked.to(torch.float32)