from __future__ import annotations import math import os from dataclasses import asdict, dataclass, field from typing import Any from adam.model_profiles import ModelProfile from adam.models import SystemSnapshot @dataclass(slots=True) class SettingsRecommendation: profile_id: str epochs: int settings: dict[str, Any] = field(default_factory=dict) reasons: list[str] = field(default_factory=list) warnings: list[str] = field(default_factory=list) summary: str = "" estimated_vram_gb: float | None = None risk_level: str = "normal" confidence: str = "conservative" def to_dict(self) -> dict[str, Any]: return asdict(self) def _field_default(profile: ModelProfile, key: str, fallback: Any) -> Any: return profile.training.get(key, {}).get("default", fallback) def _clamp_to_schema(profile: ModelProfile, key: str, value: Any) -> Any: spec = profile.training.get(key, {}) kind = str(spec.get("type", "text")) try: if kind in {"int", "slider"}: numeric = int(value) return max(int(spec.get("min", numeric)), min(numeric, int(spec.get("max", numeric)))) if kind == "float": numeric = float(value) return max(float(spec.get("min", numeric)), min(numeric, float(spec.get("max", numeric)))) except (TypeError, ValueError): return spec.get("default", value) if kind == "choice": options = list(spec.get("options", [])) return value if value in options else (options[0] if options else value) return value def estimate_vram_gb(profile: ModelProfile, resolution: int | str, batch_size: int, base_model_gb: float = 0.0) -> float: """Broad VRAM estimate used only for warnings and conservative defaults.""" architecture = profile.architecture.casefold() pixels = (max(64, resolution) / 512) ** 2 if profile.id == "lora" or "lora" in architecture: base = max(6.0, base_model_gb * 1.8) return base + pixels * max(1, batch_size) * 1.2 if profile.id == "oasis" or "action_conditioned" in architecture: width, height = (resolution, resolution) if isinstance(resolution, str) and "x" in resolution: try: width, height = (int(part) for part in resolution.lower().split("x", 1)) except ValueError: width, height = (256, 144) pixels = (max(width, height) / 512) ** 2 return 4.0 + pixels * max(1, batch_size) * 2.4 if "flow" in architecture: return 2.8 + pixels * max(1, batch_size) * 1.0 if "diffusion" in architecture: return 2.2 + pixels * max(1, batch_size) * 0.9 return 3.0 + pixels * max(1, batch_size) * 0.8 def recommend_for_profile( profile: ModelProfile, *, dataset_items: int, resolution: int | str | None = None, snapshot: SystemSnapshot | None = None, base_model_gb: float = 0.0, ) -> SettingsRecommendation: images = max(10, int(dataset_items or 10)) raw_resolution = resolution or _field_default(profile, "resolution", 256) or 256 if isinstance(raw_resolution, str) and "x" in raw_resolution: resolution = int(raw_resolution.lower().split("x", 1)[0]) else: resolution = int(raw_resolution) reasons: list[str] = [] vram_total = snapshot.vram_total_gb if snapshot and snapshot.vram_total_gb else None available_vram = ( max(0.0, snapshot.vram_total_gb - snapshot.vram_used_gb) if snapshot and snapshot.vram_total_gb else vram_total ) architecture = profile.architecture.casefold() target_exposures = 80_000 if profile.id == "lora" else 180_000 if "diffusion" in architecture else 120_000 max_epochs = 220 if profile.id == "lora" else 600 if "diffusion" in architecture else 300 epochs = max(10 if profile.id == "lora" else 25, min(max_epochs, round(target_exposures / images))) reasons.append( f"Epochs target roughly {target_exposures:,} image exposures, then clamp to the profile's safe range." ) batch_defaults = { 64: 16, 128: 12, 256: 4, 384: 2, 512: 1, 768: 1, 1024: 1, } if profile.id == "flow": batch_defaults.update({64: 12, 128: 8, 256: 4}) if profile.id == "oasis": batch_defaults.update({128: 4, 256: 2, 384: 1, 512: 1}) if profile.id == "lora": batch_defaults.update({512: 2, 768: 1, 1024: 1}) nearest = min(batch_defaults, key=lambda size: abs(size - resolution)) batch_size = batch_defaults[nearest] reasons.append(f"Batch starts from the closest resolution preset ({nearest}px).") if available_vram is not None and available_vram < 8: batch_size = max(1, batch_size // 2) reasons.append("Available VRAM is below 8 GB, so batch size is reduced conservatively.") settings: dict[str, Any] = {} for key in ("resolution", "batch_size"): if key in profile.training: value = batch_size if key == "resolution": value = raw_resolution if profile.id == "oasis" else resolution settings[key] = _clamp_to_schema( profile, key, value, ) if "learning_rate" in profile.training: settings["learning_rate"] = _clamp_to_schema( profile, "learning_rate", 0.00002 if profile.id == "oasis" else 0.0001 if profile.id in {"ddpm", "lora"} else 0.0002, ) workers = max(1, min(8, (os.cpu_count() or 4) // 2)) for key in ("dataloader_num_workers", "workers"): if key in profile.training: settings[key] = _clamp_to_schema(profile, key, workers) for key, value in { "gradient_accumulation_steps": 1, "gradient_accumulation": 1, "mixed_precision": "fp32" if profile.id == "oasis" else "fp16", "save_every": max(5, min(25, max(1, epochs // 10))), "preview_every": max(5, min(50, max(1, epochs // 10))), "training_intensity": 100, "gradient_checkpointing": resolution >= 384 or (available_vram is not None and available_vram < 8), "rank": 16, "alpha": 16, "frame_gap": 3, "sequence_context": 1, "preview_steps": 1 if profile.id == "oasis" else 50 if profile.id == "ddpm" else 10, }.items(): if key in profile.training: settings[key] = _clamp_to_schema(profile, key, value) estimated = estimate_vram_gb(profile, resolution, int(settings.get("batch_size", batch_size)), base_model_gb) warnings: list[str] = [] if available_vram is not None and estimated > available_vram * 0.9: warnings.append( f"Estimated VRAM need is about {estimated:.1f} GB, above the conservative {available_vram * 0.9:.1f} GB working limit." ) if "batch_size" in settings and int(settings["batch_size"]) > 1: settings["batch_size"] = max(1, int(settings["batch_size"]) // 2) estimated = estimate_vram_gb(profile, resolution, int(settings["batch_size"]), base_model_gb) warnings.append(f"Batch size was reduced to {settings['batch_size']} for a safer first run.") reasons.append("The initial VRAM estimate was high, so ADAM reduced the batch before applying the recipe.") if images < 20: warnings.append("Dataset is very small; expect overfitting unless this is just a smoke test.") reasons.append("Very small datasets get a warning because quality usually depends more on data cleanup than long training.") risk_level = "risky" if warnings else "normal" memory_note = ( f" using about {available_vram:.1f} GB available VRAM" if available_vram is not None else " without detected VRAM" ) summary = ( f"Recommended {epochs:,} epochs for {images:,} item(s), " f"batch {settings.get('batch_size', batch_size)} at {resolution}px{memory_note}. " "Treat this as a starting recipe, not a guarantee." ) return SettingsRecommendation( profile_id=profile.id, epochs=epochs, settings=settings, reasons=reasons, warnings=warnings, summary=summary, estimated_vram_gb=estimated, risk_level=risk_level, )