diff --git a/backend/handlers/generation_handler.py b/backend/handlers/generation_handler.py index 123092c1..a46ffbc6 100644 --- a/backend/handlers/generation_handler.py +++ b/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() diff --git a/backend/services/patches/diffusion_stage_cache.py b/backend/services/patches/diffusion_stage_cache.py index 7882b813..d60dae99 100644 --- a/backend/services/patches/diffusion_stage_cache.py +++ b/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 diff --git a/backend/state/app_settings.py b/backend/state/app_settings.py index f3a3e987..eb7666ff 100644 --- a/backend/state/app_settings.py +++ b/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 diff --git a/backend/tests/test_diffusion_stage_cache.py b/backend/tests/test_diffusion_stage_cache.py index ee7e287f..14a3386f 100644 --- a/backend/tests/test_diffusion_stage_cache.py +++ b/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: diff --git a/frontend/generated/backend-openapi.json b/frontend/generated/backend-openapi.json index f439beab..4e296c05 100644 --- a/frontend/generated/backend-openapi.json +++ b/frontend/generated/backend-openapi.json @@ -3449,7 +3449,7 @@ "title": "Activeltxmodelid" }, "diffusionStageCacheEnabled": { - "default": false, + "default": true, "title": "Diffusionstagecacheenabled", "type": "boolean" }, diff --git a/frontend/generated/backend-openapi.ts b/frontend/generated/backend-openapi.ts index 1c29d9fb..110fd1d8 100644 --- a/frontend/generated/backend-openapi.ts +++ b/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; /**