| from __future__ import annotations |
|
|
| import math |
| import os |
| from pathlib import Path |
| from typing import Any |
|
|
|
|
| IMAGE_EXTENSIONS = {".jpg", ".jpeg", ".png", ".webp", ".bmp"} |
|
|
|
|
| def dataset_image_count(raw_path: object) -> int: |
| path = Path(str(raw_path or "")).expanduser() |
| if not path.is_dir(): |
| return 0 |
| try: |
| return sum( |
| 1 for item in path.rglob("*") |
| if item.is_file() and item.suffix.casefold() in IMAGE_EXTENSIONS |
| ) |
| except OSError: |
| return 0 |
|
|
|
|
| def _available_vram_gb() -> float | None: |
| try: |
| import pynvml |
|
|
| pynvml.nvmlInit() |
| try: |
| handle = pynvml.nvmlDeviceGetHandleByIndex(0) |
| return pynvml.nvmlDeviceGetMemoryInfo(handle).total / (1024 ** 3) |
| finally: |
| pynvml.nvmlShutdown() |
| except Exception: |
| return None |
|
|
|
|
| def recommend_training_settings( |
| trainer: str, |
| image_count: int, |
| resolution: int, |
| *, |
| vram_gb: float | None = None, |
| ) -> dict[str, Any]: |
| """Return an explainable, conservative starting recipe for manual review.""" |
| trainer = str(trainer).casefold() |
| images = max(10, int(image_count)) |
| resolution = max(64, min(512, int(resolution))) |
| vram = _available_vram_gb() if vram_gb is None else vram_gb |
| cpu_workers = max(2, min(8, (os.cpu_count() or 4) // 2)) |
|
|
| if trainer == "lora": |
| epochs = max(20, min(200, round(8_000 / images))) |
| return { |
| "epochs": epochs, |
| "settings": {}, |
| "summary": ( |
| f"ORION chose {epochs} epochs from about {images:,} images. " |
| "The connected LoRA trainer continues to own its detailed recipe." |
| ), |
| } |
|
|
| exposure_target = 180_000 if trainer == "ddpm" else 120_000 |
| minimum_epochs = 25 |
| maximum_epochs = 600 if trainer == "ddpm" else 300 |
| epochs = max(minimum_epochs, min(maximum_epochs, round(exposure_target / images))) |
| batch_by_resolution = { |
| 64: 16 if trainer == "ddpm" else 12, |
| 128: 12 if trainer == "ddpm" else 8, |
| 256: 4 if trainer == "ddpm" else 4, |
| 384: 2, |
| 512: 1, |
| } |
| nearest_resolution = min(batch_by_resolution, key=lambda size: abs(size - resolution)) |
| batch = batch_by_resolution[nearest_resolution] |
| if vram is not None and vram < 8: |
| batch = max(1, batch // 2) |
| gradient_checkpointing = resolution >= 384 or (vram is not None and vram < 8) |
| settings: dict[str, Any] = { |
| "batch_size": batch, |
| "learning_rate": 0.0001 if trainer == "ddpm" else 0.0002, |
| "gradient_accumulation_steps": 1, |
| "dataloader_num_workers": cpu_workers, |
| "mixed_precision": "fp16", |
| "save_every": max(5, min(25, max(1, epochs // 10))), |
| "preview_steps": 50 if trainer == "ddpm" else 10, |
| "preview_every": max(5, min(50, max(1, epochs // 10))), |
| "training_intensity": 100, |
| "gradient_checkpointing": gradient_checkpointing, |
| } |
| if trainer == "flow": |
| settings["gradient_accumulation"] = settings["gradient_accumulation_steps"] |
| settings["workers"] = settings["dataloader_num_workers"] |
| memory_note = ( |
| f" using the detected {vram:.0f} GB GPU" if vram is not None |
| else " without assuming a specific GPU" |
| ) |
| return { |
| "epochs": epochs, |
| "settings": settings, |
| "summary": ( |
| f"ORION chose about {epochs:,} epochs ({images * epochs:,} image exposures), " |
| f"batch {batch} at {resolution}px{memory_note}. Review this starting recipe before training." |
| ), |
| } |
|
|
|
|
| def review_training_plan(plan: Any) -> dict[str, Any]: |
| """Give a conservative, explainable review without changing user settings.""" |
| findings: list[dict[str, str]] = [] |
| training_steps = [step for step in plan.steps if step.tool_id.endswith("_trainer")] |
| projected_counts = { |
| str(Path(str(step.arguments.get("output_dir", ""))).expanduser()): int( |
| step.arguments.get("image_count", 0) or 0 |
| ) |
| for step in plan.steps |
| if step.tool_id in {"dataset_collector", "youtube_video_collector"} |
| and step.arguments.get("output_dir") |
| } |
| total_steps = 0 |
| estimated_high_minutes = 0 |
|
|
| for step in training_steps: |
| args = step.arguments |
| dataset_key = str(Path(str(args.get("dataset_dir", ""))).expanduser()) |
| images = dataset_image_count(args.get("dataset_dir")) or projected_counts.get(dataset_key, 0) |
| epochs = max(1, int(args.get("epochs", 1) or 1)) |
| batch = max(1, int(args.get("batch_size", 1) or 1)) |
| accumulation = max( |
| 1, |
| int(args.get("gradient_accumulation_steps", args.get("gradient_accumulation", 1)) or 1), |
| ) |
| resolution = max(64, int(args.get("resolution", 256) or 256)) |
| exposures = images * epochs if images else 0 |
| optimizer_steps = math.ceil(images / batch / accumulation) * epochs if images else 0 |
| total_steps += optimizer_steps |
|
|
| |
| if exposures: |
| resolution_factor = (resolution / 256) ** 2 |
| estimated_high_minutes += max(1, math.ceil(exposures * resolution_factor / batch * 0.003)) |
|
|
| label = step.title or step.tool_id.replace("_", " ").title() |
| if images >= 1_000 and epochs >= 300: |
| suggested_epochs = max(25, min(150, round(180_000 / images))) |
| findings.append({ |
| "level": "warning", |
| "message": ( |
| f"{label}: {images:,} images × {epochs:,} epochs requests " |
| f"{exposures:,} image exposures. This resembles a small-dataset preset. " |
| f"Review the intent; about {suggested_epochs} epochs is a safer initial test." |
| ), |
| }) |
| elif exposures >= 1_000_000: |
| findings.append({ |
| "level": "warning", |
| "message": f"{label}: the plan exceeds 1,000,000 image exposures; confirm this is intentional.", |
| }) |
| if images >= 750 and "batch_size" in args and batch == 1 and resolution <= 256: |
| findings.append({ |
| "level": "warning", |
| "message": ( |
| f"{label}: batch size 1 at {resolution}px may leave substantial GPU capacity unused. " |
| "Try a short test with a larger batch if VRAM allows." |
| ), |
| }) |
| if epochs >= 1_000: |
| findings.append({ |
| "level": "warning", |
| "message": f"{label}: {epochs:,} epochs is unusually long and deserves explicit review.", |
| }) |
|
|
| if not training_steps: |
| return {} |
| if not findings: |
| findings.append({ |
| "level": "ready", |
| "message": "No obviously accidental training settings were found. Estimates are still approximate.", |
| }) |
| level = "warning" if any(item["level"] == "warning" for item in findings) else "ready" |
| return { |
| "agent": "ORION", |
| "level": level, |
| "headline": "Review recommended" if level == "warning" else "Plan looks reasonable", |
| "findings": findings, |
| "estimated_optimizer_steps": total_steps, |
| "estimated_high_minutes": estimated_high_minutes, |
| "settings_changed": False, |
| } |
|
|
|
|
| def apply_orion_review(plan: Any) -> dict[str, Any]: |
| review = review_training_plan(plan) |
| plan.orion_review = review |
| if not review: |
| return review |
| lines = [f"ORION — {review['headline']}"] |
| lines.extend(f"• {item['message']}" for item in review["findings"]) |
| if review["estimated_optimizer_steps"]: |
| lines.append(f"• Estimated optimizer steps: about {review['estimated_optimizer_steps']:,}") |
| plan.summary += "\n\n" + "\n".join(lines) |
| if review["level"] == "warning": |
| plan.requires_confirmation = True |
| reason = "ORION found unusual training settings. Review his findings before starting." |
| if reason not in plan.confirmation_reason: |
| plan.confirmation_reason = "; ".join(filter(None, [plan.confirmation_reason, reason])) |
| return review |
|
|
