Merge branch 'diffusion-train-tab-2' into diffusion-krea2

studio/backend/core/training/diffusion_dit_trainer.py

@@ -31,7 +31,7 @@ import os
 import random
 import time
 from contextlib import nullcontext
-from dataclasses import dataclass
+from dataclasses import dataclass, replace
 from pathlib import Path
 from typing import Any, Callable, Optional
 
@@ -48,6 +48,7 @@ from core.training.diffusion_train_common import (
     _restore_perf_flags,
     discover_image_caption_pairs,
     repo_is_prequantized,
+    resolve_train_steps,
 )
 
 # Per-family LoRA target modules (attention projections). FLUX / Qwen double-stream blocks
@@ -1081,6 +1082,10 @@ def run_dit_lora_training(
     pairs = discover_image_caption_pairs(
         cfg.data_dir, instance_prompt = cfg.instance_prompt, caption_column = cfg.caption_column
     )
+    # Resolve num_epochs -> a concrete train_steps now that the dataset size is known, and
+    # rebind cfg so every downstream read (scheduler length, the loop range, progress
+    # total_steps, steps_run) sees the same resolved value.
+    cfg = replace(cfg, train_steps = resolve_train_steps(cfg, len(pairs)), num_epochs = 0)
     _emit(on_event, "model_load_started", num_images = len(pairs))
     if _check_stop():
         out_dir = Path(cfg.output_dir).expanduser()

studio/backend/core/training/diffusion_lora_trainer.py

@@ -36,6 +36,7 @@ import gc
 import os
 import random
 import time
+from dataclasses import replace
 from pathlib import Path
 from typing import Any, Optional
 
@@ -58,6 +59,7 @@ from core.training.diffusion_train_common import (  # noqa: F401
     _restore_perf_flags,
     discover_image_caption_pairs,
     get_trainer,
+    resolve_train_steps,
 )
 
 
@@ -301,6 +303,10 @@ def run_diffusion_lora_training(
         pairs = discover_image_caption_pairs(
             cfg.data_dir, instance_prompt = cfg.instance_prompt, caption_column = cfg.caption_column
         )
+        # Resolve num_epochs -> a concrete train_steps now that the dataset size is known, and
+        # rebind cfg so every downstream read (scheduler length, the loop range, progress
+        # total_steps, steps_run) sees the same resolved value.
+        cfg = replace(cfg, train_steps = resolve_train_steps(cfg, len(pairs)), num_epochs = 0)
         _emit(on_event, "model_load_started", num_images = len(pairs))
 
         # Honour a stop requested before the (potentially large / slow) base model loads, the

studio/backend/core/training/diffusion_train_common.py

@@ -16,6 +16,7 @@ actual training loop; this module only routes a request to the right one.
 from __future__ import annotations
 
 import json
+import math
 import os
 import random
 import re
@@ -250,6 +251,9 @@ class DiffusionLoraConfig:
     instance_prompt: Optional[str] = None
     resolution: int = 1024
     train_steps: int = 500
+    # 0 = disabled (train for train_steps). > 0 overrides train_steps with a run length of
+    # num_epochs full passes over the dataset, in optimizer steps (see resolve_train_steps).
+    num_epochs: int = 0
     learning_rate: float = 1e-4
     train_batch_size: int = 1
     gradient_accumulation_steps: int = 1
@@ -303,6 +307,8 @@ class DiffusionLoraConfig:
         resolved_family = resolve_trainable_family(self.base_model, self.model_family)
         if self.train_steps < 1:
             raise ValueError("train_steps must be >= 1")
+        if not 0 <= int(self.num_epochs) <= 1000:
+            raise ValueError("num_epochs must be between 0 and 1000 (0 uses train_steps)")
         if self.train_batch_size < 1:
             raise ValueError("train_batch_size must be >= 1")
         if self.gradient_accumulation_steps < 1:
@@ -364,6 +370,18 @@ class DiffusionLoraConfig:
         )
 
 
+def resolve_train_steps(cfg: "DiffusionLoraConfig", n_images: int) -> int:
+    """The effective optimizer-step count for a run. When ``cfg.num_epochs`` is set (> 0),
+    one epoch is one full pass over the dataset in optimizer steps -- ceil(N / (batch x
+    grad_accum)) steps -- so the run is ``num_epochs`` such passes, capped at 100000. With
+    ``num_epochs == 0`` the explicit ``cfg.train_steps`` is used unchanged."""
+    if cfg.num_epochs > 0:
+        per_step = max(1, cfg.train_batch_size * cfg.gradient_accumulation_steps)
+        steps_per_epoch = max(1, math.ceil(n_images / per_step))
+        return min(100000, cfg.num_epochs * steps_per_epoch)
+    return cfg.train_steps
+
+
 def discover_image_caption_pairs(
     data_dir: str | os.PathLike[str],
     *,
@@ -601,6 +619,10 @@ def _write_lora_sidecar(sidecar_path: Path, cfg: DiffusionLoraConfig) -> None:
 _CONFIG_ALIASES = {
     "model_name": "base_model",
     "max_steps": "train_steps",
+    # The generic payload's num_epochs already matches the diffusion field name, but list it
+    # so the epochs override is threaded through the shared-payload path as explicitly as
+    # max_steps -> train_steps is.
+    "num_epochs": "num_epochs",
     "batch_size": "train_batch_size",
     "lora_r": "lora_rank",
     "lr_scheduler_type": "lr_scheduler",

studio/backend/models/training.py

@@ -691,6 +691,15 @@ class DiffusionTrainingStartRequest(BaseModel):
         1024, ge = 64, le = 2048, description = "Square training resolution (multiple of 8)"
     )
     train_steps: int = Field(500, ge = 1, le = 100000)
+    num_epochs: int = Field(
+        0,
+        ge = 0,
+        le = 1000,
+        description = (
+            "0 = use train_steps; > 0 overrides train_steps with epochs x "
+            "ceil(N / (batch x grad_accum)) optimizer steps over the N-image dataset"
+        ),
+    )
     learning_rate: float = Field(1e-4, gt = 0)
     train_batch_size: int = Field(1, ge = 1, le = 64)
     gradient_accumulation_steps: int = Field(1, ge = 1, le = 256)

studio/backend/tests/test_diffusion_lora_trainer.py

@@ -20,6 +20,7 @@ from core.training.diffusion_lora_trainer import (
     _config_from_dict,
     compute_sdxl_add_time_ids,
     discover_image_caption_pairs,
+    resolve_train_steps,
 )
 
 
@@ -102,6 +103,60 @@ def test_config_normalized_validation(kw):
         DiffusionLoraConfig(base_model = "b", data_dir = "d", output_dir = "o", **kw).normalized()
 
 
+def _cfg(**kw):
+    return DiffusionLoraConfig(base_model = "b", data_dir = "d", output_dir = "o", **kw)
+
+
+def test_resolve_train_steps_uses_train_steps_when_epochs_disabled():
+    # num_epochs == 0 leaves the explicit train_steps untouched, whatever the image count.
+    cfg = _cfg(train_steps = 300, num_epochs = 0)
+    assert resolve_train_steps(cfg, 20) == 300
+    assert resolve_train_steps(cfg, 1) == 300
+
+
+def test_resolve_train_steps_epochs_ceil_over_batch_and_grad_accum():
+    # One epoch = ceil(N / (batch x grad_accum)) optimizer steps; num_epochs multiplies it.
+    # 10 images, batch 4, grad_accum 1 -> ceil(10/4)=3 steps/epoch.
+    assert resolve_train_steps(_cfg(num_epochs = 1, train_batch_size = 4), 10) == 3
+    assert resolve_train_steps(_cfg(num_epochs = 5, train_batch_size = 4), 10) == 15
+    # grad_accum widens the effective batch: 100 images, batch 2, grad_accum 3 -> per_step=6,
+    # ceil(100/6)=17 steps/epoch, 2 epochs -> 34.
+    cfg = _cfg(num_epochs = 2, train_batch_size = 2, gradient_accumulation_steps = 3)
+    assert resolve_train_steps(cfg, 100) == 34
+    # An exact multiple does not round up: 8 images / batch 4 -> 2 steps/epoch.
+    assert resolve_train_steps(_cfg(num_epochs = 3, train_batch_size = 4), 8) == 6
+
+
+def test_resolve_train_steps_single_image_dataset():
+    # A one-image dataset is one optimizer step per epoch, so num_epochs == steps.
+    assert resolve_train_steps(_cfg(num_epochs = 7, train_batch_size = 4), 1) == 7
+
+
+def test_resolve_train_steps_caps_at_100000():
+    # The run length is capped at 100000 even for absurd epoch counts (matches the request
+    # model's train_steps ceiling), so a huge epochs x dataset never overflows the loop.
+    cfg = _cfg(num_epochs = 1000, train_batch_size = 1)
+    assert resolve_train_steps(cfg, 10_000) == 100000
+
+
+def test_config_normalized_num_epochs_bounds():
+    # 0 (disabled) and the 1..1000 range normalise; out-of-range is rejected.
+    assert _cfg(num_epochs = 0).normalized().num_epochs == 0
+    assert _cfg(num_epochs = 1000).normalized().num_epochs == 1000
+    with pytest.raises(ValueError, match = "num_epochs"):
+        _cfg(num_epochs = -1).normalized()
+    with pytest.raises(ValueError, match = "num_epochs"):
+        _cfg(num_epochs = 1001).normalized()
+
+
+def test_config_from_dict_threads_num_epochs():
+    # num_epochs flows through the shared-payload adapter onto the diffusion field.
+    cfg = _config_from_dict(
+        {"base_model": "b", "data_dir": "d", "output_dir": "o", "num_epochs": 12}
+    )
+    assert cfg.num_epochs == 12
+
+
 def test_compute_sdxl_add_time_ids():
     assert compute_sdxl_add_time_ids(1024) == (1024, 1024, 0, 0, 1024, 1024)
 

studio/backend/tests/test_diffusion_training.py

@@ -316,6 +316,29 @@ def test_route_start_forwards_extra_training_knobs(client):
     assert client._fake.started_with["lora_target_modules"] == ["to_q", "to_v"]
 
 
+def test_route_start_forwards_num_epochs(client):
+    # Epochs mode: the frontend omits train_steps and sends num_epochs; it must reach the
+    # service so the trainer can resolve it against the dataset size.
+    body = {k: v for k, v in _BODY.items() if k != "train_steps"}
+    r = client.post("/api/train/diffusion/start", json = {**body, "num_epochs": 8})
+    assert r.status_code == 200, r.text
+    assert client._fake.started_with["num_epochs"] == 8
+
+
+def test_request_model_num_epochs_bounds():
+    # The request schema mirrors DiffusionLoraConfig's 0..1000 num_epochs range.
+    from pydantic import ValidationError
+
+    from models.training import DiffusionTrainingStartRequest
+
+    base = {"base_model": "b", "data_dir": "d", "output_dir": "o"}
+    assert DiffusionTrainingStartRequest(**base).num_epochs == 0  # default = use train_steps
+    assert DiffusionTrainingStartRequest(**base, num_epochs = 1000).num_epochs == 1000
+    for bad in (-1, 1001):
+        with pytest.raises(ValidationError):
+            DiffusionTrainingStartRequest(**base, num_epochs = bad)
+
+
 def test_route_start_rejects_uncontained_paths(client):
     # An absolute path outside the Studio dataset roots is a 400, not silently accepted.
     r = client.post("/api/train/diffusion/start", json = {**_BODY, "data_dir": "/etc"})

studio/frontend/src/features/images/api.ts

@@ -270,6 +270,9 @@ export interface DiffusionTrainingStartRequest {
   instance_prompt?: string | null;
   resolution?: number;
   train_steps?: number;
+  // 0 or omitted uses train_steps. > 0 overrides train_steps with that many epochs
+  // (full passes over the dataset, in optimizer steps).
+  num_epochs?: number;
   learning_rate?: number;
   train_batch_size?: number;
   gradient_accumulation_steps?: number;

studio/frontend/src/features/images/train/diffusion-train-panel.tsx

@@ -210,6 +210,10 @@ export function DiffusionTrainPanel({
   const [instancePrompt, setInstancePrompt] = useState("");
 
   const [steps, setSteps] = useState(500);
+  // Run length is set in either steps or epochs; the trainer resolves epochs -> steps once
+  // the dataset size is known (num_epochs overrides train_steps on the backend).
+  const [durationUnit, setDurationUnit] = useState<"steps" | "epochs">("steps");
+  const [epochs, setEpochs] = useState(10);
   const [learningRate, setLearningRate] = useState(family?.defaults.lr ?? 0.0001);
   const [rank, setRank] = useState(family?.defaults.rank ?? 16);
   const [resolution, setResolution] = useState(family?.defaults.resolution ?? 768);
@@ -553,7 +557,11 @@ export function DiffusionTrainPanel({
       );
       return;
     }
-    if (steps < 1) return toast.error("Steps must be at least 1.");
+    if (durationUnit === "epochs") {
+      if (epochs < 1) return toast.error("Epochs must be at least 1.");
+    } else if (steps < 1) {
+      return toast.error("Steps must be at least 1.");
+    }
     if (rank < 1) return toast.error("LoRA rank must be at least 1.");
     if (resolution < 64 || resolution % 8 !== 0) {
       return toast.error("Resolution must be a multiple of 8 and at least 64.");
@@ -577,7 +585,10 @@ export function DiffusionTrainPanel({
         output_dir: outputDir.trim(),
         instance_prompt: instancePrompt.trim() || undefined,
         resolution,
-        train_steps: steps,
+        // Epochs mode overrides train_steps on the backend, so send num_epochs and omit
+        // train_steps (the backend default is unused when num_epochs > 0).
+        train_steps: durationUnit === "epochs" ? undefined : steps,
+        num_epochs: durationUnit === "epochs" ? epochs : undefined,
         learning_rate: learningRate,
         train_batch_size: batchSize,
         gradient_accumulation_steps: gradAccum,
@@ -610,6 +621,8 @@ export function DiffusionTrainPanel({
     instancePrompt,
     resolution,
     steps,
+    durationUnit,
+    epochs,
     learningRate,
     batchSize,
     gradAccum,
@@ -690,6 +703,40 @@ export function DiffusionTrainPanel({
     </div>
   );
 
+  // Run length: a number paired with a compact unit select (Steps / Epochs). Epochs mode
+  // trains for that many full passes over the dataset; the backend resolves it to steps.
+  const durationField = (
+    <div className="grid gap-1.5">
+      <Label className="text-xs">{durationUnit === "epochs" ? "Epochs" : "Steps"}</Label>
+      <div className="flex gap-1.5">
+        <Input
+          type="number"
+          min={1}
+          value={durationUnit === "epochs" ? epochs : steps}
+          onChange={(e) => {
+            settingsDirty.current = true;
+            const n = Number(e.target.value) || 1;
+            if (durationUnit === "epochs") setEpochs(n);
+            else setSteps(n);
+          }}
+          className="h-8 min-w-0 flex-1 text-xs"
+        />
+        <select
+          value={durationUnit}
+          onChange={(e) => {
+            settingsDirty.current = true;
+            setDurationUnit(e.target.value as "steps" | "epochs");
+          }}
+          className="h-8 w-24 rounded-md border border-input bg-background px-2 text-xs"
+          aria-label="Run length unit"
+        >
+          <option value="steps">Steps</option>
+          <option value="epochs">Epochs</option>
+        </select>
+      </div>
+    </div>
+  );
+
   const precisionLabel = (m: "nf4" | "bf16" | "int8" | "fp8" | "auto"): string => {
     if (m === "auto") return "Auto (recommended)";
     if (m === "nf4") return "nf4 (4-bit QLoRA, lowest VRAM)";
@@ -704,7 +751,7 @@ export function DiffusionTrainPanel({
   const trainingSettings = (
     <div className="flex flex-col gap-5">
       <div className="grid grid-cols-2 gap-x-6 gap-y-4 lg:grid-cols-3">
-        {numberField("Steps", steps, setSteps, 1)}
+        {durationField}
         {numberField("LoRA rank", rank, setRank, 1)}
         {numberField("Resolution", resolution, setResolution, 512, { min: 64, step: 64 })}
         {numberField("Batch", batchSize, setBatchSize, 1)}
@@ -1292,9 +1339,12 @@ export function DiffusionTrainPanel({
               finishes first.
             </AlertDialogDescription>
           </AlertDialogHeader>
-          {/* items-center + a real label on the destructive action: a bare "Stop" rendered
-              as a stubby pill between two wide ones and read as misaligned. */}
-          <AlertDialogFooter className="items-center">
+          {/* flex-wrap keeps all three buttons visible when the sm:flex-row row is wider than
+              the dialog at narrow widths (down to ~480px); it wraps instead of clipping the
+              last button past the right edge. items-center + a real label on the destructive
+              action: a bare "Stop" rendered as a stubby pill between two wide ones and read
+              as misaligned. */}
+          <AlertDialogFooter className="flex-wrap items-center">
             <AlertDialogCancel>Continue training</AlertDialogCancel>
             <AlertDialogAction variant="destructive" onClick={() => void onStop(false)}>
               Stop without saving