Add runtime/text_encoder_variants.py

runtime/text_encoder_variants.py

@@ -0,0 +1,199 @@
+"""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",
+]