From 52af944bdc515e2be4fb91c7ac9937486dc82774 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Sat, 18 Jul 2026 06:24:48 +0000 Subject: [PATCH] Test the pre-cast text-encoder load path Hermetic CPU coverage for diffusion_te_prequant: the checkpoint filename convention, family-table resolution by scheme and component with malformed entries skipped, resolution priority (path override, hosted repo, none) and the fp8-only scheme gate, the checkpoint validation matrix (wrong format, missing state_dict, wrong scheme, wrong component, wrong or missing base_model_id) with base case folding, the local-path allowlist refusal and missing-file fallback, and the assembly injection gating (mode, hosted entry, device support, family deny, load failure, successful injection). Also pins the te_prequant_repos field on both family dataclasses and that no family ships a hosted TE checkpoint until the campaign wires one. --- .../tests/test_diffusion_te_prequant.py | 237 ++++++++++++++++++ 1 file changed, 237 insertions(+) create mode 100644 studio/backend/tests/test_diffusion_te_prequant.py diff --git a/studio/backend/tests/test_diffusion_te_prequant.py b/studio/backend/tests/test_diffusion_te_prequant.py new file mode 100644 index 0000000000..dd474ed876 --- /dev/null +++ b/studio/backend/tests/test_diffusion_te_prequant.py @@ -0,0 +1,237 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Hermetic CPU tests for the pre-cast text-encoder load path. + +Mirrors tests/test_diffusion_prequant.py: resolution priority, checkpoint validation, +fallback behaviour, the local-path allowlist gate, and the pipeline-assembly injection +gating -- all without CUDA, the Hub, or a real transformers model.""" + +from __future__ import annotations + +import types + +import pytest + +import core.inference.diffusion_te_prequant as tpq +from core.inference.diffusion_te_prequant import ( + TE_PREQUANT_FORMAT, + TePrequantSource, + family_te_prequant_repo, + resolve_te_prequant_source, + te_prequant_pipe_kwargs, + te_prequant_repo_filename, +) + + +def _fam(te_prequant_repos = (), name = "ltx-2"): + return types.SimpleNamespace(name = name, te_prequant_repos = te_prequant_repos) + + +# ── resolution ─────────────────────────────────────────────────────────────── +def test_repo_filename_convention(): + assert ( + te_prequant_repo_filename("unsloth/LTX-2-FP8", "text_encoder", "fp8") + == "LTX-2-text_encoder-FP8.pt" + ) + assert ( + te_prequant_repo_filename("org/Some-Model-quantized", "text_encoder_2", "fp8") + == "Some-Model-text_encoder_2-FP8.pt" + ) + assert ( + te_prequant_repo_filename("org/PlainRepo", "text_encoder", "fp8") + == "PlainRepo-text_encoder-FP8.pt" + ) + + +def test_family_repo_by_scheme_and_component(): + fam = _fam( + te_prequant_repos = ( + ("fp8", "text_encoder", "org/hosted-fp8"), + ("fp8", "text_encoder_2", "org/hosted-2-fp8"), + ) + ) + assert family_te_prequant_repo(fam, "fp8", "text_encoder") == "org/hosted-fp8" + assert family_te_prequant_repo(fam, "fp8", "text_encoder_2") == "org/hosted-2-fp8" + assert family_te_prequant_repo(fam, "fp8", "text_encoder_3") is None + assert family_te_prequant_repo(fam, "int8", "text_encoder") is None + # A malformed entry is skipped, not fatal. + assert family_te_prequant_repo(_fam(te_prequant_repos = (("bad",),)), "fp8", "text_encoder") is None + # Families without the field resolve to None (both dataclasses default it, but a fake + # or an older family object must not break). + assert family_te_prequant_repo(types.SimpleNamespace(name = "x"), "fp8", "text_encoder") is None + + +def test_resolve_priority_and_scheme_gate(): + fam = _fam(te_prequant_repos = (("fp8", "text_encoder", "org/hosted-fp8"),)) + # Path override wins. + src = resolve_te_prequant_source(fam, "text_encoder", "fp8", path_override = "/tmp/te.pt") + assert src == TePrequantSource(kind = "path", location = "/tmp/te.pt", filename = None) + # Hosted repo second. + src = resolve_te_prequant_source(fam, "text_encoder", "fp8") + assert src.kind == "repo" and src.location == "org/hosted-fp8" + assert src.filename == "hosted-text_encoder-FP8.pt" + # Nothing configured -> None. + assert resolve_te_prequant_source(_fam(), "text_encoder", "fp8") is None + # v1 hosts the layerwise fp8 storage scheme only. + assert resolve_te_prequant_source(fam, "text_encoder", "int8") is None + assert resolve_te_prequant_source(fam, "text_encoder", "fp8_dynamic") is None + + +# ── checkpoint validation ──────────────────────────────────────────────────── +def _good_ckpt(scheme = "fp8", component = "text_encoder", base = "Lightricks/LTX-2"): + return { + "format": TE_PREQUANT_FORMAT, + "metadata": { + "scheme": scheme, + "component": component, + "base_model_id": base, + "te_class": "Gemma3ForConditionalGeneration", + }, + "state_dict": {"weight": object()}, + } + + +@pytest.mark.parametrize( + "mutate, reason", + [ + (lambda c: c.update(format = "other"), "format"), + (lambda c: c.pop("state_dict"), "state_dict"), + (lambda c: c["metadata"].update(scheme = "int8"), "scheme"), + (lambda c: c["metadata"].update(component = "text_encoder_2"), "component"), + (lambda c: c["metadata"].update(base_model_id = "other/repo"), "base"), + (lambda c: c["metadata"].pop("base_model_id"), "missing base"), + ], +) +def test_validate_rejects_mismatches(mutate, reason): + ckpt = _good_ckpt() + mutate(ckpt) + assert ( + tpq._validate_checkpoint(ckpt, "fp8", "text_encoder", "Lightricks/LTX-2", None) is False + ), reason + + +def test_validate_accepts_good_checkpoint_and_base_case_folding(): + assert tpq._validate_checkpoint(_good_ckpt(), "fp8", "text_encoder", "Lightricks/LTX-2", None) + # _same_base_model folds case like the DiT module. + assert tpq._validate_checkpoint(_good_ckpt(), "fp8", "text_encoder", "lightricks/ltx-2", None) + + +# ── loader fallback behaviour ──────────────────────────────────────────────── +def test_load_refuses_unallowlisted_local_path(monkeypatch, tmp_path): + from core.inference.diffusion_prequant import ALLOW_LOCAL_PREQUANT_PATH_ENV + + monkeypatch.delenv(ALLOW_LOCAL_PREQUANT_PATH_ENV, raising = False) + path = tmp_path / "te.pt" + path.write_bytes(b"x") + out = tpq.load_prequant_text_encoder( + "Lightricks/LTX-2", + "text_encoder", + TePrequantSource(kind = "path", location = str(path)), + dtype = None, + ) + assert out is None # refused, caller falls back to dense + + +def test_load_missing_file_returns_none(monkeypatch, tmp_path): + from core.inference.diffusion_prequant import ALLOW_LOCAL_PREQUANT_PATH_ENV + + monkeypatch.setenv(ALLOW_LOCAL_PREQUANT_PATH_ENV, str(tmp_path)) + out = tpq.load_prequant_text_encoder( + "Lightricks/LTX-2", + "text_encoder", + TePrequantSource(kind = "path", location = str(tmp_path / "absent.pt")), + dtype = None, + ) + assert out is None + + +# ── pipeline-assembly injection gating ─────────────────────────────────────── +def _target(): + return types.SimpleNamespace(device = "cuda", dtype = None) + + +def test_pipe_kwargs_empty_when_mode_not_fp8(monkeypatch): + fam = _fam(te_prequant_repos = (("fp8", "text_encoder", "org/hosted"),)) + for mode in (None, "", "off", "int8", "fp8_dynamic"): + assert te_prequant_pipe_kwargs( + fam, "Lightricks/LTX-2", te_quant_mode = mode, target = _target(), dtype = None + ) == {} + + +def test_pipe_kwargs_empty_without_hosted_entry(monkeypatch): + import core.inference.diffusion_precision as precision + + monkeypatch.setattr(precision, "te_quant_supported", lambda target, mode: True) + assert te_prequant_pipe_kwargs( + _fam(), "Lightricks/LTX-2", te_quant_mode = "fp8", target = _target(), dtype = None + ) == {} + + +def test_pipe_kwargs_empty_when_device_unsupported(monkeypatch): + import core.inference.diffusion_precision as precision + + fam = _fam(te_prequant_repos = (("fp8", "text_encoder", "org/hosted"),)) + monkeypatch.setattr(precision, "te_quant_supported", lambda target, mode: False) + assert te_prequant_pipe_kwargs( + fam, "Lightricks/LTX-2", te_quant_mode = "fp8", target = _target(), dtype = None + ) == {} + + +def test_pipe_kwargs_respects_family_deny(monkeypatch): + import core.inference.diffusion_precision as precision + + fam = _fam(te_prequant_repos = (("fp8", "text_encoder", "org/hosted"),)) + monkeypatch.setattr(precision, "te_quant_supported", lambda target, mode: True) + # The deny helper ships on the video branch's precision module; simulate it here. + monkeypatch.setattr( + precision, "_te_family_denied", lambda family, mode: family == "ltx-2", raising = False + ) + assert te_prequant_pipe_kwargs( + fam, "Lightricks/LTX-2", te_quant_mode = "fp8", target = _target(), dtype = None + ) == {} + + +def test_pipe_kwargs_injects_loaded_encoder(monkeypatch): + import core.inference.diffusion_precision as precision + + fam = _fam(te_prequant_repos = (("fp8", "text_encoder", "org/hosted"),)) + monkeypatch.setattr(precision, "te_quant_supported", lambda target, mode: True) + marker = object() + seen = {} + + def fake_load(base, component, source, **kw): + seen.update(base = base, component = component, source = source) + return marker + + monkeypatch.setattr(tpq, "load_prequant_text_encoder", fake_load) + out = te_prequant_pipe_kwargs( + fam, "Lightricks/LTX-2", te_quant_mode = "fp8", target = _target(), dtype = None + ) + assert out == {"text_encoder": marker} + assert seen["base"] == "Lightricks/LTX-2" + assert seen["source"].location == "org/hosted" + + +def test_pipe_kwargs_empty_when_load_fails(monkeypatch): + import core.inference.diffusion_precision as precision + + fam = _fam(te_prequant_repos = (("fp8", "text_encoder", "org/hosted"),)) + monkeypatch.setattr(precision, "te_quant_supported", lambda target, mode: True) + monkeypatch.setattr(tpq, "load_prequant_text_encoder", lambda *a, **k: None) + assert te_prequant_pipe_kwargs( + fam, "Lightricks/LTX-2", te_quant_mode = "fp8", target = _target(), dtype = None + ) == {} + + +# ── family field wiring ────────────────────────────────────────────────────── +def test_family_dataclasses_declare_te_prequant_field(): + from core.inference.diffusion_families import DiffusionFamily, detect_family + from core.inference.video_families import VideoFamily + + assert DiffusionFamily.__dataclass_fields__["te_prequant_repos"].default_factory is tuple + assert VideoFamily.__dataclass_fields__["te_prequant_repos"].default_factory is tuple + # No family ships a hosted TE checkpoint yet: the campaign wires entries after the + # artifacts are gate-validated and uploaded. + fam = detect_family("unsloth/FLUX.1-schnell-GGUF") + assert fam.te_prequant_repos == ()