diff --git a/studio/backend/models/training.py b/studio/backend/models/training.py index 29ddb3ff7a..11cc3e09fc 100644 --- a/studio/backend/models/training.py +++ b/studio/backend/models/training.py @@ -798,6 +798,16 @@ class DiffusionTrainingStartRequest(BaseModel): cache_variants: int = Field( 4, ge = 1, le = 16, description = "Frozen crop/flip variants per image in the latent cache" ) + cond_cache_dir: Optional[str] = Field( + None, + description = ( + "Directory for the PERSISTENT conditioning cache (latents + text embeddings), reused " + "across runs: a rerun whose images, captions and resolution are unchanged skips " + "loading the VAE and the multi-GB text encoders entirely. Studio-relative names are " + "resolved under the Studio outputs root and absolute paths must stay inside it. " + "null or blank keeps the in-memory cache, which is rebuilt every run." + ), + ) compile_transformer: Literal["off", "on", "auto"] = Field( "auto", description = "Regional torch.compile of the transformer blocks" ) diff --git a/studio/backend/routes/training.py b/studio/backend/routes/training.py index 0831955437..d41c0b6210 100644 --- a/studio/backend/routes/training.py +++ b/studio/backend/routes/training.py @@ -1375,6 +1375,11 @@ async def start_diffusion_training( from utils.paths import resolve_output_dir config["data_dir"] = str(_resolve_diffusion_data_dir(config["data_dir"])) config["output_dir"] = str(resolve_output_dir(config["output_dir"])) + # The persistent conditioning cache is another directory the TRAINER writes to, so it gets + # the same containment as output_dir rather than the trainer's cwd. Blank/None means the + # in-memory cache (the trainer's own "off"), so it must not resolve to the outputs root. + cond_cache = str(config.get("cond_cache_dir") or "").strip() + config["cond_cache_dir"] = str(resolve_output_dir(cond_cache)) if cond_cache else None except ValueError as e: raise HTTPException(status_code = 400, detail = str(e)) diff --git a/studio/backend/tests/test_diffusion_training.py b/studio/backend/tests/test_diffusion_training.py index e2331e3b93..8f0bd01f03 100644 --- a/studio/backend/tests/test_diffusion_training.py +++ b/studio/backend/tests/test_diffusion_training.py @@ -1911,3 +1911,26 @@ def test_gpu_load_admission_and_reserve_exclude_each_other(): svc.reserve() svc.reserve() svc.unreserve() + + +def test_route_start_carries_and_contains_the_conditioning_cache_dir(client): + # The trainer's persistent conditioning cache (cond_cache_dir) skips loading the VAE and the + # multi-GB text encoders on a rerun, but the start schema omitted the field, so Pydantic + # dropped it silently and every API-driven run fell back to the in-memory cache. It also has to + # be contained like output_dir: the trainer subprocess would otherwise resolve it against its + # own cwd. + from pathlib import Path + + r = client.post("/api/train/diffusion/start", json = {**_BODY, "cond_cache_dir": "cond-cache"}) + assert r.status_code == 200, r.text + resolved = client._fake.started_with["cond_cache_dir"] + assert Path(resolved).is_absolute() + assert Path(resolved).name == "cond-cache" + + # Omitted or blank keeps the in-memory cache rather than resolving to the outputs root. + r = client.post("/api/train/diffusion/start", json = _BODY) + assert r.status_code == 200, r.text + assert client._fake.started_with["cond_cache_dir"] is None + r = client.post("/api/train/diffusion/start", json = {**_BODY, "cond_cache_dir": " "}) + assert r.status_code == 200, r.text + assert client._fake.started_with["cond_cache_dir"] is None