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