# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0

"""Unit tests for VAE quantisation (``diffusion_vae_quant.py``).

Hermetic: torch + the diffusers / torchao casters are stubbed via ``sys.modules`` so
gating, the conv filter, and the apply path run without a GPU, real diffusers, or real
torchao. Mirrors tests/test_diffusion_precision.py's stubbing style.
"""

from __future__ import annotations

import sys
import types

import pytest

import core.inference.diffusion_vae_quant as vq
from core.inference.diffusion_vae_quant import (
    VAE_QUANT_AUTO,
    VAE_QUANT_FP8,
    VAE_QUANT_FP8_DYNAMIC,
    _cast_vae_fp8,
    _cast_vae_fp8_dynamic,
    normalize_vae_quant,
    quantize_vae,
    select_vae_quant_scheme,
    vae_quant_supported,
)


def _target(
    *,
    device = "cuda",
    dtype = "bfloat16",
    cc = (10, 0),
):
    return types.SimpleNamespace(device = device, dtype = dtype, _cc = cc)


class _Weight:
    """A stand-in for a conv / linear weight tensor: exposes ``.shape`` and ``.dim()``."""

    def __init__(self, shape):
        self.shape = shape

    def dim(self):
        return len(self.shape)


def _stub_torch(
    monkeypatch,
    *,
    with_fp8 = True,
    cc = (10, 0),
):
    torch = types.ModuleType("torch")
    torch.bfloat16 = "bfloat16"
    torch.float16 = "float16"
    if with_fp8:
        torch.float8_e4m3fn = "float8_e4m3fn"
    # The conv filter isinstance-checks nn.Linear / nn.Conv2d / nn.Conv3d, and the layerwise
    # caster reads torch.float8_e4m3fn, so the stub torch must expose both.
    torch.nn = types.SimpleNamespace(
        Linear = type("Linear", (), {}),
        Conv2d = type("Conv2d", (), {}),
        Conv3d = type("Conv3d", (), {}),
    )
    torch.cuda = types.SimpleNamespace(
        get_device_capability = lambda *a: cc,
        synchronize = lambda *a, **k: None,
        is_available = lambda: True,
    )
    monkeypatch.setitem(sys.modules, "torch", torch)
    return torch


def _stub_torchao(monkeypatch, captured):
    # torchao's fp8_dynamic conv config + quantize_. Records the (config, filter_fn) so the
    # PerTensor granularity and the conv filter closure can be asserted.
    tq = types.ModuleType("torchao.quantization")
    tq.quantize_ = lambda module, config, filter_fn = None: captured.update(
        module = module, config = config, filter_fn = filter_fn
    )
    tq.Float8DynamicActivationFloat8WeightConfig = lambda granularity = None: ("fp8dyn", granularity)
    tq.PerTensor = lambda: "pertensor"
    monkeypatch.setitem(sys.modules, "torchao.quantization", tq)
    return tq


def _stub_diffusers(monkeypatch, recorder):
    hooks = types.ModuleType("diffusers.hooks")
    casting = types.ModuleType("diffusers.hooks.layerwise_casting")
    casting.DEFAULT_SKIP_MODULES_PATTERN = ("norm",)
    hooks.apply_layerwise_casting = lambda module, **kw: recorder.append(("fp8", module, kw))
    monkeypatch.setitem(sys.modules, "diffusers.hooks", hooks)
    monkeypatch.setitem(sys.modules, "diffusers.hooks.layerwise_casting", casting)


def _stub_capability(monkeypatch, cc):
    """Stub the transformer module's ``_capability`` that select_vae_quant_scheme imports."""
    dtq = types.ModuleType("core.inference.diffusion_transformer_quant")
    dtq._capability = lambda: cc
    monkeypatch.setitem(sys.modules, "core.inference.diffusion_transformer_quant", dtq)
    return dtq


def _allow_vae(monkeypatch, allowed):
    """Force vae_quant_supported to accept only ``allowed`` (simulates the hardware gate)."""
    monkeypatch.setattr(vq, "vae_quant_supported", lambda target, mode: mode in allowed)


# ── normalisation ─────────────────────────────────────────────────────────────


