hao-ai-lab-FastVideo/tests/local_tests/kandinsky6_sr/reference.py

# SPDX-License-Identifier: Apache-2.0
"""Loads the k6_video reference (``kandinsky_sr``) for the Kandinsky6 SR parity tests.

Local-only tooling: the reference package is not a FastVideo dependency.  It is found via
``$KANDINSKY_SR_SRC`` (the ``k6_video/src`` directory) or a ``k6_video`` checkout next to this repository.  On a
machine without CUDA / flash-attn / loguru small stubs are installed (per-sequence SDPA instead of flash-attn varlen,
no-op logger) and the base resolutions are shrunk, so the whole reference pipeline runs on CPU with tiny random
models.  On a GPU box with the real packages nothing is stubbed.

IMPORTANT: import FastVideo (``import fastvideo...``) BEFORE calling :func:`load_reference`; diffusers probes
``flash_attn`` at import time and chokes on a stub module.
"""
from __future__ import annotations

import copy
import importlib
import os
import sys
import types
from dataclasses import dataclass
from pathlib import Path
from typing import Any

import pytest
import torch

REPO_ROOT = Path(__file__).resolve().parents[3]
sys.path.insert(0, str(REPO_ROOT / "fastvideo" / "tests" / "stages" / "kandinsky6_sr"))

import k6_sr_tiny  # noqa: E402  (tiny configs shared with the CPU unit tests)

TINY_RESOLUTIONS = k6_sr_tiny.TINY_RESOLUTIONS


def find_reference_src() -> Path | None:
    """``$KANDINSKY_SR_SRC`` (``none`` disables the lookup, forcing the tests to skip) or ``../k6_video/src``."""
    if os.environ.get("KANDINSKY_SR_SRC", "").lower() == "none":
        return None
    candidates = [os.environ.get("KANDINSKY_SR_SRC"), str(REPO_ROOT.parent / "k6_video" / "src")]
    for candidate in candidates:
        if candidate and (Path(candidate) / "kandinsky_sr").is_dir():
            return Path(candidate)
    return None


def _varlen_sdpa(q, k, v, cu_q, cu_k):
    import torch.nn.functional as F

    outs = []
    for i in range(len(cu_q) - 1):
        qi = q[int(cu_q[i]):int(cu_q[i + 1])].transpose(0, 1).unsqueeze(0)
        ki = k[int(cu_k[i]):int(cu_k[i + 1])].transpose(0, 1).unsqueeze(0)
        vi = v[int(cu_k[i]):int(cu_k[i + 1])].transpose(0, 1).unsqueeze(0)
        outs.append(F.scaled_dot_product_attention(qi, ki, vi).squeeze(0).transpose(0, 1))
    return torch.cat(outs, dim=0)


def _install_stubs() -> None:
    os.environ.setdefault("TORCH_COMPILE_DISABLE", "1")
    os.environ.setdefault("TORCHDYNAMO_DISABLE", "1")
    if "flash_attn" not in sys.modules:
        try:
            importlib.import_module("flash_attn")
        except ImportError:
            mod = types.ModuleType("flash_attn")

            def varlen(q, k, v, cu_seqlens_q, cu_seqlens_k, max_seqlen_q, max_seqlen_k, dropout_p=0.0,
                       softmax_scale=None, causal=False, return_attn_probs=False, **_):
                out = _varlen_sdpa(q, k, v, cu_seqlens_q, cu_seqlens_k)
                return (out, torch.zeros(1), None) if return_attn_probs else out

            def qkvpacked(qkv, cu_seqlens, max_seqlen, dropout_p=0.0, softmax_scale=None, causal=False,
                          return_attn_probs=False, **_):
                q, k, v = qkv.unbind(dim=1)
                out = _varlen_sdpa(q, k, v, cu_seqlens, cu_seqlens)
                return (out, torch.zeros(1), None) if return_attn_probs else out

            mod.flash_attn_varlen_func = varlen
            mod.flash_attn_varlen_qkvpacked_func = qkvpacked
            sys.modules["flash_attn"] = mod
    if "loguru" not in sys.modules:
        try:
            importlib.import_module("loguru")
        except ImportError:
            mod = types.ModuleType("loguru")

            class _Logger:

                def __getattr__(self, name):
                    return lambda *a, **k: self

            mod.logger = _Logger()
            sys.modules["loguru"] = mod


@dataclass
class Reference:
    """Handles on the reference modules used by the parity tests."""

    constants: Any
    dit_mod: Any
    dx_dit_mod: Any
    kvae_mod: Any
    sr_pipeline: Any
    stages: Any
    upscale_utils: Any
    lu_algo: Any
    lu_config: Any
    lu_factory: Any
    piflow_sampler: Any
    algo_utils: Any
    tile_grid: Any
    tiling_utils: Any
    video_io: Any
    tiny: bool


