tests/test_inrflow.py

from __future__ import annotations

import json
from pathlib import Path
from types import SimpleNamespace

import torch
from PIL import Image

from adam.generations import parse_chat_generation_request
from adam.model_plugins import ModelPluginRegistry
from adam.model_plugins_builtin.inrflow.generator import generate
from adam.model_plugins_builtin.inrflow.model import (
    INRFlowConfig,
    INRFlowModel,
    coordinate_grid,
    sample_image,
)
from adam.model_plugins_builtin.inrflow.trainer import train
from adam.registry import ToolRegistry


class FakeContext:
    def __init__(self, root: Path, tool_id: str, tool_name: str) -> None:
        self.root = root
        self.job_id = "inrflow-test-job"
        self.tool = SimpleNamespace(id=tool_id, name=tool_name)
        self.logs: list[str] = []
        self.progress_updates: list[tuple[int, str, dict]] = []
        self.previews: list[dict] = []

    def log(self, message: str) -> None:
        self.logs.append(message)

    def checkpoint(self) -> None:
        pass

    def progress(self, percent: int, message: str, **details) -> None:
        self.progress_updates.append((percent, message, details))

    def preview(self, path, **details) -> None:
        self.previews.append({"path": str(path), **details})


def test_inrflow_plugin_is_discovered_with_native_backends() -> None:
    plugins = ModelPluginRegistry(Path.cwd())
    tools = ToolRegistry(Path.cwd())

    plugin = plugins.get("inrflow")
    trainer = tools.get("inrflow_trainer")
    generator = tools.get("inrflow_generator")

    assert plugin.info["architecture"] == "inrflow_ambient_space"
    assert plugin.info["method_reference"].endswith("2412.03791")
    assert plugin.training_settings["query_points"]["default"] == 1024
    assert trainer.backend["module"].endswith("inrflow.trainer")
    assert "resume_training" in trainer.capabilities
    assert generator.model_trainers == ("inrflow",)
    assert "resolution_flexible_generation" in generator.capabilities
    assert len(generator.arguments) == len(set(generator.arguments))


def test_inrflow_generation_request_selects_distinct_provider() -> None:
    request = parse_chat_generation_request(
        'Generate an INRFlow image of tiny cars using model "Coordinate Cars" with 20 steps.'
    )
    spaced = parse_chat_generation_request("Generate an INR Flow image of landscapes")

    assert request is not None
    assert request.provider_hint == "inrflow"
    assert request.model_query == "Coordinate Cars"
    assert request.subject == "tiny cars"
    assert request.steps == 20
    assert spaced is not None and spaced.provider_hint == "inrflow"


def test_inrflow_predicts_pointwise_velocity_and_seeded_samples() -> None:
    config = INRFlowConfig(
        resolution=32,
        patch_size=8,
        hidden_size=64,
        depth=1,
        num_heads=4,
        decoder_layers=1,
        fourier_frequencies=2,
    )
    model = INRFlowModel(config)
    coordinates = coordinate_grid(32, 32).unsqueeze(0).expand(2, -1, -1)
    values = torch.randn(2, 32 * 32, 3)
    prediction = model(
        coordinates,
        values,
        torch.tensor([0.25, 0.75]),
        coordinates[:, :37],
        values[:, :37],
        height=32,
        width=32,
    )
    loss = prediction.square().mean()
    loss.backward()
    first = sample_image(
        model,
        resolution=32,
        steps=2,
        method="Euler",
        noise_scale=1.0,
        query_chunk_size=256,
        generator=torch.Generator().manual_seed(42),
    )
    second = sample_image(
        model,
        resolution=32,
        steps=2,
        method="Euler",
        noise_scale=1.0,
        query_chunk_size=256,
        generator=torch.Generator().manual_seed(42),
    )

    assert prediction.shape == (2, 37, 3)
    assert first.shape == (32, 32, 3)
    assert torch.equal(first, second)


def test_inrflow_training_and_generation_round_trip(tmp_path: Path, monkeypatch) -> None:
    monkeypatch.setattr(torch.cuda, "is_available", lambda: False)
    dataset = tmp_path / "cars"
    dataset.mkdir()
    for index, color in enumerate(((220, 40, 30), (20, 80, 220))):
        image = Image.new("RGB", (40, 28), (225, 225, 225))
        for x in range(8, 32):
            for y in range(10, 22):
                image.putpixel((x, y), color)
        image.save(dataset / f"car_{index}.png")

    output = tmp_path / "data" / "model_plugin_outputs" / "inrflow" / "Coordinate Cars"
    training_context = FakeContext(tmp_path, "inrflow_trainer", "INRFlow Trainer")
    result = train(
        training_context,
        dataset_dir=str(dataset),
        model_name="Coordinate Cars",
        epochs=1,
        output_dir=str(output),
        resolution=32,
        batch_size=2,
        mixed_precision="no",
        patch_size=8,
        hidden_size=64,
        depth=1,
        num_heads=4,
        decoder_layers=1,
        fourier_frequencies=2,
        query_points=64,
        preview_enabled=False,
    )

    assert (output / "inrflow_model.pt").is_file()
    assert (output / "inrflow_config.json").is_file()
    assert (output / "training_metadata.json").is_file()
    assert result["assets"][0]["trainer"] == "inrflow"
    metadata = json.loads((output / "training_metadata.json").read_text(encoding="utf-8"))
    assert metadata["uses_pretrained_compressor"] is False
    assert metadata["settings"]["query_points"] == 64

    continued_output = (
        tmp_path / "data" / "model_plugin_outputs" / "inrflow" / "Coordinate Cars Continued"
    )
    continuation_context = FakeContext(tmp_path, "inrflow_trainer", "INRFlow Trainer")
    train(
        continuation_context,
        dataset_dir=str(dataset),
        model_name="Coordinate Cars Continued",
        epochs=1,
        output_dir=str(continued_output),
        resume_from=str(output),
        resolution=64,
        batch_size=2,
        mixed_precision="no",
        query_points=64,
        preview_enabled=False,
    )
    continued_metadata = json.loads(
        (continued_output / "training_metadata.json").read_text(encoding="utf-8")
    )
    assert continued_metadata["completed_epochs"] == 2
    assert continued_metadata["settings"]["resolution"] == 32

    generation_context = FakeContext(tmp_path, "inrflow_generator", "INRFlow Generator")
    generated = generate(
        generation_context,
        model_name="Coordinate Cars",
        model_path=str(output),
        prompt="ambient-space comparison",
        image_count=1,
        steps=2,
        seed=123,
        sampler="Euler",
        aspect_ratio="1:1 (Coordinate Field)",
        output_resolution="64",
        query_chunk_size=256,
        preview_interval=1,
    )

    generation_folder = Path(generated["output_folder"])
    final_images = list(generation_folder.glob("*_64px.png"))
    generation_metadata = list(generation_folder.glob("generation_*.json"))
    assert len(final_images) == 1
    assert len(generation_metadata) == 1
    assert generation_context.previews[-1]["current"] == 2
    assert generation_context.progress_updates[-1][0] == 100