diff --git a/studio/backend/core/inference/diffusion.py b/studio/backend/core/inference/diffusion.py index 57268291b0..8b40fdfc55 100644 --- a/studio/backend/core/inference/diffusion.py +++ b/studio/backend/core/inference/diffusion.py @@ -182,6 +182,43 @@ def _smart_base_repo(fam: DiffusionFamily, repo_id: str) -> str: return "black-forest-labs/FLUX.2-klein-4B" +def _expand_existing_local_path(value: str) -> str: + """Expand ``~`` in ``value`` when the expanded path exists locally. + + Round 14 P2 #11: the GGUF local path branch already calls + ``Path(repo_id).expanduser()``, but the full-diffusers-repo and + base-companion-repo paths passed the literal ``~/...`` straight + into ``from_pretrained``, which treated it as a Hub id and tried + to download. Keep behaviour identical for Hub ids (no leading + ``~`` -> return as-is) and for non-existent expansions (the + diffusers loader will surface its own ``not found`` error). + """ + if not value or not isinstance(value, str) or not value.startswith("~"): + return value + candidate = Path(value).expanduser() + if candidate.exists(): + return str(candidate) + return value + + +_HF_TOKEN_RE = re.compile(r"hf_[A-Za-z0-9]{20,}") + + +def _redact_hf_tokens(value: Any) -> Any: + """Scrub embedded ``hf_xxxxxxxx`` tokens out of a string before + logging. Round 14 P2 #9: callers can wrap an authenticated URL + (``https://hf_token@huggingface.co/...``) into ``repo_id`` / + ``base_repo`` / paths; the token would otherwise reach + structured-log sinks via the load-info / load-failure log lines. + Non-strings are returned unchanged so the helper is safe to + sprinkle through ``logger.info`` / ``logger.error`` argument + lists. + """ + if not isinstance(value, str): + return value + return _HF_TOKEN_RE.sub("", value) + + def _resolve_local_gguf_child(repo_root: Path, gguf_filename: str) -> Path: """Resolve a GGUF filename inside a local repo directory safely. @@ -288,12 +325,26 @@ def detect_family( needle = (repo_id or "").lower() if not needle: return None + # Normalise mixed separator spellings (``Qwen_Image-Edit-GGUF``, + # ``Qwen-Image_Edit-GGUF``, ``Qwen.Image.Edit-GGUF``) and the + # compact concatenation (``QwenImageEdit-GGUF``) so the + # _FAMILY_EXCLUDE deny lists do not need every permutation of + # ``-``, ``_``, ``.`` and run-together spellings to keep + # Qwen-Image-Edit out of the base Qwen-Image family (round 14 + # P2 #8). + needle_norm = re.sub(r"[^a-z0-9]+", "-", needle).strip("-") + needle_compact = re.sub(r"[^a-z0-9]+", "", needle) # Scan _FAMILIES first (GGUF-supported), then _FULL_REPO_FAMILIES # so a repo like ``stabilityai/stable-diffusion-xl-base-1.0`` is # auto-detected as SDXL instead of returning None. for fam in _FAMILIES + _FULL_REPO_FAMILIES: excludes = _FAMILY_EXCLUDE.get(fam.name, ()) - if any(e in needle for e in excludes): + if any( + e in needle + or re.sub(r"[^a-z0-9]+", "-", e).strip("-") in needle_norm + or re.sub(r"[^a-z0-9]+", "", e) in needle_compact + for e in excludes + ): continue if fam.name in needle: return fam @@ -352,9 +403,24 @@ class DiffusionBackend: self._family: Optional[DiffusionFamily] = None self._repo_id: Optional[str] = None self._gguf_path: Optional[str] = None + # Original ``gguf_filename`` the caller passed in, preserved + # so delete guards can compare against subdirectory variants + # like ``BF16/model.gguf`` or ``Q4_K_M/model.gguf`` instead + # of the collapsed basename (round 14 P1 #4). The basename + # alone (``model.gguf``) loses the quant directory and lets + # /delete-cached unlink the wrong file. + self._gguf_filename: Optional[str] = None self._base_repo: Optional[str] = None self._device: Optional[str] = None self._dtype: Optional[str] = None + # True when ``enable_model_cpu_offload()`` was applied on the + # loaded pipeline. Diffusers' offload moves the active + # submodule between CPU and GPU on each step, so a CUDA + # ``torch.Generator`` mismatches the CPU-resident embeddings + # and generation crashes mid-forward (round 14 P1 #6). When + # this is True, seeded generation has to use a CPU generator + # regardless of self._device. + self._cpu_offload_enabled: bool = False self._loaded_at: Optional[float] = None self._loading: bool = False self._last_error: Optional[str] = None @@ -387,6 +453,11 @@ class DiffusionBackend: # local HF cache layout (and the system username on default # POSIX layouts) to any authenticated Studio session. with self._lock: + # UI-facing collapsed basename. Full local path leaks the + # HF cache layout + system username; the original caller- + # supplied filename (e.g. ``BF16/model.gguf``) is kept + # separately as ``active_gguf_filename`` for delete + # guards. gguf_basename = Path(self._gguf_path).name if self._gguf_path else None # Expose BOTH the resident pipeline's id AND the pending # load target. Delete guards must check both: when model A @@ -398,7 +469,7 @@ class DiffusionBackend: # user just clicked. active_repo = self._repo_id active_base = self._base_repo - active_gguf = gguf_basename + active_gguf = self._gguf_filename pending_repo = self._pending_repo_id if self._loading else None pending_base = self._pending_base_repo if self._loading else None pending_gguf = self._pending_gguf_filename if self._loading else None @@ -414,6 +485,14 @@ class DiffusionBackend: if pending_repo and pending_repo != active_repo: ui_family = None ui_pipeline_class = None + # UI-facing ``gguf_filename`` collapses to the basename + # so the Images panel does not surface internal cache / + # variant directory names. Guard-facing ``active_*`` / + # ``pending_*`` retain the full caller-supplied filename + # so /delete-cached can compare against subdirectory + # variants like ``BF16/model.gguf`` (round 14 P1 #4-5). + ui_gguf = pending_gguf or active_gguf + ui_gguf_basename = Path(ui_gguf).name if ui_gguf else None return { "is_loaded": self._pipe is not None, "is_loading": self._loading, @@ -421,7 +500,7 @@ class DiffusionBackend: "family": ui_family, "pipeline_class": ui_pipeline_class, "base_repo": pending_base or active_base, - "gguf_filename": pending_gguf or active_gguf, + "gguf_filename": ui_gguf_basename, # Guard-facing fields: every repo / path / GGUF # filename the backend owns RIGHT NOW. Delete routes # iterate both, paired so the variant-filename check @@ -550,9 +629,11 @@ class DiffusionBackend: # success. self._pending_repo_id = repo_id self._pending_base_repo = base_repo - self._pending_gguf_filename = ( - Path(gguf_filename).name if gguf_filename else None - ) + # Store the caller's full ``gguf_filename`` (e.g. + # ``BF16/model.gguf``) so the variant-aware delete + # guards have the subdirectory info. The UI side of + # status() still collapses to the basename for display. + self._pending_gguf_filename = gguf_filename if gguf_filename else None try: pipeline_cls = getattr(diffusers, fam.pipeline_class, None) if pipeline_cls is None: @@ -595,11 +676,15 @@ class DiffusionBackend: "or load a full diffusers repo (base_repo only " "applies when picking a GGUF quant)." ) - effective_base = repo_id + # ``~/models/my-flux`` must be expanded so + # diffusers' from_pretrained does not pass the + # literal tilde through to ``os.path.isdir`` and + # fall back to the Hub (round 14 P2 #11). + effective_base = _expand_existing_local_path(repo_id) with self._lock: self._pending_base_repo = effective_base elif base_repo: - effective_base = base_repo + effective_base = _expand_existing_local_path(base_repo) # Refresh pending so delete guards see the actual # base, not just caller-supplied None. with self._lock: @@ -608,13 +693,19 @@ class DiffusionBackend: effective_base = _smart_base_repo(fam, repo_id) with self._lock: self._pending_base_repo = effective_base + # ``repo_id`` / ``effective_base`` are user-supplied + # strings that can embed an ``hf_xxxxx`` token via a + # URL-style path (``https://hf_token@huggingface.co/...``). + # Scrub them BEFORE the logger formats the line so the + # token never reaches structured-log sinks (round 14 + # P2 #9). logger.info( "Loading diffusion model %s (family=%s, device=%s, dtype=%s, base=%s)", - repo_id, + _redact_hf_tokens(repo_id), fam.name, device, dtype, - effective_base, + _redact_hf_tokens(effective_base), ) transformer = None @@ -685,9 +776,11 @@ class DiffusionBackend: self._family = None self._repo_id = None self._gguf_path = None + self._gguf_filename = None self._base_repo = None self._device = None self._dtype = None + self._cpu_offload_enabled = False self._loaded_at = None _release(old) old = None @@ -735,6 +828,9 @@ class DiffusionBackend: pipe_kwargs["token"] = hf_token pipe = None + cpu_offload_enabled = bool( + enable_model_cpu_offload and device == "cuda" + ) try: pipe = pipeline_cls.from_pretrained(effective_base, **pipe_kwargs) # Device placement / offload can ALSO raise after @@ -746,7 +842,7 @@ class DiffusionBackend: # the next load attempt. Explicitly release both # pipe and transformer in the same try (round 13 # P2 #11). - if enable_model_cpu_offload and device == "cuda": + if cpu_offload_enabled: pipe.enable_model_cpu_offload() else: pipe.to(device) @@ -765,9 +861,14 @@ class DiffusionBackend: self._family = fam self._repo_id = repo_id self._gguf_path = local_gguf_path + # Preserve the full caller-supplied filename, not + # just the basename, so per-variant delete guards + # see ``BF16/model.gguf`` (round 14 P1 #4). + self._gguf_filename = gguf_filename if gguf_filename else None self._base_repo = effective_base self._device = device self._dtype = str(dtype).replace("torch.", "") + self._cpu_offload_enabled = cpu_offload_enabled self._loaded_at = time.time() # Clear loading + pending here, BEFORE returning, # so the response payload reports the resident @@ -813,7 +914,11 @@ class DiffusionBackend: # Use ``logger.error`` with the already-scrubbed # message and exc_info=False so the bearer token # cannot leak through structured logging sinks. - logger.error("Diffusion load failed for %s: %s", repo_id, exc_msg) + logger.error( + "Diffusion load failed for %s: %s", + _redact_hf_tokens(repo_id), + exc_msg, + ) raise RuntimeError( f"Failed to load diffusion model: {exc_msg}" ) from exc @@ -845,9 +950,11 @@ class DiffusionBackend: self._family = None self._repo_id = None self._gguf_path = None + self._gguf_filename = None self._base_repo = None self._device = None self._dtype = None + self._cpu_offload_enabled = False self._loaded_at = None _release(old) old = None # noqa: F841 @@ -930,14 +1037,25 @@ class DiffusionBackend: raise RuntimeError("No diffusion model is loaded.") pipe = self._pipe device = self._device or "cpu" + cpu_offload_enabled = self._cpu_offload_enabled generator = None if seed is not None: # Match the device of the pipeline so determinism holds - # across reload cycles. For CPU offload, the noise still - # has to live on the device the diffusion forward runs on. - gen_device = ( - "cuda" if device == "cuda" and torch.cuda.is_available() else "cpu" - ) + # across reload cycles. When CPU offload is enabled + # (the default on CUDA hosts), diffusers shuttles each + # submodule between CPU and GPU on every step. A CUDA + # torch.Generator then mismatches the CPU-resident + # embeddings at the start of the forward and the run + # crashes (round 14 P1 #6). Use a CPU generator in that + # case; numerical determinism for the same seed is + # preserved because the seed feeds an int rather than a + # device-local RNG state. + if cpu_offload_enabled: + gen_device = "cpu" + else: + gen_device = ( + "cuda" if device == "cuda" and torch.cuda.is_available() else "cpu" + ) generator = torch.Generator(device = gen_device).manual_seed(int(seed)) call_kwargs: dict[str, Any] = { @@ -1126,34 +1244,49 @@ def _release_other_gpu_owners_for_diffusion() -> None: # higher-level guard) cannot still kill an active export. try: from core.export import get_export_backend # type: ignore - - exp = get_export_backend() - if getattr(exp, "current_checkpoint", None): - is_export_active_fn = getattr(exp, "is_export_active", None) - export_is_active = False - if is_export_active_fn is not None: - try: - export_is_active = bool(is_export_active_fn()) - except Exception: - # Unverifiable status -> treat as 'might be - # active' and refuse to touch the subprocess. - export_is_active = True - if export_is_active: - logger.info( - "Skipping export shutdown for diffusion load: " - "is_export_active=True (route layer should have " - "rejected this request with 409)" - ) - else: - logger.info( - "Shutting down idle export subprocess before diffusion load" - ) - exp._shutdown_subprocess() - exp.current_checkpoint = None - exp.is_vision = False - exp.is_peft = False except Exception as exc: - logger.debug("export unload skipped: %s", exc) + logger.debug("export module not importable: %s", exc) + return + + try: + exp = get_export_backend() + except Exception as exc: + logger.debug("export backend not available: %s", exc) + return + + is_export_active_fn = getattr(exp, "is_export_active", None) + if is_export_active_fn is not None: + try: + export_is_active = bool(is_export_active_fn()) + except Exception: + # Unverifiable status -> treat as 'might be active' and + # refuse so a direct backend caller (test / script / + # future route that forgot the higher-level 409 guard) + # cannot still terminate an in-flight export. + export_is_active = True + if export_is_active: + # Round 14 P2 #10: the prior behaviour logged a warning + # and continued, so direct ``DiffusionBackend.load_model`` + # callers (tests, scripts) silently bypassed the route + # layer's 409. Hard-refuse instead so any code path that + # reaches this helper while an export is active sees the + # same failure mode the route returns. + raise RuntimeError( + "An export job is currently active. Stop the export " + "job before loading a diffusion image model." + ) + + if getattr(exp, "current_checkpoint", None): + try: + logger.info( + "Shutting down idle export subprocess before diffusion load" + ) + exp._shutdown_subprocess() + exp.current_checkpoint = None + exp.is_vision = False + exp.is_peft = False + except Exception as exc: + logger.debug("idle export shutdown failed: %s", exc) # Note: active training is *not* stopped here. The route layer # (`_raise_if_training_active` in routes/inference.py) refuses diff --git a/studio/backend/core/inference/llama_cpp.py b/studio/backend/core/inference/llama_cpp.py index 51da54789d..e017fb70ed 100644 --- a/studio/backend/core/inference/llama_cpp.py +++ b/studio/backend/core/inference/llama_cpp.py @@ -2616,9 +2616,66 @@ class LlamaCppBackend: Returns True if server started and health check passed. """ - # Serialise the whole load so concurrent /load calls never - # leave two llama-server processes alive (#5401 / #5161). Does - # not block /unload, /status, /load-progress. + # Publish ``_loading_model_identifier`` BEFORE any phase of + # the load can begin and clear it AFTER the load fully settles + # (success or failure, including the duplicate-state fast path + # and every internal early ``return False``). Round 14 P1 #2: + # the prior inline try/finally only wrapped the download, so + # /delete-cached and the cross-workload handoff helpers saw + # the backend as idle once the GGUF bytes had landed but the + # subprocess had not yet spawned. Mark the load as pending + # for the entire duration -- download, metadata read, + # VRAM settle, process spawn, health check, audio probe. + self._loading_model_identifier = model_identifier + try: + # Serialise the whole load so concurrent /load calls never + # leave two llama-server processes alive (#5401 / #5161). + # Does not block /unload, /status, /load-progress. + return self._load_model_impl( + gguf_path = gguf_path, + mmproj_path = mmproj_path, + hf_repo = hf_repo, + hf_variant = hf_variant, + hf_token = hf_token, + model_identifier = model_identifier, + is_vision = is_vision, + n_ctx = n_ctx, + chat_template_override = chat_template_override, + cache_type_kv = cache_type_kv, + speculative_type = speculative_type, + spec_draft_n_max = spec_draft_n_max, + n_threads = n_threads, + n_gpu_layers = n_gpu_layers, + n_parallel = n_parallel, + extra_args = extra_args, + ) + finally: + self._loading_model_identifier = None + + def _load_model_impl( + self, + *, + gguf_path: Optional[str] = None, + mmproj_path: Optional[str] = None, + hf_repo: Optional[str] = None, + hf_variant: Optional[str] = None, + hf_token: Optional[str] = None, + model_identifier: str, + is_vision: bool = False, + n_ctx: int = 4096, + chat_template_override: Optional[str] = None, + cache_type_kv: Optional[str] = None, + speculative_type: Optional[str] = None, + spec_draft_n_max: Optional[int] = None, + n_threads: Optional[int] = None, + n_gpu_layers: Optional[int] = None, + n_parallel: int = 1, + extra_args: Optional[List[str]] = None, + ) -> bool: + """Internal body of ``load_model``. Kept as a separate method + so ``load_model`` can wrap it in a single try/finally that + publishes ``_loading_model_identifier`` for the WHOLE load + instead of only the download window.""" with self._serial_load_lock: # Duplicate /load that raced past the route-level check # (the first one hadn't published _healthy=True yet). If the @@ -2693,40 +2750,25 @@ class LlamaCppBackend: # Scope HF_HUB_OFFLINE to the download block only when DNS is # dead; cleanup runs even on exception so a transient hiccup # at the start of one load cannot quarantine future loads. - # - # Publish ``_loading_model_identifier`` BEFORE entering the - # download so /delete-cached and the cross-workload handoff - # helpers can see a multi-GB pending load: previously they - # only consulted ``model_identifier``, which the success - # path sets later (see "Set identifier early" below). That - # left a window where the user could rmtree the cache the - # download was still writing to, or start /images/load - # while llama-server was about to come up on the same GPU. - # Cleared in ``finally`` so failed / cancelled loads do not - # leak the pending state. - self._loading_model_identifier = model_identifier - try: - if hf_repo: - with _hf_offline_if_dns_dead(): - model_path = self._download_gguf( + if hf_repo: + with _hf_offline_if_dns_dead(): + model_path = self._download_gguf( + hf_repo = hf_repo, + hf_variant = hf_variant, + hf_token = hf_token, + ) + # Auto-download mmproj for vision models + if is_vision and not mmproj_path: + mmproj_path = self._download_mmproj( hf_repo = hf_repo, - hf_variant = hf_variant, hf_token = hf_token, ) - # Auto-download mmproj for vision models - if is_vision and not mmproj_path: - mmproj_path = self._download_mmproj( - hf_repo = hf_repo, - hf_token = hf_token, - ) - elif gguf_path: - if not Path(gguf_path).is_file(): - raise FileNotFoundError(f"GGUF file not found: {gguf_path}") - model_path = gguf_path - else: - raise ValueError("Either gguf_path or hf_repo must be provided") - finally: - self._loading_model_identifier = None + elif gguf_path: + if not Path(gguf_path).is_file(): + raise FileNotFoundError(f"GGUF file not found: {gguf_path}") + model_path = gguf_path + else: + raise ValueError("Either gguf_path or hf_repo must be provided") # Set identifier early so _read_gguf_metadata can use it for DeepSeek detection self._model_identifier = model_identifier diff --git a/studio/backend/routes/export.py b/studio/backend/routes/export.py index 11c5004ced..d157d7d9f8 100644 --- a/studio/backend/routes/export.py +++ b/studio/backend/routes/export.py @@ -107,6 +107,42 @@ async def load_checkpoint( ), ) + backend = get_export_backend() + # Refuse to reload the export checkpoint while an export job + # is still running. ``ExportBackend.load_checkpoint`` would + # terminate the running subprocess in order to spawn a new + # one, silently corrupting the partial output the user is + # waiting on (round 13 P1 #1). Runs BEFORE the chat / + # diffusion unloads below: a 409 from this guard must not + # leave the user's chat or diffusion GPU owners freed for + # nothing (round 14 P1 #1). ``is_export_active`` may be + # absent on older / mocked backends; treat missing as "no + # async-job tracker available" and skip rather than + # fail-closed. + is_export_active_fn = getattr(backend, "is_export_active", None) + if is_export_active_fn is not None: + try: + export_is_active = bool(is_export_active_fn()) + except Exception as e: + logger.warning( + "Could not verify export status before export load: %s", e + ) + raise HTTPException( + status_code = 503, + detail = ( + "Could not verify export status before loading " + "an export checkpoint. Try again." + ), + ) from e + if export_is_active: + raise HTTPException( + status_code = 409, + detail = ( + "An export job is currently active. Stop the " + "export job before loading another checkpoint." + ), + ) + # Free GPU memory: shut down any chat backend before loading # the export checkpoint. Routes the unload through the shared # helper so we cover llama-server is_active=True and @@ -141,41 +177,6 @@ async def load_checkpoint( except Exception as e: logger.debug("diffusion unload skipped for export: %s", e) - backend = get_export_backend() - # Refuse to reload the export checkpoint while an export job - # is still running. ``ExportBackend.load_checkpoint`` would - # terminate the running subprocess in order to spawn a new - # one, silently corrupting the partial output the user is - # waiting on (round 13 P1 #1). Mirrors the symmetric guards - # already in place for chat / diffusion / training handoffs. - # ``is_export_active`` may be absent on older / mocked - # backends -- treat missing as "no async-job tracker - # available" -> skip rather than fail-closed; the - # surrounding chat / diffusion unloads have already run. - is_export_active_fn = getattr(backend, "is_export_active", None) - if is_export_active_fn is not None: - try: - export_is_active = bool(is_export_active_fn()) - except Exception as e: - logger.warning( - "Could not verify export status before export load: %s", e - ) - raise HTTPException( - status_code = 503, - detail = ( - "Could not verify export status before loading " - "an export checkpoint. Try again." - ), - ) from e - if export_is_active: - raise HTTPException( - status_code = 409, - detail = ( - "An export job is currently active. Stop the " - "export job before loading another checkpoint." - ), - ) - # load_checkpoint spawns and waits on a subprocess and can take # minutes. Run it in a worker thread so the event loop stays # free to serve the live log SSE stream concurrently. diff --git a/studio/backend/routes/models.py b/studio/backend/routes/models.py index a573f88233..ecf88f7a86 100644 --- a/studio/backend/routes/models.py +++ b/studio/backend/routes/models.py @@ -1940,6 +1940,27 @@ async def delete_finetuned_model( from routes.inference import get_llama_cpp_backend llama_backend = get_llama_cpp_backend() + # Pending HF GGUF download targeting this path: round 14 P1 #3. + # ``loading_model_identifier`` is set before the download starts + # and cleared after the subprocess settles, so the user cannot + # rmtree the directory llama.cpp is writing into mid-flight. + loading_identifier = getattr(llama_backend, "loading_model_identifier", None) + if ( + loading_identifier + and _loaded_model_matches_deleted_path( + loading_identifier, + target_path, + ) + and ( + not gguf_variant + or not getattr(llama_backend, "hf_variant", None) + or llama_backend.hf_variant.lower() == gguf_variant.lower() + ) + ): + raise HTTPException( + status_code = 409, + detail = "Cannot delete a model while it is loading", + ) if ( llama_backend.is_active and not llama_backend.is_loaded @@ -2745,14 +2766,24 @@ async def delete_cached_model( # Exact match only (case-insensitive). Prefix match would # block deleting unrelated ``org/model`` while # ``org/model-v2`` is loaded -- same surface the diffusion - # guard fixed in round 5. + # guard fixed in round 5. Per-variant deletes that target a + # DIFFERENT quant than the loaded one are allowed so the + # llama and diffusion paths stay symmetric (round 14 P1 #7). if loaded_id == needle and ( llama_backend.is_loaded or getattr(llama_backend, "is_active", False) ): - raise HTTPException( - status_code = 400, - detail = "Unload the model before deleting", + loaded_variant = (getattr(llama_backend, "hf_variant", None) or "").lower() + requested_variant = (variant or "").lower() + same_variant = ( + not requested_variant + or not loaded_variant + or requested_variant == loaded_variant ) + if same_variant: + raise HTTPException( + status_code = 400, + detail = "Unload the model before deleting", + ) except HTTPException: raise except Exception as e: diff --git a/studio/backend/tests/test_diffusion_backend.py b/studio/backend/tests/test_diffusion_backend.py index 43918fb17e..f642ff2cbd 100644 --- a/studio/backend/tests/test_diffusion_backend.py +++ b/studio/backend/tests/test_diffusion_backend.py @@ -1359,6 +1359,120 @@ def test_generate_image_with_metadata_returns_active_pipeline(monkeypatch): } +@pytest.mark.parametrize( + "repo_id", + [ + "unsloth/Qwen_Image-Edit-GGUF", + "unsloth/Qwen-Image_Edit-GGUF", + "unsloth/Qwen-ImageEdit-GGUF", + "unsloth/qwen-image_edit-2509-GGUF", + "unsloth/Qwen.Image.Edit-GGUF", + ], +) +def test_detect_family_qwen_image_edit_mixed_separators(repo_id): + """Round 14 P2 #8: every spelling of Qwen-Image-Edit must NOT + match the base Qwen-Image text-to-image family.""" + from core.inference.diffusion import detect_family + + assert detect_family(repo_id) is None + + +def test_redact_hf_tokens_removes_url_embedded_token(): + """Round 14 P2 #9: tokens embedded in user-supplied paths / + URLs must be scrubbed before logging.""" + from core.inference.diffusion import _redact_hf_tokens + + leaky = "https://hf_abcdefghij0123456789@huggingface.co/unsloth/FLUX.2-klein-4B-GGUF" + redacted = _redact_hf_tokens(leaky) + assert "hf_" not in redacted + assert "" in redacted + # Non-strings pass through unchanged so the helper is safe in + # logger argument lists where families / dtypes mix in. + assert _redact_hf_tokens(None) is None + assert _redact_hf_tokens(42) == 42 + + +def test_status_preserves_active_gguf_subdir(monkeypatch): + """Round 14 P1 #4: status() must surface the original caller- + supplied gguf_filename (``BF16/model.gguf``) instead of the + collapsed basename.""" + import core.inference.diffusion as d + + backend = d.DiffusionBackend() + backend._pipe = object() + backend._repo_id = "unsloth/FLUX.2-klein-4B-GGUF" + backend._gguf_path = "/cache/models/unsloth/FLUX.2-klein-4B-GGUF/BF16/model.gguf" + backend._gguf_filename = "BF16/model.gguf" + backend._family = d.DiffusionFamily( + name = "flux.2-klein", + pipeline_class = "Flux2KleinPipeline", + transformer_class = "Flux2Transformer2DModel", + base_repo = "black-forest-labs/FLUX.2-klein-4B", + aliases = (), + ) + + s = backend.status() + assert s["active_gguf_filename"] == "BF16/model.gguf" + # UI-facing field still collapses to the basename. + assert s["gguf_filename"] == "model.gguf" + + +def test_generator_uses_cpu_when_cpu_offload_enabled(monkeypatch): + """Round 14 P1 #6: seeded CUDA generation must NOT create a + CUDA torch.Generator when the pipeline was loaded with CPU + offload enabled, otherwise it crashes mid-forward.""" + import core.inference.diffusion as d + + backend = d.DiffusionBackend() + + class _FakePipe: + def __init__(self): + self.last_kwargs = None + + def __call__(self, **kwargs): + self.last_kwargs = kwargs + from PIL import Image + + return SimpleNamespace(images = [Image.new("RGB", (8, 8))]) + + fake_pipe = _FakePipe() + backend._pipe = fake_pipe + backend._device = "cuda" + backend._cpu_offload_enabled = True + + captured_devices: list[str] = [] + + class _FakeGenerator: + def __init__(self, device): + captured_devices.append(device) + + def manual_seed(self, seed): + return self + + class _FakeTorchCuda: + @staticmethod + def is_available(): + return True + + fake_torch = SimpleNamespace( + Generator = _FakeGenerator, cuda = _FakeTorchCuda + ) + monkeypatch.setitem(sys.modules, "torch", fake_torch) + + backend._generate_image_unlocked(prompt = "x", seed = 7, width = 8, height = 8) + assert captured_devices == ["cpu"] + + +def test_smart_base_repo_uses_windows_leaf_only_already_set_separator_round14(): + """Sanity: relative paths still work after the Windows fix.""" + from core.inference.diffusion import _smart_base_repo, detect_family + + repo = "owner/FLUX.2-klein-9B-GGUF" + fam = detect_family(repo) + assert fam is not None + assert _smart_base_repo(fam, repo) == "black-forest-labs/FLUX.2-klein-9B" + + def test_generate_image_with_metadata_blocks_concurrent_unload(monkeypatch): """Round 13 P2 #9: _generate_lock serialises the forward AND the meta snapshot, so a queued unload cannot wipe state in between.""" diff --git a/studio/backend/tests/test_diffusion_routes.py b/studio/backend/tests/test_diffusion_routes.py index d3659b89ee..2818576659 100644 --- a/studio/backend/tests/test_diffusion_routes.py +++ b/studio/backend/tests/test_diffusion_routes.py @@ -84,6 +84,9 @@ class _FakeBackend: "active_base_repo": ( "black-forest-labs/FLUX.2-klein" if self._loaded else None ), + # Round 14: guard-facing GGUF filename is now the full + # caller-supplied value, but this fake never sets one so + # both active and pending stay None. "active_gguf_filename": None, "pending_repo_id": None, "pending_base_repo": None,