|
| 1 | +import { useEffect, useState, useCallback, useMemo } from 'react'; |
| 2 | +import { useSseProgress } from './useSseProgress.js'; |
| 3 | +import toast from '../components/ui/Toast'; |
| 4 | +import { getImageModelStatuses, getVideoModelStatuses } from '../services/apiImageVideo.js'; |
| 5 | + |
| 6 | +// Sentinel `modelId` used to drive a text-encoder download instead of a model |
| 7 | +// download. The video form passes this to `start()`; the URL builder below |
| 8 | +// rewrites to the dedicated /text-encoder/download endpoint. Exported so |
| 9 | +// callers and the hook agree on the magic string. |
| 10 | +export const TEXT_ENCODER_DOWNLOAD_ID = '__text_encoder__'; |
| 11 | + |
| 12 | +const buildDownloadUrl = (kind, modelId) => { |
| 13 | + if (!modelId) return null; |
| 14 | + if (kind === 'video' && modelId === TEXT_ENCODER_DOWNLOAD_ID) { |
| 15 | + return '/api/video-gen/text-encoder/download'; |
| 16 | + } |
| 17 | + return `/api/${kind}-gen/models/${encodeURIComponent(modelId)}/download`; |
| 18 | +}; |
| 19 | + |
| 20 | +// Model download-status hook. Drives the inline "Available · 7.8 GB" / |
| 21 | +// "Download (~8 GB)" badge next to the image/video gen model picker. |
| 22 | +// |
| 23 | +// `kind` selects the endpoint family ('image' | 'video'). `start(modelId)` |
| 24 | +// opens an EventSource against the download endpoint. When the stream |
| 25 | +// emits a terminal frame we automatically refetch /models/status so the |
| 26 | +// badge flips to "Available" without the caller wiring that up. |
| 27 | +export function useModelDownloadStatus({ kind = 'image' } = {}) { |
| 28 | + const [statuses, setStatuses] = useState(null); |
| 29 | + const [extra, setExtra] = useState({}); // video: { textEncoder: {...} } |
| 30 | + const [loading, setLoading] = useState(false); |
| 31 | + const [activeModelId, setActiveModelId] = useState(null); |
| 32 | + |
| 33 | + const fetchStatuses = useCallback(async () => { |
| 34 | + setLoading(true); |
| 35 | + // Best-effort: a failure leaves the badge in its loading state. The form |
| 36 | + // still works because lazy download is the existing fallback. |
| 37 | + const body = await (kind === 'video' ? getVideoModelStatuses() : getImageModelStatuses()) |
| 38 | + .catch(() => null); |
| 39 | + if (body == null) { |
| 40 | + setStatuses([]); |
| 41 | + setExtra({}); |
| 42 | + } else if (kind === 'video') { |
| 43 | + setStatuses(Array.isArray(body?.models) ? body.models : []); |
| 44 | + setExtra({ textEncoder: body?.textEncoder || null }); |
| 45 | + } else { |
| 46 | + setStatuses(Array.isArray(body) ? body : []); |
| 47 | + setExtra({}); |
| 48 | + } |
| 49 | + setLoading(false); |
| 50 | + }, [kind]); |
| 51 | + |
| 52 | + useEffect(() => { fetchStatuses(); }, [fetchStatuses]); |
| 53 | + |
| 54 | + // EventSource for the active download. `null` URL = idle (useSseProgress's |
| 55 | + // `enabled: false` cleanup tears the connection down on cancel). |
| 56 | + const downloadUrl = buildDownloadUrl(kind, activeModelId); |
| 57 | + const sse = useSseProgress(downloadUrl, { enabled: !!downloadUrl }); |
| 58 | + |
| 59 | + // Refetch on natural stream close. useSseProgress flips `closed:true` once |
| 60 | + // per subscription and resets to false when the URL changes; that single |
| 61 | + // transition is the safe signal — no extra `wasClosed` ref needed. |
| 62 | + // A terminal error frame (gated repo, missing HF token, broken venv) is |
| 63 | + // routed to a toast here because the active-badge state vanishes the moment |
| 64 | + // we clear `activeModelId`; without this, the UI silently snaps back to the |
| 65 | + // Download button and the actionable server message is lost. |
| 66 | + useEffect(() => { |
| 67 | + if (sse.closed) { |
| 68 | + if (sse.latest?.type === 'error' && sse.latest?.message) { |
| 69 | + toast.error(sse.latest.message); |
| 70 | + } |
| 71 | + fetchStatuses(); |
| 72 | + setActiveModelId(null); |
| 73 | + } |
| 74 | + }, [sse.closed, sse.latest, fetchStatuses]); |
| 75 | + |
| 76 | + const start = useCallback((modelId) => { |
| 77 | + setActiveModelId(modelId); |
| 78 | + }, []); |
| 79 | + |
| 80 | + // Manual cancel: refetch directly because `sse.close()` followed by |
| 81 | + // setActiveModelId(null) clears the URL, which causes useSseProgress to |
| 82 | + // reset `closed → false` before the close-effect can observe `true`. |
| 83 | + // Without this direct refetch the badge would stay stuck on its pre-cancel |
| 84 | + // state until the next page mount. |
| 85 | + const cancel = useCallback(() => { |
| 86 | + sse.close(); |
| 87 | + setActiveModelId(null); |
| 88 | + fetchStatuses(); |
| 89 | + }, [sse, fetchStatuses]); |
| 90 | + |
| 91 | + // Memoize the active model's enriched status so a new SSE frame doesn't |
| 92 | + // hand a fresh object to every non-active model's badge (only the active |
| 93 | + // one re-renders). For inactive models we return the raw entry, which is |
| 94 | + // referentially stable across frames. |
| 95 | + const activeStatus = useMemo(() => { |
| 96 | + if (!activeModelId) return null; |
| 97 | + const list = Array.isArray(statuses) ? statuses : []; |
| 98 | + const entry = list.find((s) => s.id === activeModelId); |
| 99 | + if (!entry) return null; |
| 100 | + return { ...entry, downloading: true, progress: sse.latest }; |
| 101 | + }, [activeModelId, statuses, sse.latest]); |
| 102 | + |
| 103 | + const getStatus = useCallback((modelId) => { |
| 104 | + if (modelId === activeModelId) return activeStatus; |
| 105 | + const list = Array.isArray(statuses) ? statuses : []; |
| 106 | + return list.find((s) => s.id === modelId) || null; |
| 107 | + }, [statuses, activeModelId, activeStatus]); |
| 108 | + |
| 109 | + return { |
| 110 | + statuses, |
| 111 | + extra, |
| 112 | + loading, |
| 113 | + refresh: fetchStatuses, |
| 114 | + start, |
| 115 | + cancel, |
| 116 | + getStatus, |
| 117 | + activeModelId, |
| 118 | + progress: sse.latest, |
| 119 | + downloading: !!activeModelId, |
| 120 | + }; |
| 121 | +} |
0 commit comments