diff --git a/studio/backend/core/inference/sd_cpp_backend.py b/studio/backend/core/inference/sd_cpp_backend.py index 09d144f58e..962c1d13aa 100644 --- a/studio/backend/core/inference/sd_cpp_backend.py +++ b/studio/backend/core/inference/sd_cpp_backend.py @@ -220,6 +220,11 @@ class _SdState: mode: str = "server" # Token kept so LoRA adapters selected at generate time can be fetched from the Hub. hf_token: Optional[str] = None + # The single-file GGUF basename this load committed. Kept so companion resolution + # (sd_cpp_text_encoders_for) reproduces the load identity -- some variants pick their + # encoder by filename (FLUX.2-klein-9B -> Qwen3-8B), and a local *klein-9B*.gguf carries + # that keyword only in the basename, not the repo id. + gguf_filename: Optional[str] = None def _memory_policy(memory_mode: Optional[str], cpu_offload: bool) -> str: @@ -570,6 +575,7 @@ class SdCppDiffusionBackend: server = server, mode = mode, hf_token = hf_token, + gguf_filename = gguf_filename, ) with self._lock: if self._load_token != _load_token: @@ -697,10 +703,15 @@ class SdCppDiffusionBackend: repos = [state.repo_id, state.base_repo] if fam.sd_cpp_vae: repos.append(fam.sd_cpp_vae[0]) - # Same per-variant encoder selection as _asset_specs, keyed on the loaded repo id, so the - # cache-deletion guard protects the encoder repo this load actually downloaded (the 9B - # variant's Qwen3-8B, not the 4B default). - repos.extend(terepo for terepo, _f, _k in sd_cpp_text_encoders_for(fam, state.repo_id)) + # Same per-variant encoder selection as _asset_specs, keyed on the loaded repo id AND + # GGUF filename, so the cache-deletion guard protects the encoder repo this load actually + # downloaded (the 9B variant's Qwen3-8B, not the 4B default) -- a local *klein-9B*.gguf + # carries that keyword only in the basename, so dropping the filename would fall back to + # the 4B default and protect the wrong repo. + repos.extend( + terepo + for terepo, _f, _k in sd_cpp_text_encoders_for(fam, state.repo_id, state.gguf_filename) + ) return tuple(dict.fromkeys(r for r in repos if r)) # ── Generate ─────────────────────────────────────────────────────────── diff --git a/studio/backend/tests/test_sd_cpp_backend.py b/studio/backend/tests/test_sd_cpp_backend.py index 9fdf628399..28536294de 100644 --- a/studio/backend/tests/test_sd_cpp_backend.py +++ b/studio/backend/tests/test_sd_cpp_backend.py @@ -99,6 +99,31 @@ def test_loaded_repo_ids_includes_native_companions(): assert b.loaded_repo_ids() == () +def test_loaded_repo_ids_tracks_variant_encoder_by_gguf_filename(): + # FLUX.2-klein-9B pairs with Qwen3-8B, and a local *klein-9B*.gguf carries that keyword only in + # the basename (not the repo id). loaded_repo_ids() must reproduce the committed load identity + # (repo id + GGUF filename), else it falls back to the 4B default encoder and the cache-deletion + # guard protects the wrong repo -- leaving the 8B encoder this load actually downloaded deletable + # while one-shot sd-cli still re-reads it every generation. + b = SdCppDiffusionBackend(engine = _FakeEngine()) + fam = detect_family("flux.2-klein") + b._state = bk._SdState( + repo_id = "local/my-klein-checkpoints", # no variant keyword; it lives in the filename + base_repo = fam.base_repo, + family = fam, + device = "cpu", + files = SdCppModelFiles(diffusion_model = "/m/FLUX.2-klein-9B-Q4_K_M.gguf"), + vae_format = fam.sd_cpp_vae_format, + sampling_method = fam.sd_cpp_sampling_method, + flow_shift = fam.sd_cpp_flow_shift, + mode = "oneshot", + gguf_filename = "FLUX.2-klein-9B-Q4_K_M.gguf", + ) + ids = set(b.loaded_repo_ids()) + assert "Comfy-Org/vae-text-encorder-for-flux-klein-9b" in ids # the 8B encoder this load pulled + assert "Comfy-Org/z_image_turbo" not in ids # the 4B default must not be protected instead + + class _FakeServer: """Stands in for SdCppServer: records the spawn + one img_gen per whole batch."""