runtime/text_encoder_variants.py

"""Test-only loaders for alternate Qwen3-VL text-encoder artifacts.

The released mixed NVFP4/FP8 loader remains immutable.  This module reuses its
validated tensor mapping and module implementation while accepting the
ComfyUI scaled-FP8 policy (252 language projections, all FP8 E4M3).
"""

from __future__ import annotations

from collections import Counter
from pathlib import Path
from typing import Any

import torch
from accelerate import init_empty_weights
from safetensors import safe_open
from transformers import AutoConfig, AutoProcessor, AutoTokenizer


EXPECTED_FP8_SHA256 = (
    "54bd5144df0bbc25dd6ccadfcb826b521445a1b06ae5a42570bdd2974ca87094"
)
EXPECTED_FP8_PROJECTION_COUNT = 252


def _base_loader() -> Any:
    import quant_text_encoder

    return quant_text_encoder


def _install_scaled_fp8_linears(
    hf_module: torch.nn.Module,
    artifact_path: Path,
) -> dict[str, Any]:
    base = _base_loader()
    format_counts: Counter[str] = Counter()
    installed: list[str] = []
    storage_bytes = 0
    original_bf16_bytes = 0
    full_precision_matrix_mult_counts: Counter[bool] = Counter()

    with safe_open(str(artifact_path), framework="pt", device="cpu") as handle:
        config_keys = sorted(
            key for key in handle.keys() if key.endswith(".comfy_quant")
        )
        for config_key in config_keys:
            layer_key = config_key.removesuffix(".comfy_quant")
            quant_config = base._decode_quant_config(
                handle.get_tensor(config_key)
            )
            quant_format = quant_config["format"]
            if quant_format != "float8_e4m3fn":
                raise RuntimeError(
                    f"{layer_key}: expected float8_e4m3fn, got {quant_format}"
                )
            full_precision_matrix_mult_counts[
                bool(quant_config.get("full_precision_matrix_mult", False))
            ] += 1
            module_name = base._artifact_layer_to_hf_module(layer_key)
            parent, leaf = base._resolve_parent(hf_module, module_name)
            original = getattr(parent, leaf)
            if not isinstance(original, torch.nn.Linear):
                raise TypeError(
                    f"{module_name}: expected torch.nn.Linear, got "
                    f"{type(original).__name__}"
                )
            if original.bias is not None:
                raise ValueError(f"{module_name}: quantized projection has a bias")

            replacement = base.PublishedQuantLinear(
                in_features=original.in_features,
                out_features=original.out_features,
                quant_format=quant_format,
                qdata=handle.get_tensor(f"{layer_key}.weight"),
                weight_scale=handle.get_tensor(f"{layer_key}.weight_scale"),
                weight_scale_2=None,
            )
            setattr(parent, leaf, replacement)
            format_counts[quant_format] += 1
            installed.append(module_name)
            original_bf16_bytes += (
                original.in_features * original.out_features * 2
            )
            storage_bytes += replacement.qdata.nbytes
            storage_bytes += replacement.weight_scale.nbytes

    summary = {
        "installed_module_count": len(installed),
        "format_counts": dict(format_counts),
        "full_precision_matrix_mult_counts": {
            str(key).lower(): value
            for key, value in sorted(
                full_precision_matrix_mult_counts.items()
            )
        },
        "original_projection_bf16_bytes": original_bf16_bytes,
        "packed_projection_bytes": storage_bytes,
        "projection_saving_bytes": original_bf16_bytes - storage_bytes,
        "projection_saving_gib": (
            (original_bf16_bytes - storage_bytes) / float(1 << 30)
        ),
    }
    if (
        summary["installed_module_count"] != EXPECTED_FP8_PROJECTION_COUNT
        or summary["format_counts"]
        != {"float8_e4m3fn": EXPECTED_FP8_PROJECTION_COUNT}
        or summary["full_precision_matrix_mult_counts"] != {"false": 252}
    ):
        raise RuntimeError(f"unexpected scaled-FP8 policy: {summary}")
    return summary


def load_scaled_fp8_text_encoder(
    *,
    text_encoder_dir: str | Path,
    artifact_path: str | Path,
    tokenizer_max_length: int,
    dit_structure: dict[str, Any],
    use_packed_text_infer: bool,
    attn_type: str = "flash2",
) -> tuple[torch.nn.Module, dict[str, Any]]:
    """Build Mage's Qwen wrapper from the 252-projection scaled-FP8 file."""

    base = _base_loader()
    from mage_flow.models.modules.text_encoder import (
        CustomQwen3VLForConditionalGeneration,
        TextEncoder,
        _resolve_hf_attn_impl,
    )

    text_encoder_dir = Path(text_encoder_dir).resolve()
    artifact_path = Path(artifact_path).resolve()
    actual_sha256 = base.sha256(artifact_path)
    if actual_sha256 != EXPECTED_FP8_SHA256:
        raise RuntimeError(
            "scaled-FP8 text artifact SHA-256 mismatch: "
            f"{actual_sha256}"
        )

    config = AutoConfig.from_pretrained(
        str(text_encoder_dir),
        local_files_only=True,
    )
    hf_attn_implementation = _resolve_hf_attn_impl(attn_type)
    with init_empty_weights():
        hf_module = CustomQwen3VLForConditionalGeneration._from_config(
            config,
            attn_implementation=hf_attn_implementation,
            dtype=torch.bfloat16,
        )

    quant_summary = _install_scaled_fp8_linears(
        hf_module,
        artifact_path,
    )
    load_summary = base._load_nonquantized_weights(
        hf_module,
        artifact_path,
    )

    text_encoder = TextEncoder.__new__(TextEncoder)
    torch.nn.Module.__init__(text_encoder)
    text_encoder.model_name = str(text_encoder_dir)
    text_encoder.tokenizer_max_length = int(tokenizer_max_length)
    text_encoder.tokenizer = AutoTokenizer.from_pretrained(
        str(text_encoder_dir),
        local_files_only=True,
    )
    text_encoder.tokenizer.padding_side = "right"
    text_encoder.processor = AutoProcessor.from_pretrained(
        str(text_encoder_dir),
        local_files_only=True,
    )
    text_encoder.hf_module = hf_module.eval().requires_grad_(False)
    text_encoder.prompt_template_encode = ""
    text_encoder.prompt_template_encode_start_idx = 0
    text_encoder.dit_structure = dict(dit_structure)
    text_encoder.use_packed_text_infer = bool(use_packed_text_infer)
    text_encoder.eval().requires_grad_(False)

    return (
        text_encoder,
        {
            "artifact": str(artifact_path),
            "artifact_sha256": actual_sha256,
            "attention_backend": attn_type,
            "hf_attention_implementation": hf_attn_implementation,
            "quantized": quant_summary,
            "nonquantized": load_summary,
        },
    )


__all__ = [
    "EXPECTED_FP8_PROJECTION_COUNT",
    "EXPECTED_FP8_SHA256",
    "load_scaled_fp8_text_encoder",
]