diff --git a/studio/backend/models/models.py b/studio/backend/models/models.py index 20dea5ec12..1c7ce86c2c 100644 --- a/studio/backend/models/models.py +++ b/studio/backend/models/models.py @@ -183,6 +183,12 @@ class LocalModelInfo(BaseModel): None, description = "Unix timestamp of latest observed update", ) + task: Optional[str] = Field( + None, + description = "HF pipeline task inferred from a GGUF's architecture " + "('text-to-image' for diffusion, 'text-generation' otherwise). Lets the " + "Images picker show only diffusion GGUFs.", + ) class LocalModelListResponse(BaseModel): diff --git a/studio/backend/routes/models.py b/studio/backend/routes/models.py index 1aa13108e4..92b8824016 100644 --- a/studio/backend/routes/models.py +++ b/studio/backend/routes/models.py @@ -843,6 +843,11 @@ async def list_local_models( reverse = True, ) models = [m for m in models if not _is_hidden_model(m.id, m.path)] + # Tag each GGUF with its task so the Images picker can filter to diffusion. + models = [ + m.model_copy(update = {"task": _local_model_task(m.path, m.model_format)}) + for m in models + ] return LocalModelListResponse( models_dir = str(models_root), @@ -3117,6 +3122,12 @@ def _gguf_architecture(path: str) -> Optional[str]: return None +def _arch_to_task(arch: Optional[str]) -> Optional[str]: + if arch is None: + return None + return "text-to-image" if arch.lower() in _DIFFUSION_GGUF_ARCHS else "text-generation" + + def _repo_gguf_task(repo_info) -> Optional[str]: """HF pipeline task of a cached GGUF repo, from its architecture: 'text-to-image' for diffusion archs, else 'text-generation' (None if unreadable).""" @@ -3124,10 +3135,29 @@ def _repo_gguf_task(repo_info) -> Optional[str]: for path in _iter_gguf_paths(Path(repo_info.repo_path)): if _is_mmproj_filename(path.name): continue - arch = _gguf_architecture(str(path)) - if arch is None: + task = _arch_to_task(_gguf_architecture(str(path))) + if task is not None: + return task + except Exception: + pass + return None + + +def _local_model_task(path: str, model_format: Optional[str]) -> Optional[str]: + """Same classification for a local model: read its GGUF architecture. The + path may be the .gguf file itself or a folder containing one.""" + if model_format != "gguf": + return None + try: + p = Path(path) + if p.suffix.lower() == ".gguf" and p.is_file(): + return _arch_to_task(_gguf_architecture(str(p))) + for f in _iter_gguf_paths(p): + if _is_mmproj_filename(f.name): continue - return "text-to-image" if arch.lower() in _DIFFUSION_GGUF_ARCHS else "text-generation" + task = _arch_to_task(_gguf_architecture(str(f))) + if task is not None: + return task except Exception: pass return None diff --git a/studio/frontend/src/components/assistant-ui/model-selector/pickers.tsx b/studio/frontend/src/components/assistant-ui/model-selector/pickers.tsx index f8c7c176dd..d995887ca9 100644 --- a/studio/frontend/src/components/assistant-ui/model-selector/pickers.tsx +++ b/studio/frontend/src/components/assistant-ui/model-selector/pickers.tsx @@ -1745,20 +1745,20 @@ export function HubModelPicker({ normalizeForSearch( `${m.model_id ?? ""} ${m.display_name} ${m.id}`, ).includes(localQuery); - // A task filter (the Images page) wants diffusion GGUFs only; local folders / - // LM Studio / fine-tuned models aren't that, so drop them in that mode. + // A task filter (the Images page) wants diffusion GGUFs only, so local models + // are filtered to that task (by the GGUF architecture the backend reports). const sortedLmStudio = useMemo( () => - task - ? [] - : sortLocalModels( - lmStudioModels.filter( - (m) => - localModelMatchesFormat(m, formatFilter) && matchesLocalQuery(m), - ), - downloadedSort, - loadTimes, - ), + sortLocalModels( + lmStudioModels.filter( + (m) => + (!task || taskMatchesFilter(m.task, task)) && + localModelMatchesFormat(m, formatFilter) && + matchesLocalQuery(m), + ), + downloadedSort, + loadTimes, + ), // eslint-disable-next-line react-hooks/exhaustive-deps [lmStudioModels, downloadedSort, formatFilter, loadTimes, localQuery, task], ); @@ -1767,20 +1767,19 @@ export function HubModelPicker({ // rule). An MLX build a Mac user dropped in ./models stays selectable. const sortedLocalDir = useMemo( () => - task - ? [] - : sortLocalModels( - localDirModels.filter( - (m) => - (!chatOnly || - localModelIsGguf(m) || - (isMac && localModelIsMlx(m))) && - localModelMatchesFormat(m, formatFilter) && - matchesLocalQuery(m), - ), - downloadedSort, - loadTimes, - ), + sortLocalModels( + localDirModels.filter( + (m) => + (!task || taskMatchesFilter(m.task, task)) && + (!chatOnly || + localModelIsGguf(m) || + (isMac && localModelIsMlx(m))) && + localModelMatchesFormat(m, formatFilter) && + matchesLocalQuery(m), + ), + downloadedSort, + loadTimes, + ), // eslint-disable-next-line react-hooks/exhaustive-deps [ localDirModels, @@ -1795,16 +1794,16 @@ export function HubModelPicker({ ); const sortedCustomFolderModels = useMemo( () => - task - ? [] - : sortLocalModels( - customFolderModels.filter( - (m) => - localModelMatchesFormat(m, formatFilter) && matchesLocalQuery(m), - ), - customSort, - loadTimes, - ), + sortLocalModels( + customFolderModels.filter( + (m) => + (!task || taskMatchesFilter(m.task, task)) && + localModelMatchesFormat(m, formatFilter) && + matchesLocalQuery(m), + ), + customSort, + loadTimes, + ), // eslint-disable-next-line react-hooks/exhaustive-deps [customFolderModels, customSort, formatFilter, loadTimes, localQuery, task], ); @@ -2208,9 +2207,7 @@ export function HubModelPicker({ // On Device owns the downloaded and custom-folder models; the Unsloth tab // searches the HF listing (below). Both filter locally by the query. const showDownloaded = section === "downloaded"; - // Custom Folders (server-side scan dirs) hold arbitrary local models, not - // diffusion GGUFs — hide the whole section under a task filter (Images). - const showCustom = section === "downloaded" && !task; + const showCustom = section === "downloaded"; const showRecommendedSection = !showHfSection && section === "recommended"; const downloadedEmpty = visibleCachedGguf.length === 0 && @@ -2611,6 +2608,7 @@ export function HubModelPicker({ ) : null} + {!task && (