File size: 2,949 Bytes
c61c435 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 | from __future__ import annotations
from pathlib import Path
import json
from adam.commands import TrainingCommand
from adam.model_plugins import (
ModelPluginRegistry,
scaffold_model_plugin,
validate_settings,
)
from adam.registry import ToolRegistry
def test_builtin_model_plugins_are_discovered() -> None:
registry = ModelPluginRegistry(Path.cwd())
assert {"ddpm", "flow", "lora"}.issubset(registry.plugins)
assert registry.errors == []
assert registry.training_schema("ddpm")["resolution"]["type"] == "choice"
assert registry.generation_schema_for_tool("lora_generator")["base_model_path"]["required"]
def test_plugin_schema_validation_reports_clear_errors(tmp_path: Path) -> None:
schema = {
"batch_size": {"label": "Batch size", "type": "int", "min": 1, "max": 8},
"base_model": {"label": "Base model", "type": "path", "required": True, "must_exist": True},
}
errors = validate_settings(schema, {"batch_size": 0, "base_model": str(tmp_path / "missing.safetensors")})
assert "Batch size must be at least 1." in errors
assert "Base model must point to an existing file." in errors
def test_plugin_schema_extends_existing_tool_arguments_without_required_breakage() -> None:
registry = ToolRegistry(Path.cwd())
lora = registry.get("lora_trainer")
assert "rank" in lora.arguments
assert "alpha" in lora.arguments
assert set(lora.required_arguments) == {
"dataset_dir",
"model_name",
"epochs",
"output_dir",
"base_model",
}
def test_scaffolded_plugin_is_discovered_and_gets_standard_training_arguments(tmp_path: Path) -> None:
folder = scaffold_model_plugin(
tmp_path,
plugin_id="Neural Cellular Automata",
name="Neural Cellular Automata",
architecture="nca",
)
config = tmp_path / "config"
config.mkdir()
(config / "tools.json").write_text(json.dumps({"tools": []}), encoding="utf-8")
registry = ToolRegistry(tmp_path)
tool = registry.get("neural_cellular_automata_trainer")
assert folder.name == "neural_cellular_automata"
assert "dataset_dir" in tool.arguments
assert "output_dir" in tool.required_arguments
assert registry.model_plugins.training_schema("neural_cellular_automata")["resolution"]["default"] == 256
def test_training_command_accepts_discovered_custom_plugin(tmp_path: Path, monkeypatch) -> None:
scaffold_model_plugin(
tmp_path,
plugin_id="maskgit",
name="MaskGIT",
architecture="maskgit",
)
monkeypatch.chdir(tmp_path)
command = TrainingCommand.from_dict(
{
"action": "train",
"trainer": "maskgit",
"dataset": "D:/data",
"model_name": "Mask Test",
"epochs": 5,
"training_options": {"resolution": 256, "batch_size": 1},
}
)
assert command.trainer == "maskgit"
|