tests/conftest.py

import os
import tempfile

import pytest
import torch

from mapgs.config import load_config

CUDA = torch.cuda.is_available()
requires_cuda = pytest.mark.skipif(not CUDA, reason="needs CUDA (gsplat rasterizer)")


def tiny_overrides(root):
    return [
        "model.embed_dim=128", "model.enc_depth=1", "model.dec_depth=2", "model.n_heads=4",
        "model.tokens.n_map=128", "model.tokens.n_free=128",
        "model.tokens.gaussians_per_token=4", "model.tokens.n_dyn_per_instance=16",
        "model.feature_dim=8",
        "data.height=48", "data.width=64", "data.num_frames=8", "data.synth_dynamic_actors=1",
        f"data.root={root}",
        "train.batch_size=2", "train.amp=false", "train.num_workers=0",
        "train.warmup=2", "train.extrap_ramp_iter=2", "train.log_every=1000", "train.ckpt_every=0",
        "loss.lambda_lpips=0.0", "tt.steps=3",
    ]


@pytest.fixture(scope="session")
def device():
    return "cuda" if CUDA else "cpu"


@pytest.fixture(scope="session")
def data_root(tmp_path_factory):
    return str(tmp_path_factory.mktemp("synthetic"))


@pytest.fixture(scope="session")
def cfg(data_root):
    return load_config(overrides=tiny_overrides(data_root))


@pytest.fixture(scope="session")
def dataset(cfg):
    if not CUDA:
        pytest.skip("synthetic GT rendering needs CUDA")
    from mapgs.data import SyntheticDataset
    return SyntheticDataset(cfg, "train", n_scenes=4, n_sup_views=4, device="cuda")