tests/test_ddpm_adapter.py

from __future__ import annotations

import io
import threading
from pathlib import Path

from adam.config import ConfigManager
from adam.executor import ToolContext
from adam.registry import ToolSpec
from adam.tools import ddpm_adapter


def _context(root: Path, logs: list[str]) -> ToolContext:
    running = threading.Event()
    running.set()
    return ToolContext(
        root=root,
        job_id="DDPMTEST",
        tool=ToolSpec("ddpm_trainer", "DDPM Trainer", "test", "Training", "train_ddpm"),
        cancel_event=threading.Event(),
        run_event=running,
        progress_callback=lambda *_args: None,
        log_callback=logs.append,
    )


def test_resolution_change_branches_from_pipeline_instead_of_resuming_checkpoint(tmp_path: Path, monkeypatch) -> None:
    trainer_root = tmp_path / "DDPM"
    model = trainer_root / "output" / "Anime"
    checkpoint = model / "checkpoint-40"
    dataset = tmp_path / "dataset"
    for folder in (checkpoint / "unet", model / "unet", model / "scheduler", dataset):
        folder.mkdir(parents=True)
    (trainer_root / "train.py").write_text("# fake trainer", encoding="utf-8")
    (model / "model_index.json").write_text("{}", encoding="utf-8")
    (checkpoint / "optimizer.bin").write_bytes(b"optimizer")
    (checkpoint / "scheduler.bin").write_bytes(b"scheduler")
    (checkpoint / "unet" / "diffusion_pytorch_model.safetensors").write_bytes(b"weights")
    (checkpoint / "unet" / "config.json").write_text('{"sample_size": 64}', encoding="utf-8")
    for index in range(2):
        (dataset / f"{index}.png").write_bytes(b"image")

    ConfigManager(tmp_path).update({"tool_folders": {"ddpm_trainer": str(trainer_root)}})
    monkeypatch.setattr(ddpm_adapter.importlib.util, "find_spec", lambda _name: object())
    commands: list[list[str]] = []

    class FakeProcess:
        stdout = io.StringIO("PROGRESS_JSON:{\"event\":\"done\"}\n")
        returncode = 0

        def poll(self):
            return 0

        def terminate(self):
            self.returncode = -15

    def fake_popen(command, **_kwargs):
        commands.append(command)
        return FakeProcess()

    monkeypatch.setattr(ddpm_adapter.subprocess, "Popen", fake_popen)

    logs: list[str] = []
    result = ddpm_adapter.train_ddpm(
        _context(tmp_path, logs),
        dataset_dir=str(dataset),
        model_name="Anime",
        epochs=5,
        output_dir=str(trainer_root / "output" / "Anime V2"),
        resume_from=str(checkpoint),
        resolution=256,
    )

    command = commands[0]
    assert "--pretrained_model_path" in command
    assert str(model.resolve()) in command
    assert "--resume_from_checkpoint" not in command
    assert result["output_folder"] != str(model.resolve())
    assert any("Changing DDPM canvas from 64x64 to 256x256" in line for line in logs)


def test_widescreen_training_uses_native_256_by_144_canvas(tmp_path: Path, monkeypatch) -> None:
    trainer_root = tmp_path / "DDPM"
    output = trainer_root / "output" / "Wide"
    dataset = tmp_path / "dataset"
    dataset.mkdir(parents=True)
    trainer_root.mkdir(parents=True)
    (trainer_root / "train.py").write_text("# fake trainer", encoding="utf-8")
    for index in range(2):
        (dataset / f"{index}.png").write_bytes(b"image")

    ConfigManager(tmp_path).update({"tool_folders": {"ddpm_trainer": str(trainer_root)}})
    monkeypatch.setattr(ddpm_adapter.importlib.util, "find_spec", lambda _name: object())
    commands: list[list[str]] = []

    class FakeProcess:
        stdout = io.StringIO("PROGRESS_JSON:{\"event\":\"done\"}\n")
        returncode = 0

        def poll(self):
            return 0

    monkeypatch.setattr(
        ddpm_adapter.subprocess,
        "Popen",
        lambda command, **_kwargs: commands.append(command) or FakeProcess(),
    )

    ddpm_adapter.train_ddpm(
        _context(tmp_path, []),
        dataset_dir=str(dataset),
        model_name="Wide",
        epochs=1,
        output_dir=str(output),
        resolution=256,
        training_aspect_ratio="16:9 (Widescreen)",
        resize_mode="fit",
    )

    command = commands[0]
    assert command[command.index("--resolution_width") + 1] == "256"
    assert command[command.index("--resolution_height") + 1] == "144"
    assert command[command.index("--resize_mode") + 1] == "fit"


def test_video_dataset_uses_nested_accepted_frames_only(tmp_path: Path, monkeypatch) -> None:
    trainer_root = tmp_path / "DDPM"
    output = trainer_root / "output" / "Video"
    dataset = tmp_path / "Video Dataset"
    accepted = dataset / "frames" / "video_0001"
    rejected = dataset / "rejected" / "video_0001"
    accepted.mkdir(parents=True)
    rejected.mkdir(parents=True)
    trainer_root.mkdir(parents=True)
    (trainer_root / "train.py").write_text("# fake trainer", encoding="utf-8")
    for index in range(2):
        (accepted / f"accepted_{index}.png").write_bytes(b"image")
    (rejected / "rejected_0.png").write_bytes(b"image")

    ConfigManager(tmp_path).update({"tool_folders": {"ddpm_trainer": str(trainer_root)}})
    monkeypatch.setattr(ddpm_adapter.importlib.util, "find_spec", lambda _name: object())
    commands: list[list[str]] = []

    class FakeProcess:
        stdout = io.StringIO("PROGRESS_JSON:{\"event\":\"done\"}\n")
        returncode = 0

        def poll(self):
            return 0

    monkeypatch.setattr(
        ddpm_adapter.subprocess,
        "Popen",
        lambda command, **_kwargs: commands.append(command) or FakeProcess(),
    )

    ddpm_adapter.train_ddpm(
        _context(tmp_path, []), dataset_dir=str(dataset), model_name="Video",
        epochs=1, output_dir=str(output),
    )

    command = commands[0]
    assert command[command.index("--train_data_dir") + 1] == str((dataset / "frames").resolve())