tests/test_pixelrow.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.pixelrow.generator import generate
from adam.model_plugins_builtin.pixelrow.model import PixelRowConfig, PixelRowModel
from adam.model_plugins_builtin.pixelrow.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 = "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_pixelrow_plugin_is_discovered_with_native_backends() -> None:
    plugins = ModelPluginRegistry(Path.cwd())
    tools = ToolRegistry(Path.cwd())

    plugin = plugins.get("pixelrow")
    trainer = tools.get("pixelrow_trainer")
    generator = tools.get("pixelrow_generator")

    assert plugin.info["architecture"] == "autoregressive_rows"
    assert plugin.training_settings["resolution"]["default"] == 64
    assert plugin.generation_settings["save_progress_frames"]["default"] is True
    assert trainer.backend["module"].endswith("pixelrow.trainer")
    assert "resume_training" in trainer.capabilities
    assert generator.model_trainers == ("pixelrow",)
    assert len(generator.arguments) == len(set(generator.arguments))


def test_pixelrow_generation_request_selects_the_provider() -> None:
    request = parse_chat_generation_request(
        'Generate a PixelRow image of cars using model "Tiny Cars" with 64 steps.'
    )

    assert request is not None
    assert request.provider_hint == "pixelrow"
    assert request.model_query == "Tiny Cars"
    assert request.subject == "cars"
    assert request.steps == 64


def test_pixelrow_model_is_row_causal_and_seeded() -> None:
    config = PixelRowConfig(
        resolution=32,
        hidden_size=64,
        recurrent_layers=1,
        row_channels=16,
        color_bins=16,
    )
    model = PixelRowModel(config)
    images = torch.rand(2, 32, 3, 32) * 2 - 1

    logits = model(images)
    changed_future = images.clone()
    changed_future[:, 10:] = torch.rand_like(changed_future[:, 10:]) * 2 - 1
    changed_logits = model(changed_future)
    loss, parts = model.loss(images, edge_loss_weight=0.05)
    first = model.generate(rows=4, top_k=4, generator=torch.Generator().manual_seed(7))
    second = model.generate(rows=4, top_k=4, generator=torch.Generator().manual_seed(7))

    assert logits.shape == (2, 32, 3, 16, 32)
    assert torch.equal(logits[:, :10], changed_logits[:, :10])
    assert loss.isfinite()
    assert parts["categorical"] > 0
    assert first.shape == (32, 3, 32)
    assert torch.equal(first, second)
    assert torch.count_nonzero(first[4:]) == 0


def test_pixelrow_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" / "pixelrow" / "Tiny Cars"
    training_context = FakeContext(tmp_path, "pixelrow_trainer", "PixelRow Trainer")
    result = train(
        training_context,
        dataset_dir=str(dataset),
        model_name="Tiny Cars",
        epochs=1,
        output_dir=str(output),
        resolution=32,
        batch_size=2,
        hidden_size=64,
        recurrent_layers=1,
        row_channels=16,
        color_bins=16,
        edge_loss_weight=0.0,
        mixed_precision="no",
        preview_enabled=False,
    )

    assert (output / "pixelrow_model.pt").is_file()
    assert (output / "pixelrow_config.json").is_file()
    assert (output / "training_metadata.json").is_file()
    assert result["assets"][0]["trainer"] == "pixelrow"
    assert training_context.progress_updates[-1][0] == 100

    continued_output = (
        tmp_path / "data" / "model_plugin_outputs" / "pixelrow" / "Tiny Cars Continued"
    )
    continuation_context = FakeContext(tmp_path, "pixelrow_trainer", "PixelRow Trainer")
    train(
        continuation_context,
        dataset_dir=str(dataset),
        model_name="Tiny Cars Continued",
        epochs=1,
        output_dir=str(continued_output),
        resume_from=str(output),
        resolution=64,
        batch_size=2,
        edge_loss_weight=0.0,
        mixed_precision="no",
        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, "pixelrow_generator", "PixelRow Generator")
    generated = generate(
        generation_context,
        model_name="Tiny Cars",
        model_path=str(output),
        prompt="car experiment",
        image_count=1,
        steps=3,
        seed=123,
        sampler="Categorical",
        aspect_ratio="1:1 (Native)",
        top_k=4,
        save_progress_frames=True,
        frame_interval=1,
        preview_interval=1,
    )

    generation_folder = Path(generated["output_folder"])
    final_images = list(generation_folder.glob("*_rows_3.png"))
    progress_frames = list(generation_folder.glob("row_progress/*/row_*.png"))
    metadata = list(generation_folder.glob("generation_*.json"))
    assert len(final_images) == 1
    assert len(progress_frames) == 3
    assert len(metadata) == 1
    assert generation_context.previews[-1]["current"] == 3
    assert generation_context.progress_updates[-1][0] == 100