From 115a50cc3096eb1eb6538ebcd93d7152dfe783f1 Mon Sep 17 00:00:00 2001 From: Unsloth Date: Sun, 26 Jul 2026 04:45:28 -0700 Subject: [PATCH] Route a chat pick of a diffusion model to the Images or Video page Chat cannot load one, so it was either hidden or failed on load. The unfiltered picker now lists on-device diffusion models and navigates to the page that runs them, passing the repo and quant so that page loads it. --- studio/frontend/src/app/routes/images.tsx | 8 +++ studio/frontend/src/app/routes/video.tsx | 8 +++ .../src/features/images/images-page.tsx | 23 ++++++++ .../components/model-selector/pickers.tsx | 55 ++++++++++++++++--- .../src/features/video/video-page.tsx | 23 ++++++++ tests/studio/test_model_picker_contracts.py | 16 ++++++ 6 files changed, 126 insertions(+), 7 deletions(-) diff --git a/studio/frontend/src/app/routes/images.tsx b/studio/frontend/src/app/routes/images.tsx index b177ddc4ce..96f17bf606 100644 --- a/studio/frontend/src/app/routes/images.tsx +++ b/studio/frontend/src/app/routes/images.tsx @@ -11,6 +11,14 @@ export const Route = createRoute({ getParentRoute: () => rootRoute, path: "/images", staticData: { title: "Images" }, + // A diffusion pick made from the chat picker arrives here as ?model= (+ ?quant=), + // which the page loads and then clears. + validateSearch: ( + search: Record, + ): { model?: string; quant?: string } => ({ + ...(typeof search.model === "string" ? { model: search.model } : {}), + ...(typeof search.quant === "string" ? { quant: search.quant } : {}), + }), beforeLoad: () => requireAuth(), component: () => null, }); diff --git a/studio/frontend/src/app/routes/video.tsx b/studio/frontend/src/app/routes/video.tsx index b1979ab5f8..fdeebd23f2 100644 --- a/studio/frontend/src/app/routes/video.tsx +++ b/studio/frontend/src/app/routes/video.tsx @@ -11,6 +11,14 @@ export const Route = createRoute({ getParentRoute: () => rootRoute, path: "/video", staticData: { title: "Video" }, + // A diffusion pick made from the chat picker arrives here as ?model= (+ ?quant=), + // which the page loads and then clears. + validateSearch: ( + search: Record, + ): { model?: string; quant?: string } => ({ + ...(typeof search.model === "string" ? { model: search.model } : {}), + ...(typeof search.quant === "string" ? { quant: search.quant } : {}), + }), beforeLoad: () => requireAuth(), component: () => null, }); diff --git a/studio/frontend/src/features/images/images-page.tsx b/studio/frontend/src/features/images/images-page.tsx index 7320b1984e..a923315b63 100644 --- a/studio/frontend/src/features/images/images-page.tsx +++ b/studio/frontend/src/features/images/images-page.tsx @@ -98,6 +98,7 @@ import { loadDiffusionModel, unloadDiffusionModel, } from "./api"; +import { useNavigate, useSearch } from "@tanstack/react-router"; import { useStagedDownload } from "@/features/hub/download-manager"; import { DiffusionTrainPanel } from "./train/diffusion-train-panel"; import { @@ -1811,6 +1812,28 @@ export function ImagesPage({ active = true }: { active?: boolean }) { [stage], ); + // A diffusion model picked from the chat picker arrives as ?model= on this route. Load it + // once, then clear the params so a refresh or a later manual eject does not reload it. + const routeSearch = useSearch({ strict: false }) as { + model?: string; + quant?: string; + }; + const navigateSelf = useNavigate(); + const handledRouteModel = useRef(null); + useEffect(() => { + const wanted = routeSearch.model; + if (!wanted || handledRouteModel.current === wanted) return; + handledRouteModel.current = wanted; + void navigateSelf({ to: "/images", search: {}, replace: true }); + void loadOrStage( + wanted, + routeSearch.quant + ? { kind: "gguf", filename: routeSearch.quant } + : { kind: "pipeline" }, + false, + ); + }, [routeSearch.model, routeSearch.quant, loadOrStage, navigateSelf]); + // Reload the current model with the current advanced options. const handleReapply = useCallback(() => { const l = lastLoad.current; diff --git a/studio/frontend/src/features/model-picker/components/model-selector/pickers.tsx b/studio/frontend/src/features/model-picker/components/model-selector/pickers.tsx index bd342a7ea1..a2045a19da 100644 --- a/studio/frontend/src/features/model-picker/components/model-selector/pickers.tsx +++ b/studio/frontend/src/features/model-picker/components/model-selector/pickers.tsx @@ -1,6 +1,7 @@ // SPDX-License-Identifier: AGPL-3.0-only // Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 +import { useNavigate } from "@tanstack/react-router"; import { Checkbox } from "@/components/ui/checkbox"; import { Input } from "@/components/ui/input"; import { Spinner } from "@/components/ui/spinner"; @@ -1215,15 +1216,19 @@ export const VIDEO_GEN_TASKS = ["text-to-video"] as const; // them out (they'd 400 on load). const UNSUPPORTED_DIFFUSION_TASK = "image-diffusion-unsupported"; -// Tasks that must never appear as a loadable chat model: the Images- and Video-handled -// generation tasks plus the non-loadable diffusion tag above. Keeping text-to-video here -// stops a downloaded video GGUF from showing up as a loadable chat model (it would 400). -const NON_CHAT_TASKS: readonly string[] = [ +// Generation tasks the Images / Video pages own. They are not chat-loadable, so an +// on-device pick of one routes to its page instead of loading into chat. +const DIFFUSION_PAGE_TASKS: readonly string[] = [ ...IMAGE_GEN_TASKS, ...VIDEO_GEN_TASKS, - UNSUPPORTED_DIFFUSION_TASK, ]; +/** The page that runs this task, or null when chat should handle the pick. */ +function diffusionPageForTask(task: string | null | undefined): "images" | "video" | null { + if (!task || !DIFFUSION_PAGE_TASKS.includes(task)) return null; + return (VIDEO_GEN_TASKS as readonly string[]).includes(task) ? "video" : "images"; +} + // Editing/inpaint checkpoints are tagged image-to-image but need an input image the // text-to-image backend rejects (mirrors its _EDIT_KEYWORDS). Hidden by id so they don't show // in the Images picker only to 400 on load. Keeping the image-to-image task itself is required: @@ -1258,7 +1263,9 @@ function passesTaskGate( ): boolean { if (filter) return taskMatchesFilter(repoTask, filter) && !isImageEditModel(repoId); - return !(repoTask != null && NON_CHAT_TASKS.includes(repoTask)); + // Unfiltered (chat) picker: an on-device diffusion model stays listed and routes to the + // Images/Video page on click. Only the never-loadable diffusion tag is hidden outright. + return repoTask !== UNSUPPORTED_DIFFUSION_TASK; } // Module-level caches so re-mounting the popover shows results instantly @@ -1463,7 +1470,7 @@ export function HubModelPicker({ loraModels = [], externalModels = [], value, - onSelect, + onSelect: onSelectProp, onFoldersChange, onBrowseHub, onModelsChange, @@ -2197,6 +2204,40 @@ export function HubModelPicker({ [customFolderModels, customSort, formatFilter, loadTimes, localQuery, task], ); + // Chat cannot load a diffusion model, but the Images/Video pages can: rather than hiding + // an on-device one or letting it 400, a pick routes to the page that runs it, which then + // loads it. Task-scoped pickers (already on those pages) select normally. + const navigateToPage = useNavigate(); + const diffusionTaskById = useMemo(() => { + const byId = new Map(); + const put = (id: string | null | undefined, t: string | null | undefined) => { + if (id && t) byId.set(id.toLowerCase(), t); + }; + for (const c of cachedGguf) put(c.repo_id, c.task); + for (const c of cachedModels) put(c.repo_id, c.task); + for (const m of lmStudioModels) put(m.model_id ?? m.id, m.task); + for (const m of localDirModels) put(m.model_id ?? m.id, m.task); + for (const m of customFolderModels) put(m.model_id ?? m.id, m.task); + return byId; + }, [cachedGguf, cachedModels, lmStudioModels, localDirModels, customFolderModels]); + + const onSelect = useCallback( + (id: string, meta: ModelSelectorChangeMeta) => { + if (!task) { + const page = diffusionPageForTask(diffusionTaskById.get(id.toLowerCase())); + if (page) { + void navigateToPage({ + to: `/${page}`, + search: { model: id, quant: meta.ggufVariant ?? undefined }, + }); + return; + } + } + onSelectProp(id, meta); + }, + [task, diffusionTaskById, navigateToPage, onSelectProp], + ); + // Fine-tuned models for the On Device "Fine-tuned" section: flat, query- // filtered, newest first. const fineTunedRows = useMemo(() => { diff --git a/studio/frontend/src/features/video/video-page.tsx b/studio/frontend/src/features/video/video-page.tsx index 98d1d34ead..c85013599c 100644 --- a/studio/frontend/src/features/video/video-page.tsx +++ b/studio/frontend/src/features/video/video-page.tsx @@ -56,6 +56,7 @@ import { ParamSlider } from "@/features/chat"; import { ModelLoadDescription } from "@/features/chat/components/model-load-status"; import { getHfToken, hfApiToken } from "@/features/hub/stores/hf-token-store"; import { formatBytes, formatEta } from "@/features/hub/lib/format"; +import { useNavigate, useSearch } from "@tanstack/react-router"; import { useStagedDownload } from "@/features/hub/download-manager"; import { cn } from "@/lib/utils"; import { toast } from "@/lib/toast"; @@ -1139,6 +1140,28 @@ export function VideoPage({ active = true }: { active?: boolean }) { [stage], ); + // A diffusion model picked from the chat picker arrives as ?model= on this route. Load it + // once, then clear the params so a refresh or a later manual eject does not reload it. + const routeSearch = useSearch({ strict: false }) as { + model?: string; + quant?: string; + }; + const navigateSelf = useNavigate(); + const handledRouteModel = useRef(null); + useEffect(() => { + const wanted = routeSearch.model; + if (!wanted || handledRouteModel.current === wanted) return; + handledRouteModel.current = wanted; + void navigateSelf({ to: "/video", search: {}, replace: true }); + void loadOrStage( + wanted, + routeSearch.quant + ? { kind: "gguf", filename: routeSearch.quant } + : { kind: "pipeline" }, + false, + ); + }, [routeSearch.model, routeSearch.quant, loadOrStage, navigateSelf]); + // Reload the current model with the current advanced options. const handleReapply = useCallback(() => { const l = lastLoad.current; diff --git a/tests/studio/test_model_picker_contracts.py b/tests/studio/test_model_picker_contracts.py index 1a4bf151bd..9285b72ff7 100644 --- a/tests/studio/test_model_picker_contracts.py +++ b/tests/studio/test_model_picker_contracts.py @@ -651,3 +651,19 @@ def test_local_model_sections_respect_the_task_filter(): assert "passesTaskGate(m.task" in block.group(0), ( f"{memo} does not apply the task gate" ) + + +def test_chat_picker_routes_diffusion_picks_to_their_page(): + """Chat cannot load a diffusion model. Rather than hiding an on-device one or letting + it 400, the unfiltered picker routes the pick to the Images/Video page, which loads it.""" + src = _read("features/model-picker/components/model-selector/pickers.tsx") + gate = re.search(r"function passesTaskGate\(.*?\n\}", src, re.S) + assert gate, "passesTaskGate not found" + # The chat branch no longer drops the generation tasks outright. + assert "UNSUPPORTED_DIFFUSION_TASK" in gate.group(0) + wrapper = re.search(r"const onSelect = useCallback\(.*?\n \);", src, re.S) + assert wrapper, "the routing wrapper around onSelect is missing" + body = wrapper.group(0) + assert "diffusionPageForTask" in body and "navigateToPage" in body + # Task-scoped pickers (already on those pages) must select normally. + assert "if (!task)" in body