def test_normalize_vae_quant():
    assert normalize_vae_quant(None) is None
    assert normalize_vae_quant("") is None
    assert normalize_vae_quant("none") is None
    # "off" disables (like the TE / transformer normalisers) -> dense.
    assert normalize_vae_quant("off") is None
    # "auto" passes through for select_vae_quant_scheme to resolve.
    assert normalize_vae_quant("AUTO") == VAE_QUANT_AUTO
    assert normalize_vae_quant("FP8") == VAE_QUANT_FP8
    # Hyphens fold to underscores so "fp8-dynamic" is accepted.
    assert normalize_vae_quant("FP8-Dynamic") == VAE_QUANT_FP8_DYNAMIC
    # int8 / nvfp4 have no VAE scheme -> rejected.
    with pytest.raises(ValueError):
        normalize_vae_quant("int8")
    with pytest.raises(ValueError):
        normalize_vae_quant("nvfp4")


# ── gating ────────────────────────────────────────────────────────────────────


def test_vae_quant_supported_fp8_requires_cuda_bf16_and_fp8(monkeypatch):
    _stub_torch(monkeypatch, with_fp8 = True, cc = (8, 9))
    assert vae_quant_supported(_target(), VAE_QUANT_FP8) is True
    assert vae_quant_supported(_target(device = "cpu"), VAE_QUANT_FP8) is False
    assert vae_quant_supported(_target(dtype = "float16"), VAE_QUANT_FP8) is False
    # No fp8 dtype at all -> unsupported.
    _stub_torch(monkeypatch, with_fp8 = False, cc = (8, 9))
    assert vae_quant_supported(_target(), VAE_QUANT_FP8) is False


def test_vae_quant_supported_fp8_dynamic_requires_sm89(monkeypatch):
    # Compute fp8 conv (torch._scaled_mm) needs fp8-GEMM silicon: Ada sm_89+ / Hopper / Blackwell.
    _stub_torch(monkeypatch, cc = (8, 9))
    assert vae_quant_supported(_target(), VAE_QUANT_FP8_DYNAMIC) is True
    _stub_torch(monkeypatch, cc = (9, 0))
    assert vae_quant_supported(_target(), VAE_QUANT_FP8_DYNAMIC) is True
    # Ampere (8.0) has no fp8 GEMM.
    _stub_torch(monkeypatch, cc = (8, 0))
    assert vae_quant_supported(_target(), VAE_QUANT_FP8_DYNAMIC) is False


# ── auto ladder (select_vae_quant_scheme) ───────────────────────────────────────


def test_select_datacenter_uses_layerwise_fp8(monkeypatch):
    # ``auto`` engages layerwise fp8 ONLY (fp8_dynamic is out of the auto ladder). Even on
    # fp8-GEMM silicon it resolves to fp8, and the fp8_dynamic conv probe is never consulted.
    _stub_capability(monkeypatch, (10, 0))
    _allow_vae(monkeypatch, {VAE_QUANT_FP8_DYNAMIC, VAE_QUANT_FP8})
    monkeypatch.setattr(
        vq, "_vae_fp8_dynamic_probe", lambda *a: pytest.fail("auto must not probe fp8_dynamic")
    )
    assert select_vae_quant_scheme(_target(), "auto", family = "flux.1") == VAE_QUANT_FP8


def test_select_offload_uses_layerwise_fp8(monkeypatch):
    # Under offload ``auto`` still resolves to layerwise fp8 (the storage-only scheme that
    # survives Module.to()); the fp8_dynamic conv probe is never consulted.
    _stub_capability(monkeypatch, (10, 0))
    _allow_vae(monkeypatch, {VAE_QUANT_FP8_DYNAMIC, VAE_QUANT_FP8})
    monkeypatch.setattr(
        vq, "_vae_fp8_dynamic_probe", lambda *a: pytest.fail("probe must not run under offload")
    )
    assert select_vae_quant_scheme(_target(), "auto", offload_active = True) == VAE_QUANT_FP8


