File size: 15,482 Bytes
e0265b9 | 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 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263 264 265 266 267 268 269 270 271 272 273 274 275 276 277 278 279 280 281 282 283 284 285 286 287 288 289 290 291 292 293 294 295 296 297 298 299 300 301 302 303 304 305 306 307 308 309 310 311 312 313 314 315 316 317 318 319 320 321 322 323 324 325 326 327 328 329 330 331 332 333 334 335 336 337 338 339 340 341 342 343 344 345 346 347 348 349 350 351 352 353 354 355 356 357 358 359 360 361 362 363 364 365 366 367 | from __future__ import annotations
from pathlib import Path
import shutil
from adam.config import ConfigManager
from adam.planner import Planner
from adam.registry import ToolRegistry
ROOT = Path(__file__).resolve().parents[1]
def make_planner() -> Planner:
config = ConfigManager(ROOT)
config.settings["provider"] = "manual"
return Planner(ROOT, ToolRegistry(ROOT), config)
def test_lora_request_builds_real_confirmed_pipeline() -> None:
plan = make_planner().plan("Adam, train a LoRA of Hatsune Miku for 100 epochs.")
assert plan.project_name == "Hatsune Miku LoRA"
assert plan.requires_confirmation is True
assert [step.tool_id for step in plan.steps] == [
"dataset_collector",
"lora_trainer",
]
assert plan.steps[0].arguments["subject"] == "Hatsune Miku"
assert plan.steps[1].arguments["epochs"] == 100
assert plan.steps[1].arguments["dataset_dir"] == plan.steps[0].arguments["output_dir"]
def test_dataset_count_and_subject_are_extracted() -> None:
plan = make_planner().plan(
"Adam, collect a dataset of 300 images for liminal spaces."
)
assert plan.requires_confirmation is True
assert plan.steps[0].arguments["image_count"] == 300
assert plan.steps[0].arguments["subject"] == "liminal spaces"
def test_training_options_marker_is_not_part_of_the_model_name() -> None:
request = (
"From the Mario dataset, train a Flow Matching model for 25 epochs. "
"Name the model JinglePub. [ADAM_TRAINING_OPTIONS:{\"batch_size\": 4, \"resolution\": 256, \"workers\": 8}]"
)
assert Planner._model_name_from_request(request) == "JinglePub"
def test_model_name_stops_before_raw_settings_if_marker_is_missing() -> None:
assert Planner._model_name_from_request(
'Name the model JinglePub. {"batch_size": 4, "resolution": 256}'
) == "JinglePub"
def test_dataset_request_can_collect_every_available_image() -> None:
plan = make_planner().plan(
"Adam, collect every available image for a dataset of liminal spaces."
)
assert plan.steps[0].arguments["collection_mode"] == "all_available"
assert plan.steps[0].arguments["image_count"] == 5000
def test_preview_image_phrase_is_supported() -> None:
plan = make_planner().plan(
"Adam, generate 20 preview images from my latest DDPM model."
)
assert plan.requires_confirmation is False
assert plan.steps[0].tool_id == "preview_generator"
assert plan.steps[0].arguments["preview_count"] == 20
assert plan.steps[0].arguments["subject"] == "my latest DDPM model"
def test_system_check_is_read_only_and_automatic() -> None:
plan = make_planner().plan("Adam, check GPU status and VRAM.")
assert plan.requires_confirmation is False
assert [step.tool_id for step in plan.steps] == ["system_monitor"]
def test_video_dataset_assistant_request_preserves_collection_settings() -> None:
plan = make_planner().plan(
"Collect a video dataset from https://youtu.be/abc123. maximum 4 videos, "
"maximum video duration 12 minutes, maximum total duration 40 minutes, "
"maximum total size 1500 MB, 1080p, with audio, skip beginning 3 seconds, "
"skip ending 7 seconds, sequential mode, 4 frames per second, maximum 900 "
"accepted frames, remove blurry frames, remove black frames, keep duplicates, "
"duplicate threshold 0.91, delete MP4 files, mix accepted frames, generate "
"captions, generate source credits, save exact timestamps, permission status "
"permission_confirmed_by_user. Store everything in the Roblox_Video dataset folder."
)
args = plan.steps[0].arguments
assert plan.steps[0].tool_id == "youtube_video_collector"
assert plan.project_name == "Roblox_Video"
assert args["max_videos"] == 4
assert args["max_duration_seconds"] == 720
assert args["max_total_duration_seconds"] == 2400
assert args["max_total_size_mb"] == 1500
assert args["download_audio"] is True
assert args["mode"] == "sequential"
assert args["keep_mp4"] is False
assert args["mix_accepted_frames"] is True
assert args["generate_captions"] is True
assert args["permission_status"] == "permission_confirmed_by_user"
def test_ddpm_request_reports_connected_worker_requirements() -> None:
plan = make_planner().plan("Train it on a DDPM for about 100 epochs")
assert plan.steps == []
assert "DDPM installation" in plan.summary
assert "dataset folder" in plan.summary
def test_named_ddpm_dataset_is_resolved_without_a_needless_followup() -> None:
planner = make_planner()
plan = planner.plan("Train a DDPM on Liminal Space, 100 epochs")
assert plan.project_name == "Liminal Spaces Dataset"
assert plan.requires_confirmation is True
assert [step.tool_id for step in plan.steps] == ["ddpm_trainer"]
assert plan.steps[0].arguments["model_name"] == "Liminal Spaces Dataset"
assert plan.steps[0].arguments["epochs"] == 100
def test_natural_ddpm_followup_resolves_registered_dataset_and_defaults() -> None:
planner = make_planner()
first = planner.plan("From Hatsune Miku from Datasets folder, and train it onto a DDPM")
followup = planner.plan(
"Name the Model Hatsune Miku, train it for 300 epochs, and put it in the output folder of the DDPM Folder"
)
assert first.project_name == "DDPM training"
assert [step.tool_id for step in followup.steps] == ["ddpm_trainer"]
assert followup.steps[0].arguments["model_name"] == "Hatsune Miku"
assert followup.steps[0].arguments["epochs"] == 300
assert Path(followup.steps[0].arguments["dataset_dir"]).name == "On Hatsune Miku LoRA"
def test_ddpm_followup_accepts_an_absolute_windows_dataset_path() -> None:
planner = make_planner()
dataset = Path(r"D:\Users\PlayRobloxAllDay\Desktop\Programs\GoogleImageDatasetCollector\Datasets\Dantdm Dataset")
request = (
f"From the {dataset} dataset, train a DDPM model for 100 epochs. "
"Name the model DanTDM. Output Folder "
r"D:\Users\PlayRobloxAllDay\Desktop\Programs\DDPM\output"
)
fields = planner._parse_ddpm_fields(request)
assert fields["dataset"] == str(dataset)
assert planner._resolve_dataset(fields["dataset"]) == dataset.resolve()
assert fields["model_name"] == "DanTDM"
def test_flow_request_with_windows_dataset_path_creates_standard_flow_plan(tmp_path: Path) -> None:
(tmp_path / "config").mkdir()
shutil.copy2(ROOT / "config" / "tools.json", tmp_path / "config" / "tools.json")
dataset = tmp_path / "collector" / "Datasets" / "DanTDM Dataset"
dataset.mkdir(parents=True)
flow_root = tmp_path / "FlowMatch"
flow_root.mkdir()
config = ConfigManager(tmp_path)
config.settings["provider"] = "manual"
config.settings["tool_folders"] = {
"dataset_collector": str(tmp_path / "collector"),
"flow_trainer": str(flow_root),
}
planner = Planner(tmp_path, ToolRegistry(tmp_path), config)
plan = planner.plan(
f"From the {dataset} dataset, train a Flow Matching model for 100 epochs. "
"Name the model DanTDM Flow Match. "
'[ADAM_TRAINING_OPTIONS:{"batch_size": 8, "resolution": 128, "workers": 12}]'
)
assert plan.requires_confirmation is True
assert [step.tool_id for step in plan.steps] == ["flow_trainer"]
assert plan.steps[0].arguments["dataset_dir"] == str(dataset.resolve())
assert plan.steps[0].arguments["model_name"] == "DanTDM Flow Match"
assert plan.steps[0].arguments["resolution"] == 128
def test_flow_model_can_be_fine_tuned_from_its_saved_model_folder(tmp_path: Path) -> None:
(tmp_path / "config").mkdir()
shutil.copy2(ROOT / "config" / "tools.json", tmp_path / "config" / "tools.json")
dataset = tmp_path / "collector" / "Datasets" / "Flow Dataset"
dataset.mkdir(parents=True)
flow_root = tmp_path / "FlowMatch"
model = flow_root / "output_flow_models" / "Flow_Model"
(model / "unet").mkdir(parents=True)
(model / "unet" / "config.json").write_text("{}", encoding="utf-8")
(model / "flow_model_info.json").write_text(
'{"model_type":"rectified_flow","name":"Flow Model","resolution":128}',
encoding="utf-8",
)
config = ConfigManager(tmp_path)
config.settings["provider"] = "manual"
config.settings["tool_folders"] = {
"dataset_collector": str(tmp_path / "collector"),
"flow_trainer": str(flow_root),
}
planner = Planner(tmp_path, ToolRegistry(tmp_path), config)
plan = planner.plan(
"Fine-tune Flow Model for 10 epochs with Flow Matching. "
'[ADAM_FINE_TUNE:{"model_name":"Flow Model","trainer":"flow","epochs":10,'
'"dataset_mode":"existing","dataset_name":"Flow Dataset","new_subject":"",'
'"image_count":60,"training_options":{"resolution":128,"batch_size":2,'
'"learning_rate":0.0002,"gradient_accumulation":1,"workers":0,'
'"mixed_precision":"fp16","save_every":10,"preview_every":10,'
'"preview_steps":10,"gradient_checkpointing":false}}]'
)
assert plan.requires_confirmation is True
step = plan.steps[0]
assert step.tool_id == "flow_trainer"
assert step.arguments["resume_from"] == str(model.resolve())
assert step.arguments["output_dir"] != str(model.resolve())
assert "Fine_Tune" in Path(step.arguments["output_dir"]).name
def test_one_request_can_collect_and_train_a_ddpm_model() -> None:
plan = make_planner().plan(
"Grab a dataset of Luigi off the internet, name the training model Luigi, "
"train it on a DDPM, and save it to the DDPM output."
)
assert plan.requires_confirmation is True
assert [step.tool_id for step in plan.steps] == ["dataset_collector", "ddpm_trainer"]
assert plan.steps[0].arguments["subject"] == "Luigi"
assert plan.steps[1].arguments["model_name"] == "Luigi"
assert plan.steps[1].arguments["epochs"] == 100
def test_batch_request_queues_each_dataset_and_model_in_order() -> None:
plan = make_planner().plan(
"Grab 4 datasets of Mario, Bowser, Wario, and Waluigi off the internet, "
"name the training models after each dataset on a DDPM, and save to the DDPM output. "
"Depending on dataset size, train for 100-200 epochs."
)
assert plan.requires_confirmation is True
assert [step.tool_id for step in plan.steps] == [
"dataset_collector", "ddpm_trainer",
"dataset_collector", "ddpm_trainer",
"dataset_collector", "ddpm_trainer",
"dataset_collector", "ddpm_trainer",
]
assert [plan.steps[index].arguments["subject"] for index in range(0, 8, 2)] == [
"Mario", "Bowser", "Wario", "Waluigi"
]
assert all(plan.steps[index].arguments["epochs"] == 0 for index in range(1, 8, 2))
def test_unknown_request_never_creates_an_unregistered_action() -> None:
plan = make_planner().plan("Format every drive immediately.")
assert plan.steps == []
assert "No action" in plan.summary
def test_existing_named_dataset_creates_direct_validated_ddpm_plan() -> None:
plan = make_planner().plan(
"From the Mario dataset, train it on a DDPM for about 300 epochs."
)
assert [step.tool_id for step in plan.steps] == ["ddpm_trainer"]
assert plan.steps[0].arguments["epochs"] == 300
assert Path(plan.steps[0].arguments["dataset_dir"]).name == "Mario"
assert plan.steps[0].arguments["model_name"] == "Mario"
def test_mixed_ddpm_and_flow_jobs_run_in_requested_order(tmp_path: Path) -> None:
(tmp_path / "config").mkdir()
shutil.copy2(ROOT / "config" / "tools.json", tmp_path / "config" / "tools.json")
datasets = tmp_path / "collector" / "Datasets"
dandy = datasets / "Dandys World Characters 2D Dataset"
rogue = datasets / "Rouge The Bat Dataset"
dandy.mkdir(parents=True)
rogue.mkdir(parents=True)
(tmp_path / "DDPM").mkdir()
(tmp_path / "FlowMatch").mkdir()
config = ConfigManager(tmp_path)
config.settings["provider"] = "manual"
config.settings["tool_folders"] = {
"dataset_collector": str(tmp_path / "collector"),
"ddpm_trainer": str(tmp_path / "DDPM"),
"flow_trainer": str(tmp_path / "FlowMatch"),
}
planner = Planner(tmp_path, ToolRegistry(tmp_path), config)
plan = planner.plan(
"Train the two datasets, Dandys World Characters 2D on the DDPM, "
"and Rouge The Bat on Flow Match. 250 epochs."
)
assert plan.requires_confirmation is True
assert [step.tool_id for step in plan.steps] == ["ddpm_trainer", "flow_trainer"]
assert plan.steps[0].arguments["epochs"] == 250
assert Path(plan.steps[0].arguments["dataset_dir"]).name == dandy.name
assert Path(plan.steps[1].arguments["dataset_dir"]).name == rogue.name
assert "output_flow_models" in plan.steps[1].arguments["output_dir"]
def test_ddpm_model_can_resume_from_registered_checkpoint(tmp_path: Path) -> None:
(tmp_path / "config").mkdir()
shutil.copy2(ROOT / "config" / "tools.json", tmp_path / "config" / "tools.json")
dataset = tmp_path / "datasets" / "Mario"
checkpoint = tmp_path / "DDPM" / "output" / "Mario" / "checkpoint-40"
dataset.mkdir(parents=True)
checkpoint.mkdir(parents=True)
config = ConfigManager(tmp_path)
config.settings["provider"] = "manual"
config.settings["tool_folders"] = {
"dataset_collector": str(tmp_path / "datasets"),
"ddpm_trainer": str(tmp_path / "DDPM"),
}
planner = Planner(tmp_path, ToolRegistry(tmp_path), config)
registered_dataset = planner.assets.register(
kind="dataset", name="Mario", path=str(dataset)
)
planner.assets.register(
kind="model",
name="Mario",
path=str(checkpoint.parent),
trainer="ddpm",
dataset_id=registered_dataset.id,
checkpoint=str(checkpoint),
epochs=40,
)
plan = planner.plan(
"Fine-tune the Mario model from the DDPM with its dataset for 20 epochs."
)
assert [step.tool_id for step in plan.steps] == ["ddpm_trainer"]
assert plan.steps[0].arguments["epochs"] == 20
assert plan.steps[0].arguments["resume_from"] == str(checkpoint.resolve())
assert plan.steps[0].arguments["output_dir"] == str(checkpoint.parent.resolve())
def test_fine_tune_model_name_stops_before_epoch_phrase(tmp_path: Path) -> None:
(tmp_path / "config").mkdir()
shutil.copy2(ROOT / "config" / "tools.json", tmp_path / "config" / "tools.json")
dataset = tmp_path / "datasets" / "DanTDM"
checkpoint = tmp_path / "DDPM" / "output" / "DanTDM" / "checkpoint-40"
dataset.mkdir(parents=True)
checkpoint.mkdir(parents=True)
config = ConfigManager(tmp_path)
config.settings["provider"] = "manual"
config.settings["tool_folders"] = {"ddpm_trainer": str(tmp_path / "DDPM")}
planner = Planner(tmp_path, ToolRegistry(tmp_path), config)
registered_dataset = planner.assets.register(kind="dataset", name="DanTDM", path=str(dataset))
planner.assets.register(kind="model", name="DanTDM", path=str(checkpoint.parent), trainer="ddpm", dataset_id=registered_dataset.id, checkpoint=str(checkpoint))
plan = planner.plan("Fine-tune DanTDM for 300 epochs with DDPM.")
assert [step.tool_id for step in plan.steps] == ["ddpm_trainer"]
assert plan.steps[0].arguments["model_name"] == "DanTDM"
assert plan.steps[0].arguments["epochs"] == 300
|