From 21052db120ff7975a50dc920e02730342f0ab989 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Sun, 12 Jul 2026 13:42:24 +0000 Subject: [PATCH] Publish image generation active state before pre-denoise setup generate() assigned self._gen only at the pipe() call, after deferred compile, LoRA resolution/application, and ControlNet download/build had run. Across that setup window generate_progress() reported inactive even though _generate_lock was held, so a reloaded page's mount probe showed idle and let a second generate queue behind the first. Publish an active step-0 _GenState the moment the generation lock is acquired, before the setup work, and clear it in the outer finally so a setup-time error cannot leave the UI stuck active. Mirrors the video backend's queued phase and the training start guard. --- studio/backend/core/inference/diffusion.py | 14 +++++ .../backend/tests/test_diffusion_backend.py | 62 +++++++++++++++++++ 2 files changed, 76 insertions(+) diff --git a/studio/backend/core/inference/diffusion.py b/studio/backend/core/inference/diffusion.py index 2473c5defa..47dd293024 100644 --- a/studio/backend/core/inference/diffusion.py +++ b/studio/backend/core/inference/diffusion.py @@ -2216,6 +2216,15 @@ class DiffusionBackend: raise RuntimeError(DIFFUSION_NOT_LOADED_MSG) # Register under _lock so unload()/a load can signal THIS generation. self._active_generate_cancel = cancel + # Publish an active (step 0) progress state the moment the lock is held, BEFORE + # the slow pre-denoise setup (deferred compile, LoRA resolution/application, + # ControlNet download/build). Without this generate_progress() reports inactive + # across that window, so a reload's mount probe shows idle even though this + # generation holds _generate_lock; the user then starts a second generate that + # merely blocks behind this one (and can duplicate the result). The per-step + # callback swaps in its own _GenState at denoise start; this is the queued phase. + # Mirrors the video backend's queued state and the training start guard. + self._gen = _GenState(total_steps = steps) try: # The local `state` ref keeps the pipe alive even if unload() nulls _state. generator = torch.Generator(device = state.device) @@ -2550,6 +2559,11 @@ class DiffusionBackend: with self._lock: if self._active_generate_cancel is cancel: self._active_generate_cancel = None + # Drop the published progress state. The normal path already nulled it after + # the denoise; this also covers a setup-time error that skips that inner + # finally. Safe under _generate_lock: no other generation can have installed + # its own _gen while this one runs. + self._gen = None def generate_progress(self) -> dict[str, Any]: """Live per-step progress for an in-flight generation (lock-free read).""" diff --git a/studio/backend/tests/test_diffusion_backend.py b/studio/backend/tests/test_diffusion_backend.py index ddbf79e23c..358d6aafeb 100644 --- a/studio/backend/tests/test_diffusion_backend.py +++ b/studio/backend/tests/test_diffusion_backend.py @@ -482,6 +482,68 @@ def test_load_generate_unload_gguf(fake_runtime, tmp_path): assert backend.is_loaded is False +def test_generate_progress_active_during_setup(fake_runtime, tmp_path, monkeypatch): + # A generation must report active from the moment it holds the lock, BEFORE the slow + # pre-denoise setup (deferred compile / LoRA resolution / ControlNet build) runs. + # Otherwise a reload's mount probe sees idle while the lock is held and lets a second + # generate queue behind the first. _apply_loras runs inside that setup window, so probing + # generate_progress() from there exercises the gap the reviewer flagged. + (tmp_path / "model.gguf").write_bytes(b"weights") + backend = DiffusionBackend() + backend.load_pipeline( + str(tmp_path), + gguf_filename = "model.gguf", + base_repo = "base/repo", + family_override = "z-image", + hf_token = "hf_secret", + ) + + seen = {} + + def fake_apply(self, state, loras, cancel): + seen["progress"] = self.generate_progress() + + monkeypatch.setattr(DiffusionBackend, "_apply_loras", fake_apply) + + # Idle before the run. + assert backend.generate_progress()["active"] is False + + gen = backend.generate(prompt = "a sloth", steps = 4) + assert len(gen["images"]) == 1 + + # Active was published during setup, with the requested step total and step 0. + assert seen["progress"]["active"] is True + assert seen["progress"]["total_steps"] == 4 + assert seen["progress"]["step"] == 0 + + # And it is cleared once the generation returns. + assert backend.generate_progress()["active"] is False + + +def test_generate_progress_cleared_on_setup_error(fake_runtime, tmp_path, monkeypatch): + # A setup-time failure skips the inner finally that nulls _gen, so the outer finally must + # clear the published progress; otherwise a crashed generation leaves the UI stuck "active". + (tmp_path / "model.gguf").write_bytes(b"weights") + backend = DiffusionBackend() + backend.load_pipeline( + str(tmp_path), + gguf_filename = "model.gguf", + base_repo = "base/repo", + family_override = "z-image", + hf_token = "hf_secret", + ) + + def boom(self, state, loras, cancel): + raise RuntimeError("setup failed") + + monkeypatch.setattr(DiffusionBackend, "_apply_loras", boom) + + with pytest.raises(RuntimeError, match = "setup failed"): + backend.generate(prompt = "a sloth", steps = 4) + + assert backend.generate_progress()["active"] is False + + def test_dense_speed_auto_defers_compile_to_third_generation(fake_runtime, tmp_path, monkeypatch): # Dense models with speed unset stay bit-identical eager for the first two generations; the # 3rd engages the `default` profile mid-session (repeated use amortises the one-time compile),