def test_select_force_fp32_stays_dense(monkeypatch):
    # A force-fp32 (Wan) family never quantises, for auto or an explicit request.
    _stub_capability(monkeypatch, (10, 0))
    _allow_vae(monkeypatch, {VAE_QUANT_FP8_DYNAMIC, VAE_QUANT_FP8})
    assert select_vae_quant_scheme(_target(), "auto", force_fp32 = True) is None
    assert select_vae_quant_scheme(_target(), "fp8", force_fp32 = True) is None


def test_select_family_deny_skips_scheme(monkeypatch):
    # The only scheme ``auto`` walks is layerwise fp8, so denying fp8 for a family leaves it
    # dense (None). (This is the SDXL case in the real deny list: fp8 marginal -> stay dense.)
    _stub_capability(monkeypatch, (10, 0))
    _allow_vae(monkeypatch, {VAE_QUANT_FP8_DYNAMIC, VAE_QUANT_FP8})
    monkeypatch.setattr(vq, "_vae_fp8_dynamic_probe", lambda device, ndim = 2: True)
    monkeypatch.setattr(vq, "_VAE_FAMILY_SCHEME_DENY", {"badfam": frozenset({VAE_QUANT_FP8})})
    assert select_vae_quant_scheme(_target(), "auto", family = "badfam") is None


def test_select_no_capability_is_none(monkeypatch):
    _stub_capability(monkeypatch, None)
    _allow_vae(monkeypatch, {VAE_QUANT_FP8_DYNAMIC, VAE_QUANT_FP8})
    assert select_vae_quant_scheme(_target(), "auto") is None


def test_select_support_gate_uses_fp8(monkeypatch):
    # fp8-capable silicon with fp8 supported -> ``auto`` resolves to layerwise fp8.
    _stub_capability(monkeypatch, (8, 9))
    _allow_vae(monkeypatch, {VAE_QUANT_FP8})
    assert select_vae_quant_scheme(_target(), "auto") == VAE_QUANT_FP8


def test_select_auto_never_uses_fp8_dynamic(monkeypatch):
    # Even with fp8_dynamic hardware-supported and its conv probe passing, ``auto`` stays on
    # layerwise fp8: fp8_dynamic is deliberately kept out of the auto ladder (explicit opt-in).
    _stub_capability(monkeypatch, (10, 0))
    _allow_vae(monkeypatch, {VAE_QUANT_FP8_DYNAMIC, VAE_QUANT_FP8})
    monkeypatch.setattr(vq, "_vae_fp8_dynamic_probe", lambda device, ndim = 2: True)
    assert select_vae_quant_scheme(_target(), "auto") == VAE_QUANT_FP8


def test_select_explicit_passthrough(monkeypatch):
    # An explicit request is returned as-is (quantize_vae re-gates it); no ladder walk,
    # so no transformer-module stub is needed.
    assert select_vae_quant_scheme(_target(), "fp8") == VAE_QUANT_FP8
    assert select_vae_quant_scheme(_target(), "fp8_dynamic") == VAE_QUANT_FP8_DYNAMIC
    assert select_vae_quant_scheme(_target(), None) is None
    assert select_vae_quant_scheme(_target(), "none") is None


def test_select_explicit_family_deny_returns_none(monkeypatch):
    monkeypatch.setattr(vq, "_VAE_FAMILY_SCHEME_DENY", {"badfam": frozenset({VAE_QUANT_FP8})})
    assert select_vae_quant_scheme(_target(), "fp8", family = "badfam") is None


# ── conv-aware filter (_cast_vae_fp8_dynamic) ───────────────────────────────────


