diff --git a/studio/backend/core/training/diffusion_training_service.py b/studio/backend/core/training/diffusion_training_service.py index 166d2a9193..bdefcfe25b 100644 --- a/studio/backend/core/training/diffusion_training_service.py +++ b/studio/backend/core/training/diffusion_training_service.py @@ -34,8 +34,7 @@ _TERMINAL = ("complete", "error") def _default_target(*, event_queue: Any, stop_queue: Any, config: dict) -> None: # Imported lazily so this module (and the route layer) stays torch-free at import. from .diffusion_lora_trainer import run_diffusion_training_process - - run_diffusion_training_process(event_queue=event_queue, stop_queue=stop_queue, config=config) + run_diffusion_training_process(event_queue = event_queue, stop_queue = stop_queue, config = config) def _idle_state() -> dict[str, Any]: @@ -94,20 +93,23 @@ class DiffusionTrainingService: if self._proc is not None and self._proc.is_alive(): raise RuntimeError("A diffusion training job is already running.") if self._pump is not None and self._pump.is_alive(): - self._pump.join(timeout=5.0) + self._pump.join(timeout = 5.0) job_id = uuid.uuid4().hex event_queue = self._ctx.Queue() self._stop_queue = self._ctx.Queue() self._proc = self._ctx.Process( - target=self._target, - kwargs={"event_queue": event_queue, "stop_queue": self._stop_queue, "config": config}, - daemon=True, + target = self._target, + kwargs = { + "event_queue": event_queue, + "stop_queue": self._stop_queue, + "config": config, + }, + daemon = True, ) self._proc.start() try: from utils.process_lifetime import adopt_pid - adopt_pid(self._proc.pid) # bind to parent lifetime (no zombie on exit) except Exception: # noqa: BLE001 -- lifetime binding is best-effort pass @@ -115,11 +117,15 @@ class DiffusionTrainingService: now = time.time() self._state = _idle_state() self._state.update( - active=True, job_id=job_id, status="running", - message="Starting diffusion LoRA training...", started_at=now, updated_at=now, + active = True, + job_id = job_id, + status = "running", + message = "Starting diffusion LoRA training...", + started_at = now, + updated_at = now, ) self._pump = threading.Thread( - target=self._pump_loop, args=(event_queue, self._proc), daemon=True + target = self._pump_loop, args = (event_queue, self._proc), daemon = True ) self._pump.start() return job_id @@ -149,7 +155,7 @@ class DiffusionTrainingService: def _pump_loop(self, event_queue: Any, proc: Any) -> None: while True: try: - ev = event_queue.get(timeout=1.0) + ev = event_queue.get(timeout = 1.0) except Exception: # noqa: BLE001 -- Empty (timeout) or a closed queue if not proc.is_alive(): # Drain anything buffered, then decide if it exited cleanly. @@ -163,9 +169,10 @@ class DiffusionTrainingService: with self._lock: if self._state.get("status") not in ("completed", "stopped", "error"): self._state.update( - active=False, status="error", - message="Training process exited unexpectedly.", - updated_at=time.time(), + active = False, + status = "error", + message = "Training process exited unexpectedly.", + updated_at = time.time(), ) _ = drained return @@ -182,33 +189,33 @@ class DiffusionTrainingService: s = self._state s["updated_at"] = time.time() if etype == "model_load_started": - s.update(in_model_load=True, status="running", message="Loading base model...") + s.update(in_model_load = True, status = "running", message = "Loading base model...") if ev.get("num_images") is not None: s["num_images"] = ev.get("num_images") elif etype == "model_load_completed": - s.update(in_model_load=False, message="Training...") + s.update(in_model_load = False, message = "Training...") elif etype == "progress": s.update( - status="running", - step=ev.get("step", s["step"]), - total_steps=ev.get("total_steps", s["total_steps"]), - loss=ev.get("loss", s["loss"]), - avg_loss=ev.get("avg_loss", s["avg_loss"]), - learning_rate=ev.get("learning_rate", s["learning_rate"]), - message="Training...", + status = "running", + step = ev.get("step", s["step"]), + total_steps = ev.get("total_steps", s["total_steps"]), + loss = ev.get("loss", s["loss"]), + avg_loss = ev.get("avg_loss", s["avg_loss"]), + learning_rate = ev.get("learning_rate", s["learning_rate"]), + message = "Training...", ) elif etype == "complete": s.update( - active=False, - status="stopped" if ev.get("stopped") else "completed", - output_dir=ev.get("output_dir"), - lora_path=ev.get("lora_path"), - message="Stopped (partial adapter saved)." + active = False, + status = "stopped" if ev.get("stopped") else "completed", + output_dir = ev.get("output_dir"), + lora_path = ev.get("lora_path"), + message = "Stopped (partial adapter saved)." if ev.get("stopped") else "Training complete.", ) elif etype == "error": - s.update(active=False, status="error", message=str(ev.get("message", "error"))) + s.update(active = False, status = "error", message = str(ev.get("message", "error"))) _service: Optional[DiffusionTrainingService] = None diff --git a/studio/backend/models/training.py b/studio/backend/models/training.py index 85f0949810..9a7774d8be 100644 --- a/studio/backend/models/training.py +++ b/studio/backend/models/training.py @@ -675,28 +675,30 @@ class DiffusionTrainingStartRequest(BaseModel): rest carry the trainer's defaults. """ - model_config = ConfigDict(protected_namespaces=()) + model_config = ConfigDict(protected_namespaces = ()) - base_model: str = Field(..., description="HF repo id or local path to an SDXL pipeline") - data_dir: str = Field(..., description="Folder of training images (+ captions)") - output_dir: str = Field(..., description="Directory to write the LoRA .safetensors into") + base_model: str = Field(..., description = "HF repo id or local path to an SDXL pipeline") + data_dir: str = Field(..., description = "Folder of training images (+ captions)") + output_dir: str = Field(..., description = "Directory to write the LoRA .safetensors into") instance_prompt: Optional[str] = Field( - None, description="Dreambooth caption applied to images without their own caption" + None, description = "Dreambooth caption applied to images without their own caption" ) - resolution: int = Field(1024, ge=64, le=2048, description="Square training resolution (multiple of 8)") - train_steps: int = Field(500, ge=1, le=100000) - 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) - lora_rank: int = Field(16, ge=1, le=320) - lora_alpha: Optional[int] = Field(None, ge=1, le=640, description="Defaults to lora_rank") - lora_dropout: float = Field(0.0, ge=0.0, le=1.0) + resolution: int = Field( + 1024, ge = 64, le = 2048, description = "Square training resolution (multiple of 8)" + ) + train_steps: int = Field(500, ge = 1, le = 100000) + 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) + lora_rank: int = Field(16, ge = 1, le = 320) + lora_alpha: Optional[int] = Field(None, ge = 1, le = 640, description = "Defaults to lora_rank") + lora_dropout: float = Field(0.0, ge = 0.0, le = 1.0) seed: int = Field(42) mixed_precision: Literal["bf16", "fp16", "no"] = Field("bf16") - snr_gamma: Optional[float] = Field(5.0, description="Min-SNR loss weighting; null disables") + snr_gamma: Optional[float] = Field(5.0, description = "Min-SNR loss weighting; null disables") gradient_checkpointing: bool = Field(True) lr_scheduler: str = Field("constant") - lr_warmup_steps: int = Field(0, ge=0) + lr_warmup_steps: int = Field(0, ge = 0) center_crop: bool = Field(False) random_flip: bool = Field(True) caption_column: str = Field("text") diff --git a/studio/backend/routes/training.py b/studio/backend/routes/training.py index b2ea6dfaa4..c96189bef2 100644 --- a/studio/backend/routes/training.py +++ b/studio/backend/routes/training.py @@ -1031,8 +1031,7 @@ async def stream_training_progress( @router.post("/diffusion/start", response_model = DiffusionTrainingStartResponse) async def start_diffusion_training( - body: DiffusionTrainingStartRequest, - current_subject: str = Depends(get_current_subject), + body: DiffusionTrainingStartRequest, current_subject: str = Depends(get_current_subject) ): """Start an SDXL LoRA training job from an image + caption dataset.""" from core.training.diffusion_training_service import get_diffusion_training_service @@ -1047,8 +1046,11 @@ async def start_diffusion_training( raise HTTPException(status_code = 409, detail = str(e)) except Exception as e: raise log_and_http_error( - e, 500, "Failed to start diffusion training", - event = "diffusion_training.start_failed", log = logger, + e, + 500, + "Failed to start diffusion training", + event = "diffusion_training.start_failed", + log = logger, ) return DiffusionTrainingStartResponse(job_id = job_id, status = "running") @@ -1066,5 +1068,4 @@ async def stop_diffusion_training(current_subject: str = Depends(get_current_sub async def diffusion_training_status(current_subject: str = Depends(get_current_subject)): """Poll the current diffusion training job's status/progress (JSON).""" from core.training.diffusion_training_service import get_diffusion_training_service - return DiffusionTrainingStatusResponse(**get_diffusion_training_service().status()) diff --git a/studio/backend/tests/test_diffusion_training.py b/studio/backend/tests/test_diffusion_training.py index ae95e6ce7e..ffc231673c 100644 --- a/studio/backend/tests/test_diffusion_training.py +++ b/studio/backend/tests/test_diffusion_training.py @@ -32,8 +32,8 @@ class _FakeQueue: def put(self, x): self._q.put(x) - def get(self, timeout=None): - return self._q.get(timeout=timeout) # raises queue.Empty on timeout + def get(self, timeout = None): + return self._q.get(timeout = timeout) # raises queue.Empty on timeout def get_nowait(self): return self._q.get_nowait() @@ -50,7 +50,7 @@ class _FakeProc: self.pid = 4321 def start(self): - self._thread = threading.Thread(target=self._target, kwargs=self._kwargs, daemon=True) + self._thread = threading.Thread(target = self._target, kwargs = self._kwargs, daemon = True) self._thread.start() def is_alive(self): @@ -69,10 +69,24 @@ def _happy_target(*, event_queue, stop_queue, config): event_queue.put({"type": "model_load_started", "num_images": 3}) event_queue.put({"type": "model_load_completed"}) event_queue.put( - {"type": "progress", "step": 1, "total_steps": 2, "loss": 0.5, "avg_loss": 0.5, "learning_rate": 1e-4} + { + "type": "progress", + "step": 1, + "total_steps": 2, + "loss": 0.5, + "avg_loss": 0.5, + "learning_rate": 1e-4, + } ) event_queue.put( - {"type": "progress", "step": 2, "total_steps": 2, "loss": 0.4, "avg_loss": 0.45, "learning_rate": 1e-4} + { + "type": "progress", + "step": 2, + "total_steps": 2, + "loss": 0.4, + "avg_loss": 0.45, + "learning_rate": 1e-4, + } ) event_queue.put( { @@ -86,8 +100,10 @@ def _happy_target(*, event_queue, stop_queue, config): def _stoppable_target(*, event_queue, stop_queue, config): event_queue.put({"type": "model_load_completed"}) - stop_queue.get(timeout=5.0) # block until stop() signals - event_queue.put({"type": "complete", "output_dir": config["output_dir"], "lora_path": "x", "stopped": True}) + stop_queue.get(timeout = 5.0) # block until stop() signals + event_queue.put( + {"type": "complete", "output_dir": config["output_dir"], "lora_path": "x", "stopped": True} + ) def _crashing_target(*, event_queue, stop_queue, config): @@ -98,7 +114,11 @@ def _crashing_target(*, event_queue, stop_queue, config): _CFG = {"base_model": "b", "data_dir": "d", "output_dir": "/tmp/out", "train_steps": 2} -def _wait_status(svc, *terminal, timeout=3.0): +def _wait_status( + svc, + *terminal, + timeout = 3.0, +): end = time.time() + timeout while time.time() < end: st = svc.status() @@ -109,7 +129,7 @@ def _wait_status(svc, *terminal, timeout=3.0): def test_service_happy_path(): - svc = DiffusionTrainingService(ctx=_FakeCtx(), target=_happy_target) + svc = DiffusionTrainingService(ctx = _FakeCtx(), target = _happy_target) job_id = svc.start(dict(_CFG)) assert job_id st = _wait_status(svc, "completed") @@ -122,7 +142,7 @@ def test_service_happy_path(): def test_service_rejects_bad_config_before_spawn(): - svc = DiffusionTrainingService(ctx=_FakeCtx(), target=_happy_target) + svc = DiffusionTrainingService(ctx = _FakeCtx(), target = _happy_target) with pytest.raises(ValueError): svc.start({**_CFG, "train_steps": 0}) # Nothing was spawned; still idle. @@ -130,7 +150,7 @@ def test_service_rejects_bad_config_before_spawn(): def test_service_rejects_second_concurrent_job(): - svc = DiffusionTrainingService(ctx=_FakeCtx(), target=_stoppable_target) + svc = DiffusionTrainingService(ctx = _FakeCtx(), target = _stoppable_target) svc.start(dict(_CFG)) _wait_status(svc, "running") with pytest.raises(RuntimeError): @@ -140,7 +160,7 @@ def test_service_rejects_second_concurrent_job(): def test_service_stop_marks_stopped(): - svc = DiffusionTrainingService(ctx=_FakeCtx(), target=_stoppable_target) + svc = DiffusionTrainingService(ctx = _FakeCtx(), target = _stoppable_target) svc.start(dict(_CFG)) _wait_status(svc, "running") assert svc.stop() is True @@ -152,7 +172,7 @@ def test_service_stop_marks_stopped(): def test_service_crash_without_terminal_event_is_error(): - svc = DiffusionTrainingService(ctx=_FakeCtx(), target=_crashing_target) + svc = DiffusionTrainingService(ctx = _FakeCtx(), target = _crashing_target) svc.start(dict(_CFG)) st = _wait_status(svc, "error") assert st["status"] == "error" @@ -160,7 +180,7 @@ def test_service_crash_without_terminal_event_is_error(): def test_apply_event_transitions(): - svc = DiffusionTrainingService(ctx=_FakeCtx(), target=_happy_target) + svc = DiffusionTrainingService(ctx = _FakeCtx(), target = _happy_target) svc._apply_event({"type": "model_load_started", "num_images": 5}) assert svc.status()["in_model_load"] is True and svc.status()["num_images"] == 5 svc._apply_event({"type": "model_load_completed"}) @@ -212,25 +232,32 @@ def client(monkeypatch): "core.training.diffusion_training_service.get_diffusion_training_service", lambda: fake ) app = FastAPI() - app.include_router(training_router, prefix="/api/train") + app.include_router(training_router, prefix = "/api/train") app.dependency_overrides[get_current_subject] = lambda: "test-user" c = TestClient(app) c._fake = fake # type: ignore[attr-defined] return c -_BODY = {"base_model": "stabilityai/sdxl-turbo", "data_dir": "/data", "output_dir": "/out", "train_steps": 10} +_BODY = { + "base_model": "stabilityai/sdxl-turbo", + "data_dir": "/data", + "output_dir": "/out", + "train_steps": 10, +} def test_route_start_ok(client): - r = client.post("/api/train/diffusion/start", json=_BODY) + r = client.post("/api/train/diffusion/start", json = _BODY) assert r.status_code == 200, r.text assert r.json() == {"job_id": "job-123", "status": "running"} assert client._fake.started_with["base_model"] == "stabilityai/sdxl-turbo" def test_route_start_missing_required_is_422(client): - r = client.post("/api/train/diffusion/start", json={"base_model": "x"}) # no data_dir/output_dir + r = client.post( + "/api/train/diffusion/start", json = {"base_model": "x"} + ) # no data_dir/output_dir assert r.status_code == 422 @@ -239,7 +266,7 @@ def test_route_start_bad_config_maps_to_400(client, monkeypatch): raise ValueError("resolution must be a multiple of 8") client._fake.start = _raise # type: ignore[assignment] - r = client.post("/api/train/diffusion/start", json=_BODY) + r = client.post("/api/train/diffusion/start", json = _BODY) assert r.status_code == 400 assert "multiple of 8" in r.json()["detail"] @@ -249,12 +276,12 @@ def test_route_start_conflict_maps_to_409(client): raise RuntimeError("A diffusion training job is already running.") client._fake.start = _raise # type: ignore[assignment] - r = client.post("/api/train/diffusion/start", json=_BODY) + r = client.post("/api/train/diffusion/start", json = _BODY) assert r.status_code == 409 def test_route_status_and_stop(client): - client.post("/api/train/diffusion/start", json=_BODY) + client.post("/api/train/diffusion/start", json = _BODY) s = client.get("/api/train/diffusion/status") assert s.status_code == 200 and s.json()["status"] == "running" st = client.post("/api/train/diffusion/stop")