adam/progressive_training.py

"""Resolution-curriculum helpers shared by the DDPM and Flow adapters.

The adapters deliberately use conservative batch caps instead of trying an
out-of-memory probe in a real training job.  A user can still turn the policy
off and enter every batch setting manually.
"""

from __future__ import annotations

from dataclasses import dataclass
from typing import Any

from adam.executor import ToolExecutionError


@dataclass(frozen=True, slots=True)
class ResolutionStage:
    resolution: int
    epochs: int


_BATCH_CAPS = {
    "ddpm": {64: 16, 128: 12, 256: 4, 384: 2, 512: 1},
    "flow": {64: 12, 128: 8, 256: 4, 384: 2, 512: 1},
}


def parse_stages(value: Any, *, trainer: str, total_epochs: int) -> list[ResolutionStage]:
    """Validate a JSON-friendly progressive-resolution schedule.

    Stages are intentionally a list of small dictionaries so plans remain easy
    to inspect and edit in saved job JSON.
    """
    if not isinstance(value, list) or len(value) < 2:
        raise ToolExecutionError("Progressive training needs at least two resolution stages.")
    multiple = 16 if trainer == "flow" else 8
    stages: list[ResolutionStage] = []
    previous = 0
    for raw in value:
        if not isinstance(raw, dict):
            raise ToolExecutionError("Each progressive stage must include resolution and epochs.")
        try:
            resolution = int(raw.get("resolution", 0))
            epochs = int(raw.get("epochs", 0))
        except (TypeError, ValueError) as exc:
            raise ToolExecutionError("Progressive stage resolution and epochs must be whole numbers.") from exc
        if not 64 <= resolution <= 512 or resolution % multiple:
            raise ToolExecutionError(
                f"{trainer.upper()} progressive resolutions must be 64–512 and divisible by {multiple}."
            )
        if resolution <= previous:
            raise ToolExecutionError("Progressive stages must increase from lower to higher resolution.")
        if epochs < 1:
            raise ToolExecutionError("Each progressive stage needs at least one epoch.")
        stages.append(ResolutionStage(resolution, epochs))
        previous = resolution
    if sum(stage.epochs for stage in stages) != int(total_epochs):
        raise ToolExecutionError(
            f"Progressive stage epochs total {sum(stage.epochs for stage in stages):,}, "
            f"but training length is {int(total_epochs):,}."
        )
    return stages


def suggested_stages(final_resolution: int, total_epochs: int) -> list[ResolutionStage]:
    """Return an editable low-to-high schedule that always matches the budget."""
    resolutions = [size for size in (64, 128, 256, 384, 512) if size <= int(final_resolution)]
    if len(resolutions) < 2:
        return [ResolutionStage(int(final_resolution), int(total_epochs))]
    if int(total_epochs) < len(resolutions):
        resolutions = resolutions[-int(total_epochs):]
        return [ResolutionStage(resolution, 1) for resolution in resolutions]
    # Front-load inexpensive structure learning while reserving final-resolution
    # refinement.  Normalizing lets the same policy work for any epoch budget.
    weights = [0.60, 0.20, 0.10, 0.06, 0.04][-len(resolutions):]
    allocation = [max(1, round(total_epochs * weight / sum(weights))) for weight in weights]
    difference = int(total_epochs) - sum(allocation)
    allocation[0] += difference
    return [ResolutionStage(resolution, epochs) for resolution, epochs in zip(resolutions, allocation)]


def stage_batch_settings(
    *, trainer: str, stage_resolution: int, final_resolution: int,
    final_batch_size: int, base_accumulation: int, auto_batch: bool,
) -> tuple[int, int]:
    """Choose a conservative physical batch and matching accumulation count."""
    requested_batch = max(1, min(64, int(final_batch_size)))
    requested_accumulation = max(1, min(64, int(base_accumulation)))
    if not auto_batch:
        return requested_batch, requested_accumulation
    caps = _BATCH_CAPS[trainer]
    cap = caps[min(caps, key=lambda size: abs(size - int(stage_resolution)))]
    # Scale from the user's final-stage batch according to image area, then
    # enforce a trainer-specific safe cap. This never raises the 512px batch.
    scaled = round(requested_batch * (int(final_resolution) / int(stage_resolution)) ** 2)
    batch_size = max(1, min(cap, 64, scaled))
    effective_batch = requested_batch * requested_accumulation
    accumulation = max(1, min(64, -(-effective_batch // batch_size)))
    return batch_size, accumulation


def stage_summary(stages: list[ResolutionStage]) -> str:
    return ", ".join(f"{stage.resolution}px × {stage.epochs}" for stage in stages)