def test_fp8_dynamic_conv_filter(monkeypatch):
    # The real filter closure the PerTensor fp8 conv config receives: %16-channel Conv2d/Conv3d
    # and Linear quantise; off-%16 channels and the conv_out / norm_out head stay dense.
    torch = _stub_torch(monkeypatch)
    captured: dict = {}
    _stub_torchao(monkeypatch, captured)
    _cast_vae_fp8_dynamic(object(), _target())
    # PerTensor granularity (NOT the DiT's per-row) is what the config was built with.
    assert captured["config"] == ("fp8dyn", "pertensor")
    ff = captured["filter_fn"]
    nn = torch.nn

    def _mod(
        cls,
        shape,
        kernel_size = None,
    ):
        m = cls()
        m.weight = _Weight(shape)
        if kernel_size is not None:
            m.kernel_size = kernel_size
        return m

    # Conv2d (4D) / Conv3d (5D) with both channel dims a multiple of 16 quantise.
    assert ff(_mod(nn.Conv2d, (128, 128, 3, 3), (3, 3)), "decoder.up.0.resnets.0.conv1") is True
    assert ff(_mod(nn.Conv3d, (64, 64, 3, 3, 3), (3, 3, 3)), "decoder.mid_block.conv3d") is True
    # nn.Linear (mid-block attention projections) also quantise (no kernel_size attr).
    assert ff(_mod(nn.Linear, (512, 512)), "decoder.mid_block.attentions.0.to_q") is True
    # POINTWISE (1x1 / 1x1x1) convs excluded even at %16 channels: torchao 0.17's fp8 conv
    # kernel rejects them ("Activation and filter channels must match") -> crash at decode.
    assert (
        ff(_mod(nn.Conv2d, (128, 128, 1, 1), (1, 1)), "decoder.mid_block.attentions.0.proj_conv")
        is False
    )
    assert ff(_mod(nn.Conv3d, (64, 64, 1, 1, 1), (1, 1, 1)), "decoder.time_mix.conv") is False
    # Channels not a multiple of 16 excluded (torchao would skip them regardless): the RGB
    # in/out head (C=3) and any off-16 dim.
    assert ff(_mod(nn.Conv2d, (128, 3, 3, 3), (3, 3)), "encoder.conv_in") is False
    assert ff(_mod(nn.Conv2d, (24, 128, 3, 3), (3, 3)), "decoder.up.1.upsamplers.0.conv") is False
    # conv_out / proj_out / norm_out excluded by NAME even with %16 channels.
    assert ff(_mod(nn.Conv2d, (16, 128, 3, 3)), "decoder.conv_out") is False
    assert ff(_mod(nn.Conv2d, (128, 128, 3, 3)), "decoder.conv_norm_out") is False
    assert ff(_mod(nn.Conv2d, (128, 128, 3, 3)), "decoder.norm_out.conv") is False
    assert ff(_mod(nn.Linear, (512, 512)), "decoder.proj_out") is False

    # A non-conv/linear module (e.g. a GroupNorm) is excluded outright.
    class _GroupNorm:
        pass

    gn = _GroupNorm()
    gn.weight = _Weight((128,))
    assert ff(gn, "decoder.mid_block.resnets.0.norm1") is False
    # A weight with dim() < 2 (a 1D param) is excluded.
    assert ff(_mod(nn.Conv2d, (128,)), "decoder.some.bias_only") is False


def test_cast_vae_fp8_layerwise_skips_head_and_norms(monkeypatch):
    # The layerwise storage cast passes the decoder head + norm tokens through to diffusers'
    # skip_modules_pattern (on top of the diffusers default), so they stay dense.
    _stub_torch(monkeypatch)
    recorder: list = []
    _stub_diffusers(monkeypatch, recorder)
    vae = object()
    _cast_vae_fp8(vae, _target())
    assert len(recorder) == 1
    _, mod, kw = recorder[0]
    assert mod is vae
    assert kw["storage_dtype"] == "float8_e4m3fn"
    assert kw["compute_dtype"] == "bfloat16"
    skip = kw["skip_modules_pattern"]
    # The diffusers default is preserved and the keep-dense tokens are appended.
    assert "norm" in skip
    for tok in ("conv_out", "proj_out", "conv_norm_out", "norm_out"):
        assert tok in skip


# ── apply (quantize_vae) ────────────────────────────────────────────────────────


def test_quantize_vae_disabled_returns_none(monkeypatch):
    pipe = types.SimpleNamespace(vae = object())
    assert quantize_vae(pipe, _target(), mode = None) is None
    assert quantize_vae(pipe, _target(), mode = "none") is None