_reference: Reference | None = None


def load_reference(tiny: bool | None = None) -> Reference:
    """Import the reference (skips the calling test when it is unavailable).

    ``tiny`` (default: True unless CUDA is available) shrinks ``constants.RESOLUTIONS`` in place to 64/96 px bases;
    it only takes effect on the first call (tests needing the trained bases patch the dict, see the tiling tests).
    """
    global _reference
    src = find_reference_src()
    if src is None:
        pytest.skip("k6_video reference (kandinsky_sr) not found: set KANDINSKY_SR_SRC to the k6_video/src directory")
    if tiny is None:
        tiny = not torch.cuda.is_available()
    if _reference is None:
        _install_stubs()
        if str(src) not in sys.path:
            sys.path.insert(0, str(src))
        try:
            imp = importlib.import_module
            _reference = Reference(
                constants=imp("kandinsky_sr.constants"),
                dit_mod=imp("kandinsky_sr.core.components.model.dit"),
                dx_dit_mod=imp("kandinsky_sr.core.components.model.dx_dit"),
                kvae_mod=imp("kandinsky_sr.core.components.video_kvae.cached_model"),
                sr_pipeline=imp("kandinsky_sr.pipeline.sr_pipeline"),
                stages=imp("kandinsky_sr.pipeline.stages"),
                upscale_utils=imp("kandinsky_sr.pipeline.upscale_utils"),
                lu_algo=imp("kandinsky_sr.core.algo.latent_upscaler"),
                lu_config=imp("kandinsky_sr.core.components.latent_upscaler.config"),
                lu_factory=imp("kandinsky_sr.core.components.latent_upscaler.model.factory"),
                piflow_sampler=imp("kandinsky_sr.core.algo.piflow_sampler"),
                algo_utils=imp("kandinsky_sr.core.algo.utils"),
                tile_grid=imp("kandinsky_sr.pipeline.tile_grid"),
                tiling_utils=imp("kandinsky_sr.core.algo.tiling_utils"),
                video_io=imp("kandinsky_sr.pipeline.video_io"),
                tiny=tiny,
            )
        except Exception as exc:  # missing optional dependency of the reference on this machine
            pytest.skip(f"k6_video reference cannot be imported here: {exc!r}")
        if tiny:
            _reference.constants.RESOLUTIONS.clear()
            _reference.constants.RESOLUTIONS.update(TINY_RESOLUTIONS)
    return _reference


def randomize(module: torch.nn.Module, seed: int, std: float = 0.05) -> None:
    generator = torch.Generator().manual_seed(seed)
    with torch.no_grad():
        for parameter in module.parameters():
            parameter.copy_(torch.randn(parameter.shape, generator=generator) * std)


def build_reference_dit(ref: Reference, piflow: dict | None = k6_sr_tiny.TINY_PIFLOW, dit_cfg: dict | None = None,
                        seed: int = 0):
    cfg = copy.deepcopy(dit_cfg or k6_sr_tiny.TINY_DIT_CFG)
    cfg.pop("sr_params", None)  # bundle metadata, not constructor arguments of the reference model
    cfg.pop("attribute_overrides", None)
    cfg["patch_size"] = tuple(cfg["patch_size"])
    cfg["axes_dims"] = tuple(cfg["axes_dims"])
    cfg["attention_params"] = {512: {"type": "flash"}, "512": {"type": "flash"}}
    if piflow is not None:
        dit = ref.dx_dit_mod.DXDiTWrapper(cfg, out_visual_dim=cfg["out_visual_dim"], n_grid=piflow["n_grid"]).eval()
        dit.piflow_params = dict(piflow)
    else:
        dit = ref.dit_mod.get_dit(cfg).eval()
    randomize(dit, seed + 2)
    return dit


def load_into_port(reference_module: torch.nn.Module, port_module: torch.nn.Module, prefix: str = "") -> None:
    """Load ``reference.state_dict()`` (the unprefixed official Diffusers key layout) into a port module, strictly."""
    from fastvideo.models.loader.utils import get_param_names_mapping, hf_to_custom_state_dict

    state = {f"{prefix}{key}": value for key, value in reference_module.state_dict().items()}
    mapping = get_param_names_mapping(getattr(port_module, "param_names_mapping", {}))
    custom, _ = hf_to_custom_state_dict(state, mapping)
    port_module.load_state_dict(custom, strict=True)