diff --git a/studio/backend/routes/models.py b/studio/backend/routes/models.py index f9a6e5833c..82b018d741 100644 --- a/studio/backend/routes/models.py +++ b/studio/backend/routes/models.py @@ -304,11 +304,26 @@ def _has_non_gguf_weights(path: Path) -> bool: return False +def _local_pipeline_index(d: Path) -> bool: + """True when *d* is a standard diffusers PIPELINE root: component weights/configs live in + subdirs (``transformer/``, ``vae/``, ...) under a top-level ``model_index.json``, so + ``_is_model_directory`` (which wants a root config + loose weights) rejects it.""" + try: + return (d / "model_index.json").is_file() + except OSError: + return False + + def _scan_models_dir(models_dir: Path, *, limit: int | None = None) -> List[LocalModelInfo]: if not models_dir.exists() or not models_dir.is_dir(): return [] - _is_self_model = _is_model_directory(models_dir) + # A scan folder can point directly at a diffusers PIPELINE dir, not only at a parent of + # model repos. _is_model_directory rejects such a root (weights live in transformer/, vae/, + # ... not beside a root config.json), so without this the child scan below surfaces the + # component subdirs as bogus models and hides the real pipeline. The Images/Video load path + # loads a local pipeline dir, so admit the root as one model (task tagging classifies it). + _is_self_model = _is_model_directory(models_dir) or _local_pipeline_index(models_dir) if _is_self_model: try: @@ -343,7 +358,7 @@ def _scan_models_dir(models_dir: Path, *, limit: int | None = None) -> List[Loca # root, so the checks above miss it. The Images/Video load path accepts such a # local pipeline dir, so admit it here too (task tagging then classifies it via # _local_is_diffusers); otherwise it is hidden from the On Device picker. - has_pipeline_index = (child / "model_index.json").is_file() + has_pipeline_index = _local_pipeline_index(child) has_model_files = has_gguf or has_non_gguf_weights or has_config or has_pipeline_index except OSError: # Skip unreadable children rather than failing the scan. @@ -3261,6 +3276,25 @@ def _repo_gguf_task(repo_info) -> Optional[str]: return None +def _local_family_needles(model: "LocalModelInfo") -> tuple[str, ...]: + """Family-detection hints for a local (non-GGUF) checkpoint: its model id, display name, and + leaf directory name, plus -- for a bare single-file directory -- the sole checkpoint's + filename. A generically named folder holding one loadable ``qwen-image-*.safetensors`` / + ``ltx-*.safetensors`` identifies its family only from that filename, and the load route already + resolves that sole file via ``resolve_local_single_file``, so feed the same name here or a + task-scoped Images/Video picker (which rejects ``task: null``) hides the on-device model. Only + the basename is used (not the parent path), so a family token in a parent dir can't match.""" + needles = [model.model_id, model.display_name, Path(model.id).name] + try: + from core.inference.diffusion import resolve_local_single_file + single = resolve_local_single_file(model.path) + if single: + needles.append(single) + except Exception: + pass + return tuple(n for n in needles if n) + + def _local_model_task(model: "LocalModelInfo") -> Optional[str]: """Classify a local model into an HF pipeline task so the Images picker can filter. @@ -3293,12 +3327,8 @@ def _local_model_task(model: "LocalModelInfo") -> Optional[str]: try: from core.inference.video import _is_trusted_video_repo from core.inference.video_families import detect_video_family - for needle in (model.model_id, model.display_name, Path(model.id).name): - if ( - needle - and detect_video_family(needle) is not None - and _is_trusted_video_repo(path) - ): + for needle in _local_family_needles(model): + if detect_video_family(needle) is not None and _is_trusted_video_repo(path): return _VIDEO_GEN_TASK except Exception: pass @@ -3311,8 +3341,9 @@ def _local_is_diffusers(model: "LocalModelInfo") -> bool: ``_repo_is_diffusers`` heuristics: a full pipeline carries a top-level ``model_index.json``, while single-file / safetensors image checkpoints ship none, so fall back to the model id resolving to a known diffusion family (the same resolver the - Images backend loads from). Family detection uses the clean model id / name, not the - on-disk path, so a parent directory keyword can't spuriously match.""" + Images backend loads from). Family detection uses the clean model id / name and the sole + checkpoint's filename (via _local_family_needles), not the on-disk path, so a parent + directory keyword can't spuriously match while a filename-only family is still caught.""" try: p = Path(model.path) if p.is_dir() and (p / "model_index.json").is_file(): @@ -3321,20 +3352,21 @@ def _local_is_diffusers(model: "LocalModelInfo") -> bool: pass try: from core.inference.diffusion_families import detect_family - for needle in (model.model_id, model.display_name, Path(model.id).name): - if needle and detect_family(needle) is not None: + for needle in _local_family_needles(model): + if detect_family(needle) is not None: return True except Exception: pass # A single-file VIDEO checkpoint (LTX / Wan / Hunyuan .safetensors, no model_index.json) has no # pipeline index and no image family, so the checks above miss it. The video load route loads it # as a single_file (routes/video.py), so it must be surfaced or _local_model_task returns - # task=null and the picker hides it. Match clean id / name needles (not the raw path) so a - # parent-dir token can't spuriously match; _local_model_task then routes it to text-to-video. + # task=null and the picker hides it. Match clean id / name / checkpoint-filename needles (not + # the raw path) so a parent-dir token can't spuriously match; _local_model_task then routes it + # to text-to-video. try: from core.inference.video_families import detect_video_family - for needle in (model.model_id, model.display_name, Path(model.id).name): - if needle and detect_video_family(needle) is not None: + for needle in _local_family_needles(model): + if detect_video_family(needle) is not None: return True except Exception: pass @@ -3438,6 +3470,47 @@ def _repo_is_diffusers(repo_info) -> bool: return False +def _repo_pipeline_missing_denoiser(repo_info) -> bool: + """True for a diffusers-pipeline snapshot (root ``model_index.json``) whose denoiser + component (``transformer/`` or ``unet/``) carries NO weight file. This is the shape of a + companion-only prefetch: a GGUF image load pulls the base repo's VAE / text-encoder / + ``model_index.json`` into the cache but deliberately skips the multi-GB transformer (the GGUF + supplies it), so the snapshot has a pipeline manifest yet is not a loadable BF16 pipeline -- + ``from_pretrained`` on it re-downloads the missing shards. ``_cached_repo_partial`` misses this + (no cancel marker / .incomplete blob, and hf_hub_download writes no manifest), so ``/cached-models`` + would advertise it as fully on-device. The caller marks such rows partial. Best-effort: any scan + error reports not-missing so a glitch never hides a genuinely complete pipeline.""" + if not _repo_has_pipeline_index(repo_info): + return False + _DENOISER_DIRS = ("transformer", "unet") + _WEIGHT_SUFFIXES = (".safetensors", ".bin") + try: + for rev in repo_info.revisions: + snapshot = getattr(rev, "snapshot_path", None) + for f in rev.files: + name = str(getattr(f, "file_name", "") or "") + path = getattr(f, "file_path", None) + parts: tuple[str, ...] = () + if path is not None and snapshot is not None: + try: + parts = Path(path).relative_to(Path(snapshot)).parts + except ValueError: + parts = () + if not parts: + # No snapshot scoping (or file outside it): fall back to the recorded name, + # which may itself carry the component subdir (e.g. 'transformer/model...'). + parts = Path(name).parts + if ( + len(parts) >= 2 + and parts[0].lower() in _DENOISER_DIRS + and parts[-1].lower().endswith(_WEIGHT_SUFFIXES) + ): + return False + return True + except Exception: + return False + + def _cached_repo_partial(repo_id: str, repo_cache_dir: Optional[Path] = None) -> bool: """Whether the cached model snapshot is incomplete (cancelled/partial download). Reuses the hub inventory scan's snapshot-partial detector (cancel marker, legacy @@ -3517,7 +3590,12 @@ async def list_cached_models( ) key = repo_id.lower() existing = seen_lower.get(key) - is_partial = _cached_repo_partial(repo_id, Path(repo_info.repo_path)) + # A companion-only prefetch (root model_index.json + VAE / text-encoder but no + # transformer/ shards, pulled to back a GGUF load) is not a loadable BF16 + # pipeline; treat it as partial so the picker does not advertise it as on-device. + is_partial = _cached_repo_partial( + repo_id, Path(repo_info.repo_path) + ) or _repo_pipeline_missing_denoiser(repo_info) # Prefer the most COMPLETE snapshot, then largest. The picker drops partial # rows, so a partial copy in one cache root must not shadow a smaller complete # copy in another (size only breaks ties among equal completeness). diff --git a/studio/backend/tests/test_cached_gguf_routes.py b/studio/backend/tests/test_cached_gguf_routes.py index ee45a44085..4e1aa3a74d 100644 --- a/studio/backend/tests/test_cached_gguf_routes.py +++ b/studio/backend/tests/test_cached_gguf_routes.py @@ -445,7 +445,13 @@ def test_list_cached_models_tags_diffusers_pipeline_as_text_to_image(monkeypatch text-to-image so the chat picker hides it, while a plain checkpoint isn't.""" diffusion = _repo( "Tongyi-MAI/Z-Image-Turbo", - [_file("model_index.json", 1_000), _file("text_encoder/model.safetensors", 9_000)], + [ + _file("model_index.json", 1_000), + _file("text_encoder/model.safetensors", 9_000), + # A complete pipeline carries its denoiser weights; without them the row is a + # companion-only prefetch and would be marked partial (see the dedicated test). + _file("transformer/diffusion_pytorch_model.safetensors", 9_000), + ], tmp_path / "models--Tongyi-MAI--Z-Image-Turbo", ) checkpoint = _repo( @@ -468,6 +474,43 @@ def test_list_cached_models_tags_diffusers_pipeline_as_text_to_image(monkeypatch } +def test_list_cached_models_marks_companion_only_pipeline_partial(monkeypatch, tmp_path): + """A GGUF image load prefetches its companion base repo's VAE / text-encoder / model_index.json + but deliberately skips the multi-GB transformer (the GGUF supplies it). That snapshot carries a + root model_index.json yet is not a loadable BF16 pipeline, so it must be marked partial (the + picker drops partial rows) rather than advertised as fully on-device. A sibling repo that DOES + ship its transformer shards stays complete.""" + companion_only = _repo( + "black-forest-labs/FLUX.1-dev", + [ + _file("model_index.json", 1_000), + _file("vae/diffusion_pytorch_model.safetensors", 9_000), + _file("text_encoder/model.safetensors", 9_000), + ], + tmp_path / "models--black-forest-labs--FLUX.1-dev", + ) + complete = _repo( + "Tongyi-MAI/Z-Image-Turbo", + [ + _file("model_index.json", 1_000), + _file("text_encoder/model.safetensors", 9_000), + _file("transformer/diffusion_pytorch_model.safetensors", 9_000), + ], + tmp_path / "models--Tongyi-MAI--Z-Image-Turbo", + ) + + monkeypatch.setattr( + models_route, + "_all_hf_cache_scans", + lambda: [SimpleNamespace(repos = [companion_only, complete])], + ) + + result = asyncio.run(models_route.list_cached_models(current_subject = "test-user")) + by_repo = {c["repo_id"]: c for c in result["cached"]} + assert by_repo["black-forest-labs/FLUX.1-dev"].get("partial") is True + assert by_repo["Tongyi-MAI/Z-Image-Turbo"].get("partial") is None + + def test_list_cached_gguf_includes_vision_repo_with_main_gguf_and_mmproj(monkeypatch, tmp_path): """A vision GGUF repo (main weight + mmproj) is a GGUF repo; reported size is the main weight only, since mmproj is filtered at classification.""" diff --git a/studio/backend/tests/test_local_model_format.py b/studio/backend/tests/test_local_model_format.py index 1fbd7efcd6..dcc948ab30 100644 --- a/studio/backend/tests/test_local_model_format.py +++ b/studio/backend/tests/test_local_model_format.py @@ -136,6 +136,23 @@ def test_scan_models_dir_surfaces_diffusers_pipeline_folder(tmp_path): assert rows["my-pipeline"].model_format is None +def test_scan_models_dir_surfaces_root_diffusers_pipeline(tmp_path): + # A custom scan folder can point DIRECTLY at a diffusers pipeline (not a parent of repos). + # Its weights live in component subdirs under a root model_index.json, so _is_model_directory + # rejects the root; without admitting it the scan would surface the component subdirs + # (transformer/, vae/) as bogus models and hide the real pipeline. Treat the root as one model. + root = tmp_path / "my-local-pipeline" + _touch(root / "model_index.json") + _touch(root / "transformer" / "config.json") + _touch(root / "transformer" / "diffusion_pytorch_model.safetensors") + _touch(root / "vae" / "diffusion_pytorch_model.safetensors") + + rows = models_route._scan_models_dir(root) + + assert [r.path for r in rows] == [str(root)] + assert rows[0].model_format is None + + # ── Images picker task tag for local (non-GGUF) diffusers models ────────────── from models.models import LocalModelInfo # noqa: E402 @@ -211,6 +228,25 @@ def test_local_task_tags_video_single_file_checkpoint(tmp_path): ) +def test_local_task_tags_single_file_by_checkpoint_filename(tmp_path): + # A generically named folder holding one loadable checkpoint whose FILENAME identifies the + # family (the parent dir does not) is loadable -- the route resolves the sole file via + # resolve_local_single_file -- so tag it from the filename or the task-scoped picker hides it. + d = tmp_path / "downloads" + _touch(d / "qwen-image-2509.safetensors") # family only in the filename, no model_index.json + m = _local(d, id = str(d), display_name = "downloads") + assert models_route._local_is_diffusers(m) is True + assert models_route._local_model_task(m) == "text-to-image" + + +def test_local_task_tags_video_single_file_by_checkpoint_filename(tmp_path): + # Same, for a video family whose token lives only in the sole checkpoint's filename. + d = tmp_path / "clips" + _touch(d / "ltx-2.3-distilled.safetensors") # ltx family only in the filename + m = _local(d, id = str(d), display_name = "clips") + assert models_route._local_model_task(m) == models_route._VIDEO_GEN_TASK + + def test_local_task_ignores_family_token_in_parent_path(tmp_path): # model.id is the full on-disk path for a scanned On-Device model, and the family-token # matcher treats any path segment as a hint. A family token in a PARENT dir (e.g.