def test_quantize_vae_force_fp32_stays_dense(monkeypatch):
    # A force-fp32 (Wan) family never casts, for an explicit scheme or auto.
    _stub_torch(monkeypatch, cc = (10, 0))
    monkeypatch.setattr(vq, "_cast_vae_fp8_dynamic", lambda v, t: pytest.fail("must not cast"))
    monkeypatch.setattr(vq, "_cast_vae_fp8", lambda v, t: pytest.fail("must not cast"))
    pipe = types.SimpleNamespace(vae = object())
    assert quantize_vae(pipe, _target(), mode = "fp8", force_fp32 = True) is None
    assert quantize_vae(pipe, _target(), mode = "auto", force_fp32 = True) is None


def test_quantize_vae_offload_skips_fp8_dynamic(monkeypatch):
    # Explicit fp8_dynamic under offload is skipped (torchao tensors reject Module.to());
    # layerwise fp8 still engages.
    _stub_torch(monkeypatch, cc = (10, 0))
    _allow_vae(monkeypatch, {VAE_QUANT_FP8_DYNAMIC, VAE_QUANT_FP8})
    monkeypatch.setattr(
        vq, "_cast_vae_fp8_dynamic", lambda v, t: pytest.fail("torchao must not run under offload")
    )
    pipe = types.SimpleNamespace(vae = object())
    assert quantize_vae(pipe, _target(), mode = "fp8_dynamic", offload_active = True) is None
    fp8_calls: list = []
    monkeypatch.setattr(vq, "_cast_vae_fp8", lambda v, t: fp8_calls.append(v))
    assert quantize_vae(pipe, _target(), mode = "fp8", offload_active = True) == VAE_QUANT_FP8
    assert len(fp8_calls) == 1


def test_quantize_vae_unsupported_hw_is_noop(monkeypatch):
    # An explicit scheme on hardware that does not support it applies nothing.
    _stub_torch(monkeypatch, cc = (8, 0))
    _allow_vae(monkeypatch, set())
    monkeypatch.setattr(vq, "_cast_vae_fp8", lambda v, t: pytest.fail("must not cast"))
    pipe = types.SimpleNamespace(vae = object())
    assert quantize_vae(pipe, _target(), mode = "fp8") is None


def test_quantize_vae_none_vae_is_noop(monkeypatch):
    # A pipeline with no VAE attribute is a best-effort no-op even when the mode is supported.
    _stub_torch(monkeypatch, cc = (10, 0))
    _allow_vae(monkeypatch, {VAE_QUANT_FP8})
    pipe = types.SimpleNamespace()  # no .vae
    assert quantize_vae(pipe, _target(), mode = "fp8") is None


def test_quantize_vae_explicit_fp8_applies(monkeypatch):
    _stub_torch(monkeypatch, cc = (10, 0))
    _allow_vae(monkeypatch, {VAE_QUANT_FP8})
    calls: list = []
    monkeypatch.setattr(vq, "_cast_vae_fp8", lambda v, t: calls.append(v))
    vae = object()
    pipe = types.SimpleNamespace(vae = vae)
    assert quantize_vae(pipe, _target(), mode = "fp8") == VAE_QUANT_FP8
    assert calls == [vae]


def test_quantize_vae_tolerates_caster_failure(monkeypatch):
    # The caster raising leaves the VAE dense (best-effort) -> None.
    _stub_torch(monkeypatch, cc = (10, 0))
    _allow_vae(monkeypatch, {VAE_QUANT_FP8})

    def _boom(v, t):
        raise RuntimeError("fp8 unsupported for this layer")

    monkeypatch.setattr(vq, "_cast_vae_fp8", _boom)
    pipe = types.SimpleNamespace(vae = object())
    assert quantize_vae(pipe, _target(), mode = "fp8") is None


def test_quantize_vae_auto_resolves_and_applies(monkeypatch):
    # End-to-end: mode="auto" resolves via the ladder (layerwise fp8) then applies that caster.
    _stub_torch(monkeypatch, cc = (10, 0))
    _stub_capability(monkeypatch, (10, 0))
    _allow_vae(monkeypatch, {VAE_QUANT_FP8_DYNAMIC, VAE_QUANT_FP8})
    monkeypatch.setattr(
        vq, "_cast_vae_fp8_dynamic", lambda v, t: pytest.fail("auto must not use fp8_dynamic")
    )
    calls: list = []
    monkeypatch.setattr(vq, "_cast_vae_fp8", lambda v, t: calls.append(v))
    vae = object()
    pipe = types.SimpleNamespace(vae = vae)
    # large VAE -> clears the size gate
    monkeypatch.setattr(vq, "_vae_param_bytes", lambda v: 3_000_000_000)
    assert quantize_vae(pipe, _target(), mode = "auto", family = "flux.1") == VAE_QUANT_FP8
    assert calls == [vae]


