Release reused diffusion cache before VAE decode

backend/handlers/generation_handler.py

@@ -96,11 +96,10 @@ def start_generation(self, generation_id: str) -> None:
             raise RuntimeError("No active GPU pipeline")
         self.state.generation_starting_since = None
 
-        # EXPERIMENTAL: push the live Settings toggle, then drop any transformer
-        # cached from the previous generation before this one starts -- otherwise
-        # it stays resident while this generation's own text encoder/VAE/etc.
-        # build, double-booking VRAM. See that module's GENERATION-SCOPED
-        # docstring section for the RTX 5090 repro (~42GB reported on a 32GB card).
+        # Push the live Settings toggle, then drop any transformer left by a
+        # cancelled/failed generation before this one starts. A normal two-stage
+        # generation retires the cache immediately after the reuse hit, before VAE
+        # decode; this is the defensive backstop. See diffusion_stage_cache.py.
         diffusion_stage_cache.set_enabled(self.state.app_settings.diffusion_stage_cache_enabled)
         diffusion_stage_cache.evict()
 

backend/services/patches/diffusion_stage_cache.py

@@ -30,18 +30,15 @@
 (``**kwargs: object,  # noqa: ARG002`` in single_gpu_model_builder.py) --
 confirmed inert for the path we cache, not assumed.
 
-GENERATION-SCOPED, not session-scoped (found live on an RTX 5090): letting the
-cache survive PAST the generation that built it collides with every other
-component that builds fresh per call too (text encoder, VAE, upsampler, audio
-decoder/vocoder) -- the next generation's text-encoder build then has to
-coexist in VRAM with the still-resident transformer from the PREVIOUS
-generation, instead of the transformer having already been freed by then.
-Observed: a second generation's peak VRAM was reported at ~41.8 GB on a
-31.82 GB card (Windows CUDA fell back to slow shared memory, backend liveness
-probe failed, total generation time regressed to 143s -- worse than no cache
-at all). Fix: ``handlers.generation_handler.GenerationHandler.start_generation``
-/``start_api_generation`` call :func:`evict` before marking a new generation as
-running, so the cache never survives past the generation it was built for.
+REUSE-SCOPED, not generation- or session-scoped: the cache exists only to bridge
+the first and second identical diffusion stages. A live 720p Apple Silicon run
+showed why the narrower lifetime matters: leaving the reused ~37 GiB transformer
+resident through tiled VAE decode raised MPS driver memory to 46.7 GiB versus
+44.5 GiB with the cache disabled, and left 35.4 GiB torch-allocated after output
+encoding. The cache now evicts immediately when a hit's denoising context exits,
+before decoder construction. The generation-start eviction remains a defensive
+backstop for a pipeline that builds a cacheable first stage but never reaches an
+identical second stage (cancellation, error, or a different stage configuration).
 
 NON-CACHEABLE TRANSITIONS also evict (found live on the same RTX 5090, IC-LoRA
 this time): IC-LoRA's ``use_lora_in_stage_2`` forces stage_2 onto the streaming
@@ -90,7 +87,7 @@
 produces a different key (cache miss, falls back to a normal rebuild), never
 a false hit.
 
-Toggle: ``AppSettings.diffusion_stage_cache_enabled`` (default off, surfaced in
+Toggle: ``AppSettings.diffusion_stage_cache_enabled`` (default on, surfaced in
 Settings next to Torch Compile). ``GenerationHandler.start_generation``/
 ``start_api_generation`` push the live setting into :func:`set_enabled` on
 every generation, so flipping it in Settings takes effect on the next
@@ -218,11 +215,18 @@ def evict() -> None:
         _evict_locked()
 
 
-def _mark_free() -> None:
-    """Mark that a caller is done with a cached transformer it checked out."""
+def _release(*, evict_after_use: bool) -> None:
+    """Release one checkout and optionally retire the now-consumed cache entry.
+
+    A hit means the cache fulfilled its only purpose: carrying one identical
+    transformer from the first diffusion stage into the second. Retire it while
+    leaving that second context so VAE/audio decode cannot overlap its weights.
+    """
     global _in_use
     with _lock:
         _in_use = max(0, _in_use - 1)
+        if evict_after_use:
+            _evict_locked()
 
 
 _orig_transformer_ctx = DiffusionStage._transformer_ctx  # noqa: SLF001
@@ -263,7 +267,7 @@ def _cached_transformer_ctx(self: DiffusionStage, **kwargs: object) -> Iterator[
     try:
         yield model
     finally:
-        _mark_free()
+        _release(evict_after_use=hit)
 
 
 DiffusionStage._transformer_ctx = _cached_transformer_ctx  # type: ignore[method-assign]  # noqa: SLF001

backend/state/app_settings.py

@@ -49,7 +49,7 @@ class SettingsPatchModel(SettingsBaseModel):
 
 class AppSettings(SettingsBaseModel):
     use_torch_compile: bool = False
-    diffusion_stage_cache_enabled: bool = False
+    diffusion_stage_cache_enabled: bool = True
     ltx_api_key: str = ""
     user_prefers_ltx_api_video_generations: bool = False
     fal_api_key: str = ""
@@ -126,7 +126,7 @@ def _is_settings_model_annotation(annotation: object) -> TypeGuard[type[Settings
 
 class SettingsResponse(SettingsBaseModel):
     use_torch_compile: bool = False
-    diffusion_stage_cache_enabled: bool = False
+    diffusion_stage_cache_enabled: bool = True
     has_ltx_api_key: bool = False
     user_prefers_ltx_api_video_generations: bool = False
     has_fal_api_key: bool = False

backend/tests/test_diffusion_stage_cache.py

@@ -84,6 +84,27 @@ def test_cache_hit_reuses_model_without_rebuilding() -> None:
     assert stage_1.build_count == 1
     assert stage_2.build_count == 0, "stage_2 should reuse stage_1's cached build"
     assert model_1 is model_2
+    assert model_1.freed_to == "meta", "a consumed cache hit must free before decoder construction"
+    assert dsc._cached_model is None
+    assert dsc._cached_key is None
+
+
+def test_third_identical_stage_rebuilds_after_consumed_hit() -> None:
+    stage_1 = _FakeStage(_single_gpu_builder("ckpt.safetensors"))
+    stage_2 = _FakeStage(_single_gpu_builder("ckpt.safetensors"))
+    stage_3 = _FakeStage(_single_gpu_builder("ckpt.safetensors"))
+
+    with dsc._cached_transformer_ctx(stage_1):
+        pass
+    with dsc._cached_transformer_ctx(stage_2):
+        pass
+    with dsc._cached_transformer_ctx(stage_3):
+        pass
+
+    assert stage_1.build_count == 1
+    assert stage_2.build_count == 0
+    assert stage_3.build_count == 1
+    assert dsc._cached_model is not None
 
 
 def test_cache_miss_on_different_checkpoint_evicts_old_model() -> None:

frontend/generated/backend-openapi.json

@@ -3449,7 +3449,7 @@
             "title": "Activeltxmodelid"
           },
           "diffusionStageCacheEnabled": {
-            "default": false,
+            "default": true,
             "title": "Diffusionstagecacheenabled",
             "type": "boolean"
           },

frontend/generated/backend-openapi.ts

@@ -1946,7 +1946,7 @@ export interface components {
             activeLtxModelId?: ("ltx-2.3-22b-distilled-1.1" | "ltx-2.3-22b-distilled") | null;
             /**
              * Diffusionstagecacheenabled
-             * @default false
+             * @default true
              */
             diffusionStageCacheEnabled: boolean;
             /**