def test_quantize_vae_auto_size_gate_skips_small(monkeypatch):
    # auto leaves a small (image) VAE dense: halving ~0.2GB saves ~nothing and only slows decode.
    _stub_torch(monkeypatch, cc = (10, 0))
    _stub_capability(monkeypatch, (10, 0))
    _allow_vae(monkeypatch, {VAE_QUANT_FP8_DYNAMIC, VAE_QUANT_FP8})
    monkeypatch.setattr(vq, "_vae_param_bytes", lambda v: 200_000_000)  # ~0.2 GB image VAE
    monkeypatch.setattr(vq, "_cast_vae_fp8", lambda v, t: pytest.fail("small VAE must stay dense"))
    monkeypatch.setattr(
        vq, "_cast_vae_fp8_dynamic", lambda v, t: pytest.fail("small VAE must stay dense")
    )
    pipe = types.SimpleNamespace(vae = object())
    assert quantize_vae(pipe, _target(), mode = "auto", family = "flux.1") is None


def test_quantize_vae_auto_size_gate_allows_large(monkeypatch):
    # auto quantises a large (video Conv3d) VAE: halving ~2.5GB saves ~1.2GB at ~2% decode.
    _stub_torch(monkeypatch, cc = (10, 0))
    _stub_capability(monkeypatch, (10, 0))
    _allow_vae(monkeypatch, {VAE_QUANT_FP8_DYNAMIC, VAE_QUANT_FP8})
    monkeypatch.setattr(vq, "_vae_param_bytes", lambda v: 2_500_000_000)  # ~2.5 GB video VAE
    calls: list = []
    monkeypatch.setattr(vq, "_cast_vae_fp8", lambda v, t: calls.append(v))
    pipe = types.SimpleNamespace(vae = object())
    assert quantize_vae(pipe, _target(), mode = "auto", family = "hunyuanvideo-1.5") == VAE_QUANT_FP8
    assert len(calls) == 1


def test_quantize_vae_explicit_bypasses_size_gate(monkeypatch):
    # An explicit request quantises even a small VAE -- the user opted in directly.
    _stub_torch(monkeypatch, cc = (10, 0))
    _allow_vae(monkeypatch, {VAE_QUANT_FP8})
    monkeypatch.setattr(vq, "_vae_param_bytes", lambda v: 200_000_000)  # small, but explicit
    calls: list = []
    monkeypatch.setattr(vq, "_cast_vae_fp8", lambda v, t: calls.append(v))
    pipe = types.SimpleNamespace(vae = object())
    assert quantize_vae(pipe, _target(), mode = "fp8", family = "flux.1") == VAE_QUANT_FP8
    assert len(calls) == 1


def test_quantize_vae_explicit_fp8_dynamic_probe_gates(monkeypatch):
    # An explicit fp8_dynamic request runs the conv smoke probe: a build whose torchao lacks a
    # working fp8 conv path (probe False) stays dense; a passing probe applies the caster.
    _stub_torch(monkeypatch, cc = (10, 0))
    _allow_vae(monkeypatch, {VAE_QUANT_FP8_DYNAMIC, VAE_QUANT_FP8})
    monkeypatch.setattr(vq, "_vae_fp8_dynamic_probe", lambda device, ndim = 2: False)
    monkeypatch.setattr(
        vq, "_cast_vae_fp8_dynamic", lambda v, t: pytest.fail("probe failed: must not cast")
    )
    pipe = types.SimpleNamespace(vae = object())
    assert quantize_vae(pipe, _target(), mode = "fp8_dynamic", family = "flux.2-klein") is None
    monkeypatch.setattr(vq, "_vae_fp8_dynamic_probe", lambda device, ndim = 2: True)
    calls: list = []
    monkeypatch.setattr(vq, "_cast_vae_fp8_dynamic", lambda v, t: calls.append(v))
    assert (
        quantize_vae(pipe, _target(), mode = "fp8_dynamic", family = "flux.2-klein")
        == VAE_QUANT_FP8_DYNAMIC
    )
    assert len(calls) == 1


def test_real_family_deny_list_policy(monkeypatch):
    # The shipped _VAE_FAMILY_SCHEME_DENY (from the B200 sweep), exercised through the real
    # select path -- no deny-list monkeypatch. Confirms the per-family auto/explicit outcomes.
    _stub_capability(monkeypatch, (10, 0))
    _allow_vae(monkeypatch, {VAE_QUANT_FP8_DYNAMIC, VAE_QUANT_FP8})
    # SDXL denies BOTH schemes -> auto stays dense (its layerwise fp8 was marginal).
    assert select_vae_quant_scheme(_target(), "auto", family = "sdxl") is None
    assert select_vae_quant_scheme(_target(), "fp8", family = "sdxl") is None
    # Qwen-Image: auto -> layerwise fp8 (safe); explicit fp8_dynamic refused (catastrophic).
    assert select_vae_quant_scheme(_target(), "auto", family = "qwen-image") == VAE_QUANT_FP8
    assert select_vae_quant_scheme(_target(), "fp8_dynamic", family = "qwen-image") is None
    # FLUX.1 / LTX-2 keep layerwise fp8 on auto but deny explicit fp8_dynamic.
    assert select_vae_quant_scheme(_target(), "auto", family = "ltx-2") == VAE_QUANT_FP8
    assert select_vae_quant_scheme(_target(), "fp8_dynamic", family = "flux.1") is None
    # FLUX.2 / Hunyuan keep fp8_dynamic available as an explicit opt-in (measured in-bar).
    assert (
        select_vae_quant_scheme(_target(), "fp8_dynamic", family = "flux.2-klein")
        == VAE_QUANT_FP8_DYNAMIC
    )
    assert (
        select_vae_quant_scheme(_target(), "fp8_dynamic", family = "hunyuanvideo-1.5")
        == VAE_QUANT_FP8_DYNAMIC
    )


# ── conv-dimensionality probe gating ─────────────────────────────────────────────


def _conv_vae(torch, *ndims):
    """A fake VAE whose .modules() yields stub Conv2d/Conv3d instances for ``ndims``."""
    mods = [(torch.nn.Conv3d if n == 3 else torch.nn.Conv2d)() for n in ndims]
    return types.SimpleNamespace(modules = lambda: iter(mods))


def test_vae_conv_ndims(monkeypatch):
    torch = _stub_torch(monkeypatch)
    assert vq._vae_conv_ndims(_conv_vae(torch, 2)) == (2,)
    assert vq._vae_conv_ndims(_conv_vae(torch, 3)) == (3,)
    assert vq._vae_conv_ndims(_conv_vae(torch, 2, 3)) == (2, 3)
    # A VAE that cannot be inspected (no .modules()) falls back to the 2D probe.
    assert vq._vae_conv_ndims(object()) == (2,)


def test_quantize_vae_fp8_dynamic_probes_conv3d_for_video_vae(monkeypatch):
    # A Conv3d (video) VAE must pass the 3D conv probe: a build whose Conv2d path works but
    # Conv3d is broken would else crash at the first decode. Probe ok for 2D but not 3D -> dense.
    torch = _stub_torch(monkeypatch, cc = (10, 0))
    _allow_vae(monkeypatch, {VAE_QUANT_FP8_DYNAMIC, VAE_QUANT_FP8})
    probed: list = []

    def _probe(device, ndim = 2):
        probed.append(ndim)
        return ndim == 2

    monkeypatch.setattr(vq, "_vae_fp8_dynamic_probe", _probe)
    monkeypatch.setattr(
        vq, "_cast_vae_fp8_dynamic", lambda v, t: pytest.fail("3D probe failed: must not cast")
    )
    pipe = types.SimpleNamespace(vae = _conv_vae(torch, 2, 3))
    assert quantize_vae(pipe, _target(), mode = "fp8_dynamic", family = "hunyuanvideo-1.5") is None
    assert 3 in probed
    # A build whose 3D path also works casts the video VAE.
    monkeypatch.setattr(vq, "_vae_fp8_dynamic_probe", lambda device, ndim = 2: True)
    pipe = types.SimpleNamespace(vae = _conv_vae(torch, 2, 3))
    calls: list = []
    monkeypatch.setattr(vq, "_cast_vae_fp8_dynamic", lambda v, t: calls.append(v))
    assert (
        quantize_vae(pipe, _target(), mode = "fp8_dynamic", family = "hunyuanvideo-1.5")
        == VAE_QUANT_FP8_DYNAMIC
    )
    assert len(calls) == 1


# ── partial in-place quant detection (fails the load, not a silent dense report) ───
class _TorchaoLikeTensor:
    """Detection keys on the tensor class's module path ("torchao" in __module__)."""


_TorchaoLikeTensor.__module__ = "torchao.dtypes.affine_quantized_tensor"


class _PartiallyQuantizedVae:
    def __init__(self):
        self._swapped = False

    def named_parameters(self):
        if self._swapped:
            yield ("decoder.up_blocks.0.conv.weight", _TorchaoLikeTensor())
        yield ("decoder.up_blocks.1.conv.weight", types.SimpleNamespace())


def test_quantize_vae_partial_cast_failure_fails_load(monkeypatch):
    # fp8_dynamic's quantize_ swaps weights module-by-module: a mid-pass failure that left
    # torchao params behind must raise instead of reporting a dense fallback.
    _stub_torch(monkeypatch, cc = (10, 0))
    _allow_vae(monkeypatch, {VAE_QUANT_FP8})

    def _swap_one_then_boom(v, t):
        v._swapped = True
        raise RuntimeError("mid-pass conv failure")

    monkeypatch.setattr(vq, "_cast_vae_fp8", _swap_one_then_boom)
    pipe = types.SimpleNamespace(vae = _PartiallyQuantizedVae())
    with pytest.raises(RuntimeError, match = "partially quantized"):
        quantize_vae(pipe, _target(), mode = "fp8")


# ── layerwise fp8 partial mutation (torchao detector is blind to diffusers hooks) ──
class _LayerwiseCastVae:
    """A VAE an apply_layerwise_casting pass mutated (installed an fp8-storage upcast hook on a
    submodule) before the caster raised. It carries NO torchao params, so raise_if_partially_
    quantized would miss the partial state -- _has_layerwise_casting must catch it."""

    def __init__(self):
        registry = types.SimpleNamespace(
            get_hook = lambda name: object() if name == "layerwise_casting" else None
        )
        self._sub = types.SimpleNamespace(_diffusers_hook = registry)

    def modules(self):
        return [self, self._sub]

    def named_parameters(self):
        return iter(())


def test_quantize_vae_layerwise_partial_cast_fails_load(monkeypatch):
    _stub_torch(monkeypatch, cc = (10, 0))
    _allow_vae(monkeypatch, {VAE_QUANT_FP8})

    def _boom(v, t):
        raise RuntimeError("layerwise casting failed mid-pass")

    monkeypatch.setattr(vq, "_cast_vae_fp8", _boom)
    pipe = types.SimpleNamespace(vae = _LayerwiseCastVae())
    with pytest.raises(RuntimeError, match = "leftover fp8 hooks"):
        quantize_vae(pipe, _target(), mode = "fp8")


def test_quantize_vae_clean_layerwise_failure_stays_dense(monkeypatch):
    # A failure with NO leftover hook (raised before mutating anything) still falls back to dense,
    # preserving the storage-only fp8 contract -- the fail-closed path is scoped to real mutation.
    _stub_torch(monkeypatch, cc = (10, 0))
    _allow_vae(monkeypatch, {VAE_QUANT_FP8})

    class _CleanVae:
        def modules(self):
            return [self]

        def named_parameters(self):
            return iter(())

    def _boom(v, t):
        raise RuntimeError("fp8 unsupported before any mutation")

    monkeypatch.setattr(vq, "_cast_vae_fp8", _boom)
    pipe = types.SimpleNamespace(vae = _CleanVae())
    assert quantize_vae(pipe, _target(), mode = "fp8") is None