From ee6d33fa324893aaa40c117e817f5044639c3235 Mon Sep 17 00:00:00 2001 From: imagineer99 Date: Thu, 19 Feb 2026 08:13:16 +0000 Subject: [PATCH 01/23] feat: add download progress indicators for dataset preview and training overlay --- .../features/studio/sections/dataset-preview-dialog.tsx | 8 +++++++- .../src/features/studio/sections/progress-section-lib.ts | 6 ++++++ studio/frontend/src/features/studio/training-view.tsx | 2 ++ studio/frontend/src/features/training/types/runtime.ts | 2 ++ 4 files changed, 17 insertions(+), 1 deletion(-) diff --git a/studio/frontend/src/features/studio/sections/dataset-preview-dialog.tsx b/studio/frontend/src/features/studio/sections/dataset-preview-dialog.tsx index 648adfbc89..e135f301af 100644 --- a/studio/frontend/src/features/studio/sections/dataset-preview-dialog.tsx +++ b/studio/frontend/src/features/studio/sections/dataset-preview-dialog.tsx @@ -63,6 +63,7 @@ export function DatasetPreviewDialog({ const mappingOk = !!manualMapping.input && !!manualMapping.output; const leftLabel = isVlm ? "Image" : "Input"; const rightLabel = isVlm ? "Text" : "Output"; + const isHfDataset = !!datasetName && datasetName.includes("/"); useEffect(() => { if (!manualMapping.input || !manualMapping.output) return; @@ -266,8 +267,13 @@ export function DatasetPreviewDialog({

- Loading preview... + {isHfDataset ? "Fetching dataset preview from Hugging Face..." : "Loading preview..."}

+ {isHfDataset && ( +

+ This may take a moment for large datasets +

+ )} )} diff --git a/studio/frontend/src/features/studio/sections/progress-section-lib.ts b/studio/frontend/src/features/studio/sections/progress-section-lib.ts index d22212f70b..bff5daed81 100644 --- a/studio/frontend/src/features/studio/sections/progress-section-lib.ts +++ b/studio/frontend/src/features/studio/sections/progress-section-lib.ts @@ -2,6 +2,8 @@ import type { TrainingPhase } from "@/features/training"; export const phaseLabel: Record = { idle: "Idle", + downloading_model: "Downloading model", + downloading_dataset: "Downloading dataset", loading_model: "Loading model", loading_dataset: "Loading dataset", configuring: "Configuring", @@ -13,6 +15,10 @@ export const phaseLabel: Record = { export const phaseColors: Record = { idle: "bg-muted text-muted-foreground", + downloading_model: + "bg-sky-100 text-sky-700 dark:bg-sky-900 dark:text-sky-300", + downloading_dataset: + "bg-sky-100 text-sky-700 dark:bg-sky-900 dark:text-sky-300", loading_model: "bg-amber-100 text-amber-700 dark:bg-amber-900 dark:text-amber-300", loading_dataset: diff --git a/studio/frontend/src/features/studio/training-view.tsx b/studio/frontend/src/features/studio/training-view.tsx index 9e83a2f662..fe616ea88a 100644 --- a/studio/frontend/src/features/studio/training-view.tsx +++ b/studio/frontend/src/features/studio/training-view.tsx @@ -19,6 +19,8 @@ export function TrainingView(): ReactElement { ); const isPreparingPhase = + runtime.phase === "downloading_model" || + runtime.phase === "downloading_dataset" || runtime.phase === "loading_model" || runtime.phase === "loading_dataset" || runtime.phase === "configuring"; diff --git a/studio/frontend/src/features/training/types/runtime.ts b/studio/frontend/src/features/training/types/runtime.ts index df24c020ca..7ebf09518d 100644 --- a/studio/frontend/src/features/training/types/runtime.ts +++ b/studio/frontend/src/features/training/types/runtime.ts @@ -1,5 +1,7 @@ export type TrainingPhase = | "idle" + | "downloading_model" + | "downloading_dataset" | "loading_model" | "loading_dataset" | "configuring" From 93b31f0db2bfcdc2e73bb458d8c1c65035cead6b Mon Sep 17 00:00:00 2001 From: samit Date: Thu, 19 Feb 2026 15:03:23 -0800 Subject: [PATCH 02/23] added optim in the frontend --- studio/frontend/src/config/training.ts | 10 +++++ .../studio/sections/params-section.tsx | 43 ++++++++++++++++++- .../studio/sections/progress-section.tsx | 7 +++ .../src/features/training/api/mappers.ts | 2 +- .../training/stores/training-config-store.ts | 6 ++- .../src/features/training/types/config.ts | 2 + 6 files changed, 67 insertions(+), 3 deletions(-) diff --git a/studio/frontend/src/config/training.ts b/studio/frontend/src/config/training.ts index 8840b9363c..9ad09ba813 100644 --- a/studio/frontend/src/config/training.ts +++ b/studio/frontend/src/config/training.ts @@ -73,10 +73,20 @@ export const TARGET_MODULES = [ "down_proj", ]; +export const OPTIMIZER_OPTIONS: ReadonlyArray<{ value: string; label: string }> = [ + { value: "adamw_8bit", label: "AdamW 8-bit" }, + { value: "paged_adamw_8bit", label: "Paged AdamW 8-bit" }, + { value: "adamw_bnb_8bit", label: "AdamW BNB 8-bit" }, + { value: "paged_adamw_32bit", label: "Paged AdamW 32-bit" }, + { value: "adamw_torch", label: "AdamW (PyTorch)" }, + { value: "adamw_torch_fused", label: "AdamW (PyTorch Fused)" }, +]; + export const DEFAULT_HYPERPARAMS = { epochs: 3, contextLength: 2048, learningRate: 2e-4, + optimizerType: "adamw_8bit", loraRank: 16, loraAlpha: 32, loraDropout: 0.05, diff --git a/studio/frontend/src/features/studio/sections/params-section.tsx b/studio/frontend/src/features/studio/sections/params-section.tsx index 2069da7a90..cf25d5ba54 100644 --- a/studio/frontend/src/features/studio/sections/params-section.tsx +++ b/studio/frontend/src/features/studio/sections/params-section.tsx @@ -20,7 +20,11 @@ import { TooltipContent, TooltipTrigger, } from "@/components/ui/tooltip"; -import { CONTEXT_LENGTHS, TARGET_MODULES } from "@/config/training"; +import { + CONTEXT_LENGTHS, + OPTIMIZER_OPTIONS, + TARGET_MODULES, +} from "@/config/training"; import { useTrainingConfigStore } from "@/features/training"; import type { GradientCheckpointing } from "@/types/training"; import { @@ -508,6 +512,43 @@ export function ParamsSection(): ReactElement { value="optimization" className="mt-3 flex flex-col gap-3" > + + Optimization algorithm. 8-bit variants reduce memory usage. + Fused is recommended for vision models.{" "} + + Read more + + + } + > + + o.value === config.optimizerType)?.label ?? + config.optimizerType; + const configItems = [ { section: "Hyperparams", @@ -133,6 +139,7 @@ export function ProgressSection(): ReactElement { ["Epochs", config.epochs], ["Batch size", config.batchSize], ["Learning rate", config.learningRate], + ["Optimizer", optimizerLabel], ["Max steps", config.maxSteps], ["Context length", config.contextLength], ["Warmup steps", config.warmupSteps], diff --git a/studio/frontend/src/features/training/api/mappers.ts b/studio/frontend/src/features/training/api/mappers.ts index 510cc542b8..78386f197b 100644 --- a/studio/frontend/src/features/training/api/mappers.ts +++ b/studio/frontend/src/features/training/api/mappers.ts @@ -40,7 +40,7 @@ export function buildTrainingStartPayload( weight_decay: config.weightDecay, random_seed: config.randomSeed, packing: config.packing, - optim: "adamw_8bit", + optim: config.optimizerType, lr_scheduler_type: "linear", use_lora: adapterMethod, lora_r: config.loraRank, diff --git a/studio/frontend/src/features/training/stores/training-config-store.ts b/studio/frontend/src/features/training/stores/training-config-store.ts index 67317dbc09..62d4808b2c 100644 --- a/studio/frontend/src/features/training/stores/training-config-store.ts +++ b/studio/frontend/src/features/training/stores/training-config-store.ts @@ -305,6 +305,7 @@ export const useTrainingConfigStore = create()( setEpochs: (epochs) => set({ epochs }), setContextLength: (contextLength) => set({ contextLength }), setLearningRate: (learningRate) => set({ learningRate }), + setOptimizerType: (optimizerType) => set({ optimizerType }), setLoraRank: (loraRank) => set({ loraRank }), setLoraAlpha: (loraAlpha) => set({ loraAlpha }), setLoraDropout: (loraDropout) => set({ loraDropout }), @@ -346,7 +347,7 @@ export const useTrainingConfigStore = create()( }, { name: "unsloth_training_config_v1", - version: 3, + version: 4, migrate: (persisted, version) => { const s = persisted as Record; if (version < 2 && s.datasetSubset == null && s.datasetConfig != null) { @@ -356,6 +357,9 @@ export const useTrainingConfigStore = create()( if (version < 3 && s.modelDefaultsAppliedFor == null) { s.modelDefaultsAppliedFor = null; } + if (version < 4 && s.optimizerType == null) { + s.optimizerType = DEFAULT_HYPERPARAMS.optimizerType; + } return s as unknown as TrainingConfigStore; }, partialize: partializePersistedState, diff --git a/studio/frontend/src/features/training/types/config.ts b/studio/frontend/src/features/training/types/config.ts index 53beabd383..8557739b06 100644 --- a/studio/frontend/src/features/training/types/config.ts +++ b/studio/frontend/src/features/training/types/config.ts @@ -30,6 +30,7 @@ export interface TrainingConfigState { epochs: number; contextLength: number; learningRate: number; + optimizerType: string; loraRank: number; loraAlpha: number; loraDropout: number; @@ -85,6 +86,7 @@ export interface TrainingConfigActions { setEpochs: (epochs: number) => void; setContextLength: (length: number) => void; setLearningRate: (rate: number) => void; + setOptimizerType: (value: string) => void; setLoraRank: (rank: number) => void; setLoraAlpha: (alpha: number) => void; setLoraDropout: (dropout: number) => void; From 68028bf7f36edbcfdff6b7c3c31b60175420b3a3 Mon Sep 17 00:00:00 2001 From: samit Date: Thu, 19 Feb 2026 23:31:46 -0800 Subject: [PATCH 03/23] added lr_scheduler type to the frontend --- studio/frontend/src/config/training.ts | 6 +++ .../studio/sections/params-section.tsx | 38 +++++++++++++++++++ .../src/features/training/api/mappers.ts | 2 +- .../src/features/training/api/models-api.ts | 2 + .../features/training/lib/model-defaults.ts | 8 ++++ .../training/stores/training-config-store.ts | 6 ++- .../src/features/training/types/config.ts | 2 + 7 files changed, 62 insertions(+), 2 deletions(-) diff --git a/studio/frontend/src/config/training.ts b/studio/frontend/src/config/training.ts index 9ad09ba813..da60328d40 100644 --- a/studio/frontend/src/config/training.ts +++ b/studio/frontend/src/config/training.ts @@ -82,11 +82,17 @@ export const OPTIMIZER_OPTIONS: ReadonlyArray<{ value: string; label: string }> { value: "adamw_torch_fused", label: "AdamW (PyTorch Fused)" }, ]; +export const LR_SCHEDULER_OPTIONS: ReadonlyArray<{ value: string; label: string }> = [ + { value: "linear", label: "Linear" }, + { value: "cosine", label: "Cosine" }, +]; + export const DEFAULT_HYPERPARAMS = { epochs: 3, contextLength: 2048, learningRate: 2e-4, optimizerType: "adamw_8bit", + lrSchedulerType: "linear", loraRank: 16, loraAlpha: 32, loraDropout: 0.05, diff --git a/studio/frontend/src/features/studio/sections/params-section.tsx b/studio/frontend/src/features/studio/sections/params-section.tsx index cf25d5ba54..edc21e0e0e 100644 --- a/studio/frontend/src/features/studio/sections/params-section.tsx +++ b/studio/frontend/src/features/studio/sections/params-section.tsx @@ -22,6 +22,7 @@ import { } from "@/components/ui/tooltip"; import { CONTEXT_LENGTHS, + LR_SCHEDULER_OPTIONS, OPTIMIZER_OPTIONS, TARGET_MODULES, } from "@/config/training"; @@ -549,6 +550,43 @@ export function ParamsSection(): ReactElement { + + How the learning rate changes over training. Linear decays + steadily; cosine decays in a curve.{" "} + + Read more + + + } + > + + ()( setContextLength: (contextLength) => set({ contextLength }), setLearningRate: (learningRate) => set({ learningRate }), setOptimizerType: (optimizerType) => set({ optimizerType }), + setLrSchedulerType: (lrSchedulerType) => set({ lrSchedulerType }), setLoraRank: (loraRank) => set({ loraRank }), setLoraAlpha: (loraAlpha) => set({ loraAlpha }), setLoraDropout: (loraDropout) => set({ loraDropout }), @@ -347,7 +348,7 @@ export const useTrainingConfigStore = create()( }, { name: "unsloth_training_config_v1", - version: 4, + version: 5, migrate: (persisted, version) => { const s = persisted as Record; if (version < 2 && s.datasetSubset == null && s.datasetConfig != null) { @@ -360,6 +361,9 @@ export const useTrainingConfigStore = create()( if (version < 4 && s.optimizerType == null) { s.optimizerType = DEFAULT_HYPERPARAMS.optimizerType; } + if (version < 5 && s.lrSchedulerType == null) { + s.lrSchedulerType = DEFAULT_HYPERPARAMS.lrSchedulerType; + } return s as unknown as TrainingConfigStore; }, partialize: partializePersistedState, diff --git a/studio/frontend/src/features/training/types/config.ts b/studio/frontend/src/features/training/types/config.ts index 8557739b06..b268a08b79 100644 --- a/studio/frontend/src/features/training/types/config.ts +++ b/studio/frontend/src/features/training/types/config.ts @@ -31,6 +31,7 @@ export interface TrainingConfigState { contextLength: number; learningRate: number; optimizerType: string; + lrSchedulerType: string; loraRank: number; loraAlpha: number; loraDropout: number; @@ -87,6 +88,7 @@ export interface TrainingConfigActions { setContextLength: (length: number) => void; setLearningRate: (rate: number) => void; setOptimizerType: (value: string) => void; + setLrSchedulerType: (value: string) => void; setLoraRank: (rank: number) => void; setLoraAlpha: (alpha: number) => void; setLoraDropout: (dropout: number) => void; From a8c1f5fa8480ad0dc207c3421b85afa426b1d9ce Mon Sep 17 00:00:00 2001 From: Roland Tannous Date: Fri, 20 Feb 2026 08:27:09 +0000 Subject: [PATCH 04/23] added NODE OPTIONS export , updating npm to 2.2.6 --- setup.sh | 2 ++ 1 file changed, 2 insertions(+) diff --git a/setup.sh b/setup.sh index 2db0c7f5a5..4ebf960b73 100755 --- a/setup.sh +++ b/setup.sh @@ -58,6 +58,7 @@ fi if [ "$NEED_NODE" = true ]; then # ── 2. Install nvm ── + export NODE_OPTIONS=--dns-result-order=ipv4first # or else fails on colab. echo "Installing nvm..." curl -so- https://raw.githubusercontent.com/nvm-sh/nvm/v0.40.1/install.sh | bash > /dev/null 2>&1 @@ -165,6 +166,7 @@ if [ "$IS_COLAB" = true ]; then -o "$LLAMA_CPP_DST" echo " Installing studio dependencies..." run_quiet "pip install studio" pip install -r "$SCRIPT_DIR/studio/backend/requirements/studio.txt" + run_quiet "pip install numpy==2.2.6" pip install --force-reinstall numpy==2.2.6 echo "✅ Python dependencies installed" else # Local: create venv (always start fresh to preserve correct install order) From 5ba8edf9fe7934a142bc54bd2076f877463f40bd Mon Sep 17 00:00:00 2001 From: samit Date: Fri, 20 Feb 2026 00:49:38 -0800 Subject: [PATCH 05/23] added the copy on mac --- .../src/components/assistant-ui/thread.tsx | 51 +++++++++++-------- studio/frontend/src/lib/copy-to-clipboard.ts | 44 ++++++++++++++++ 2 files changed, 74 insertions(+), 21 deletions(-) create mode 100644 studio/frontend/src/lib/copy-to-clipboard.ts diff --git a/studio/frontend/src/components/assistant-ui/thread.tsx b/studio/frontend/src/components/assistant-ui/thread.tsx index c109c92ead..aaaca0eeb3 100644 --- a/studio/frontend/src/components/assistant-ui/thread.tsx +++ b/studio/frontend/src/components/assistant-ui/thread.tsx @@ -8,6 +8,7 @@ import { Reasoning, ReasoningGroup } from "@/components/assistant-ui/reasoning"; import { ToolFallback } from "@/components/assistant-ui/tool-fallback"; import { TooltipIconButton } from "@/components/assistant-ui/tooltip-icon-button"; import { Button } from "@/components/ui/button"; +import { copyToClipboard } from "@/lib/copy-to-clipboard"; import { cn } from "@/lib/utils"; import { ActionBarMorePrimitive, @@ -38,7 +39,7 @@ import { RefreshCwIcon, SquareIcon, } from "lucide-react"; -import { type FC, useRef } from "react"; +import { type FC, useRef, useState } from "react"; export const Thread: FC<{ hideComposer?: boolean; hideWelcome?: boolean }> = ({ hideComposer, @@ -275,6 +276,32 @@ const AssistantMessage: FC = () => { ); }; +const COPY_RESET_MS = 2000; + +const CopyButton: FC = () => { + const aui = useAui(); + const [copied, setCopied] = useState(false); + const resetTimeoutRef = useRef | null>(null); + + const handleCopy = () => { + const text = aui.message().getCopyText(); + if (copyToClipboard(text)) { + setCopied(true); + if (resetTimeoutRef.current) clearTimeout(resetTimeoutRef.current); + resetTimeoutRef.current = setTimeout(() => { + setCopied(false); + resetTimeoutRef.current = null; + }, COPY_RESET_MS); + } + }; + + return ( + + {copied ? : } + + ); +}; + const AssistantActionBar: FC = () => { return ( { autohideFloat="single-branch" className="aui-assistant-action-bar-root col-start-3 row-start-2 -ml-1 flex gap-1 text-muted-foreground data-floating:absolute data-floating:rounded-md data-floating:border data-floating:bg-background data-floating:p-1 data-floating:shadow-sm" > - - - message.isCopied}> - - - !message.isCopied}> - - - - + @@ -352,16 +370,7 @@ const UserActionBar: FC = () => { autohide="not-last" className="aui-user-action-bar-root flex items-center" > - - - message.isCopied}> - - - !message.isCopied}> - - - - + diff --git a/studio/frontend/src/lib/copy-to-clipboard.ts b/studio/frontend/src/lib/copy-to-clipboard.ts new file mode 100644 index 0000000000..3ef3df1177 --- /dev/null +++ b/studio/frontend/src/lib/copy-to-clipboard.ts @@ -0,0 +1,44 @@ +/** + * Copy text to clipboard in a way that works on Mac/Safari. + * Uses a synchronous textarea + execCommand fallback so the copy runs in the + * same user gesture as the click (required by Safari's clipboard security). + */ +export function copyToClipboard(text: string): boolean { + if (typeof text !== "string" || text.length === 0) { + return false; + } + + // Synchronous fallback: works in Safari/Mac when clipboard API fails + // because it runs entirely within the user gesture (click) stack. + if (document.queryCommandSupported?.("copy") !== false) { + const textarea = document.createElement("textarea"); + textarea.value = text; + textarea.style.position = "fixed"; + textarea.style.top = "0"; + textarea.style.left = "0"; + textarea.style.opacity = "0"; + textarea.setAttribute("aria-hidden", "true"); + document.body.appendChild(textarea); + textarea.focus({ preventScroll: true }); + textarea.select(); + try { + const ok = document.execCommand("copy"); + document.body.removeChild(textarea); + return ok; + } catch { + document.body.removeChild(textarea); + return false; + } + } + + // Modern API only when fallback not available (e.g. non-browser) + if (typeof navigator?.clipboard?.writeText === "function") { + navigator.clipboard.writeText(text).then( + () => {}, + () => {} + ); + return true; + } + + return false; +} From 3d403c6c992885ba4ca7951a21c8b6c77f6e3da2 Mon Sep 17 00:00:00 2001 From: samit Date: Fri, 20 Feb 2026 01:29:37 -0800 Subject: [PATCH 06/23] edited the font of the new parameters --- .../src/features/studio/sections/params-section.tsx | 6 ++---- 1 file changed, 2 insertions(+), 4 deletions(-) diff --git a/studio/frontend/src/features/studio/sections/params-section.tsx b/studio/frontend/src/features/studio/sections/params-section.tsx index edc21e0e0e..de82fb0fb5 100644 --- a/studio/frontend/src/features/studio/sections/params-section.tsx +++ b/studio/frontend/src/features/studio/sections/params-section.tsx @@ -534,7 +534,7 @@ export function ParamsSection(): ReactElement { value={store.optimizerType} onValueChange={(v) => store.setOptimizerType(v)} > - + @@ -542,7 +542,6 @@ export function ParamsSection(): ReactElement { {opt.label} @@ -571,7 +570,7 @@ export function ParamsSection(): ReactElement { value={store.lrSchedulerType} onValueChange={(v) => store.setLrSchedulerType(v)} > - + @@ -579,7 +578,6 @@ export function ParamsSection(): ReactElement { {opt.label} From 759ae059dba9e868a2c8af780d317c59ea214f33 Mon Sep 17 00:00:00 2001 From: imagineer99 Date: Fri, 20 Feb 2026 11:33:20 +0000 Subject: [PATCH 07/23] Fix: model and dataset dropdowns selecting stale value on Enter --- .../studio/sections/dataset-section.tsx | 22 ++++++++++++++----- .../studio/sections/model-section.tsx | 17 ++++++++++++-- 2 files changed, 31 insertions(+), 8 deletions(-) diff --git a/studio/frontend/src/features/studio/sections/dataset-section.tsx b/studio/frontend/src/features/studio/sections/dataset-section.tsx index e3f50efb40..97bd1304f0 100644 --- a/studio/frontend/src/features/studio/sections/dataset-section.tsx +++ b/studio/frontend/src/features/studio/sections/dataset-section.tsx @@ -109,7 +109,7 @@ export function DatasetSection() { const resultIds = useMemo(() => { const ids = hfResults.map((r) => r.id); if (dataset && !ids.includes(dataset)) { - ids.unshift(dataset); + ids.push(dataset); } return ids; }, [hfResults, dataset]); @@ -159,7 +159,20 @@ export function DatasetSection() { -
+
{ + if (event.key !== "Enter") return; + if (!(event.target instanceof HTMLInputElement)) return; + event.preventDefault(); + if (hfResults.length > 0) { + handleDatasetSelect(hfResults[0].id); + } else { + const text = event.target.value.trim(); + if (text) handleDatasetSelect(text); + } + }} + > id} autoHighlight={true} > - + diff --git a/studio/frontend/src/features/studio/sections/model-section.tsx b/studio/frontend/src/features/studio/sections/model-section.tsx index 32810e0700..501a4a41e3 100644 --- a/studio/frontend/src/features/studio/sections/model-section.tsx +++ b/studio/frontend/src/features/studio/sections/model-section.tsx @@ -160,7 +160,7 @@ export function ModelSection() { const resultIds = useMemo(() => { const ids = hfResults.map((r) => r.id); if (selectedModel && !ids.includes(selectedModel)) { - ids.unshift(selectedModel); + ids.push(selectedModel); } return ids; }, [hfResults, selectedModel]); @@ -375,7 +375,20 @@ export function ModelSection() { -
+
{ + if (event.key !== "Enter") return; + if (!(event.target instanceof HTMLInputElement)) return; + event.preventDefault(); + if (hfResults.length > 0) { + handleModelSelect(hfResults[0].id); + } else { + const text = event.target.value.trim(); + if (text) handleModelSelect(text); + } + }} + > Date: Fri, 20 Feb 2026 15:35:06 +0000 Subject: [PATCH 08/23] fix: remove warmup text inference status --- studio/frontend/src/features/chat/api/chat-adapter.ts | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/studio/frontend/src/features/chat/api/chat-adapter.ts b/studio/frontend/src/features/chat/api/chat-adapter.ts index 8997570271..8e9e7df0e0 100644 --- a/studio/frontend/src/features/chat/api/chat-adapter.ts +++ b/studio/frontend/src/features/chat/api/chat-adapter.ts @@ -159,7 +159,7 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter { if (abortSignal.aborted) return; warmupToastShown = true; toast.promise(firstTokenPromise, { - loading: "Warming up model", + loading: "Generating", success: "Generating", error: (err) => err instanceof Error && err.message ? err.message : "Generation failed", From 34131da9a4f29c69a14b456c439dbaf1fec49d35 Mon Sep 17 00:00:00 2001 From: Roland Tannous Date: Fri, 20 Feb 2026 18:08:58 +0000 Subject: [PATCH 09/23] moved transformers4.57.1 to no-extra-deps --- setup.sh | 1 - studio/backend/requirements/extras-no-deps.txt | 1 + studio/backend/requirements/overrides.txt | 1 - 3 files changed, 1 insertion(+), 2 deletions(-) diff --git a/setup.sh b/setup.sh index 4ebf960b73..b54f782ded 100755 --- a/setup.sh +++ b/setup.sh @@ -166,7 +166,6 @@ if [ "$IS_COLAB" = true ]; then -o "$LLAMA_CPP_DST" echo " Installing studio dependencies..." run_quiet "pip install studio" pip install -r "$SCRIPT_DIR/studio/backend/requirements/studio.txt" - run_quiet "pip install numpy==2.2.6" pip install --force-reinstall numpy==2.2.6 echo "✅ Python dependencies installed" else # Local: create venv (always start fresh to preserve correct install order) diff --git a/studio/backend/requirements/extras-no-deps.txt b/studio/backend/requirements/extras-no-deps.txt index c2a9c6bad8..b78b479e51 100644 --- a/studio/backend/requirements/extras-no-deps.txt +++ b/studio/backend/requirements/extras-no-deps.txt @@ -11,3 +11,4 @@ git+https://github.com/meta-pytorch/OpenEnv.git executorch==1.0.1 torch-c-dlpack-ext sentence_transformers==5.2.0 +transformers==4.57.1 diff --git a/studio/backend/requirements/overrides.txt b/studio/backend/requirements/overrides.txt index 02770f3953..6852f601ed 100644 --- a/studio/backend/requirements/overrides.txt +++ b/studio/backend/requirements/overrides.txt @@ -1,6 +1,5 @@ # Torch AO overrides (installed with --force-reinstall --no-cache-dir) torchao==0.14.0 -transformers==4.57.1 pytorch_tokenizers # Kernel packages From 48e232b38cc616d8f0182c710ddbe5d236d8251d Mon Sep 17 00:00:00 2001 From: Roland Tannous Date: Fri, 20 Feb 2026 18:24:48 +0000 Subject: [PATCH 10/23] add huggingface-hub==0.36.0 due to colab error --- studio/backend/requirements/studio.txt | 1 + 1 file changed, 1 insertion(+) diff --git a/studio/backend/requirements/studio.txt b/studio/backend/requirements/studio.txt index fd7da2626d..916875e1fc 100644 --- a/studio/backend/requirements/studio.txt +++ b/studio/backend/requirements/studio.txt @@ -11,3 +11,4 @@ pyjwt easydict addict gradio>=4.0.0 +huggingface-hub==0.36.0 From 08ff8de31d8606236fd7d9c4d64bce50bd78f448 Mon Sep 17 00:00:00 2001 From: Roland Tannous Date: Fri, 20 Feb 2026 18:26:20 +0000 Subject: [PATCH 11/23] add huggingface-hub==0.36.0 due to colab error --- studio/backend/requirements/studio.txt | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/studio/backend/requirements/studio.txt b/studio/backend/requirements/studio.txt index 916875e1fc..6a732664d2 100644 --- a/studio/backend/requirements/studio.txt +++ b/studio/backend/requirements/studio.txt @@ -11,4 +11,4 @@ pyjwt easydict addict gradio>=4.0.0 -huggingface-hub==0.36.0 +huggingface-hub==0.36.0 \ No newline at end of file From a77b9717f81a18d80b11d3f23e015f52b79b0b1b Mon Sep 17 00:00:00 2001 From: Roland Tannous Date: Fri, 20 Feb 2026 19:15:03 +0000 Subject: [PATCH 12/23] removed branch from colab git clone --- Unsloth_Studio_Colab.ipynb | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/Unsloth_Studio_Colab.ipynb b/Unsloth_Studio_Colab.ipynb index fb2ca72604..7778b2254a 100644 --- a/Unsloth_Studio_Colab.ipynb +++ b/Unsloth_Studio_Colab.ipynb @@ -63,7 +63,7 @@ "\n", "import os\n", "github_token = os.environ['GITHUB_TOKEN']\n", - "!git clone -b feature/colab-notebook https://{github_token}@github.com/unslothai/new-ui-prototype.git\n", + "!git clone https://{github_token}@github.com/unslothai/new-ui-prototype.git\n", "%cd /content/new-ui-prototype\n", "\n", "# Run setup script\n", From 3fa9e773c22858103190219262b5e70a0f35f543 Mon Sep 17 00:00:00 2001 From: Manan17 Date: Fri, 20 Feb 2026 22:23:26 +0000 Subject: [PATCH 13/23] fixed the vlm's text only errors --- studio/backend/core/inference/inference.py | 38 +++++++++------------- 1 file changed, 15 insertions(+), 23 deletions(-) diff --git a/studio/backend/core/inference/inference.py b/studio/backend/core/inference/inference.py index 99a7c485e7..11408399c0 100644 --- a/studio/backend/core/inference/inference.py +++ b/studio/backend/core/inference/inference.py @@ -576,8 +576,8 @@ class InferenceBackend: tokenizer = model_info.get("tokenizer") or model_info.get("processor") top_k = self._normalize_top_k(top_k) - if is_vision: - # Vision model generation + if is_vision and image: + # Vision model generation (only when an image is actually provided) yield from self._generate_vision_response( messages, system_prompt, image, temperature, top_p, top_k, min_p, max_new_tokens, repetition_penalty, @@ -601,7 +601,7 @@ class InferenceBackend: # This modifies the tokenizer with the correct template tokenizer = get_chat_template( tokenizer, - self.active_model_name + chat_template=template_name, ) else: logger.info(f"No registered template for {self.active_model_name}, using tokenizer default") @@ -635,6 +635,9 @@ class InferenceBackend: model_info = self.models[self.active_model_name] model = model_info["model"] processor = model_info["processor"] + # FastVisionModel may return a raw tokenizer (e.g. GemmaTokenizerFast) + # instead of a Processor for some models. Safe unwrap for tokenize-only ops. + raw_tokenizer = getattr(processor, "tokenizer", processor) # Extract user message user_message = "" @@ -668,7 +671,7 @@ class InferenceBackend: else: # Text-only for vision model formatted_prompt = self.format_chat_prompt(messages, system_prompt) - inputs = processor.tokenizer(formatted_prompt, return_tensors="pt").to(self.device) + inputs = raw_tokenizer(formatted_prompt, return_tensors="pt").to(self.device) # Stream with TextIteratorStreamer + background thread try: @@ -676,7 +679,7 @@ class InferenceBackend: import threading streamer = TextIteratorStreamer( - processor.tokenizer, + raw_tokenizer, skip_prompt=True, skip_special_tokens=True, timeout=0.2, @@ -1116,24 +1119,13 @@ class InferenceBackend: return img def _clean_generated_text(self, text: str) -> str: - import re - - text = re.sub(r'<\|start_header_id\|>.*?<\|end_header_id\|>', '', text) - text = re.sub(r'<\|eot_id\|>', '', text) - text = re.sub(r'<\|begin_of_text\|>', '', text) - - text = re.sub(r'\[INST\].*?\[/INST\]', '', text) - text = re.sub(r'|', '', text) - - # Clean ChatML tokens (used by Qwen2-VL and similar models) - text = re.sub(r'<\|im_start\|>.*?<\|im_end\|>', '', text) - text = re.sub(r'<\|im_end\|>', '', text) - text = re.sub(r'<\|im_start\|>', '', text) - - text = re.sub(r'^\s*(assistant|user|system):\s*', '', text, flags=re.IGNORECASE) - text = text.strip() - - return text + """Strip leaked special tokens using the tokenizer's own token list.""" + tokenizer = self.models.get(self.active_model_name, {}).get("tokenizer") + if tokenizer: + for token in getattr(tokenizer, "all_special_tokens", []): + if token in text: + text = text.replace(token, "") + return text.strip() def _load_chat_template_info(self, model_name: str): if model_name not in self.models or not self.models[model_name].get("tokenizer"): From 0a3beade35bbcd8b8afdf3f53cbae147b142123c Mon Sep 17 00:00:00 2001 From: samit Date: Fri, 20 Feb 2026 15:49:33 -0800 Subject: [PATCH 14/23] updated to edit loading as downloading model --- .../frontend/src/features/chat/chat-page.tsx | 15 +++++++++++- .../chat/hooks/use-chat-model-runtime.ts | 24 ++++++++++++------- 2 files changed, 30 insertions(+), 9 deletions(-) diff --git a/studio/frontend/src/features/chat/chat-page.tsx b/studio/frontend/src/features/chat/chat-page.tsx index 34443ec87f..1184183088 100644 --- a/studio/frontend/src/features/chat/chat-page.tsx +++ b/studio/frontend/src/features/chat/chat-page.tsx @@ -5,6 +5,7 @@ import { } from "@/components/assistant-ui/model-selector"; import { Thread } from "@/components/assistant-ui/thread"; import { Button } from "@/components/ui/button"; +import { Spinner } from "@/components/ui/spinner"; import { SidebarProvider, SidebarTrigger, useSidebar } from "@/components/ui/sidebar"; import { Sheet, @@ -283,7 +284,8 @@ export function ChatPage(): ReactElement { const modelsFromStore = useChatRuntimeStore((state) => state.models); const lorasFromStore = useChatRuntimeStore((state) => state.loras); const modelsError = useChatRuntimeStore((state) => state.modelsError); - const { refresh, selectModel, ejectModel } = useChatModelRuntime(); + const { refresh, selectModel, ejectModel, loadingModel } = + useChatModelRuntime(); const refreshRef = useRef(refresh); const selectModelRef = useRef(selectModel); @@ -518,6 +520,17 @@ export function ChatPage(): ReactElement { contentDataTour="chat-model-selector-popover" className="max-w-[62vw] sm:max-w-none" /> + {loadingModel ? ( +
+ + + Downloading model… + +
+ ) : null}
{modelsError && (
diff --git a/studio/frontend/src/features/chat/hooks/use-chat-model-runtime.ts b/studio/frontend/src/features/chat/hooks/use-chat-model-runtime.ts index 03aba01736..175ccfba65 100644 --- a/studio/frontend/src/features/chat/hooks/use-chat-model-runtime.ts +++ b/studio/frontend/src/features/chat/hooks/use-chat-model-runtime.ts @@ -1,4 +1,4 @@ -import { useCallback } from "react"; +import { useCallback, useState } from "react"; import { toast } from "sonner"; import { getInferenceStatus, @@ -116,6 +116,11 @@ export function useChatModelRuntime() { const setCheckpoint = useChatRuntimeStore((state) => state.setCheckpoint); const clearCheckpoint = useChatRuntimeStore((state) => state.clearCheckpoint); + const [loadingModel, setLoadingModel] = useState<{ + id: string; + displayName: string; + } | null>(null); + const refresh = useCallback(async () => { setModelsError(null); try { @@ -157,6 +162,7 @@ export function useChatModelRuntime() { const displayName = model?.name || lora?.name || modelId; setModelsError(null); + setLoadingModel({ id: modelId, displayName }); try { async function performLoad(): Promise { if (params.checkpoint) { @@ -176,19 +182,20 @@ export function useChatModelRuntime() { await refresh(); } - let description = "Base model selected."; - if (isLora) { - description = "Fine-tuned (LoRA) selected."; - } + const loadPromise = performLoad().finally(() => { + setLoadingModel(null); + }); - await toast.promise(performLoad(), { - loading: `Loading ${displayName}`, + await toast.promise(loadPromise, { + loading: "Downloading model…", success: `${displayName} loaded`, error: (err) => err instanceof Error ? err.message : "Failed to load model", - description, + description: + "This may include downloading. Large models can take a while.", }); } catch (error) { + setLoadingModel(null); const message = error instanceof Error ? error.message : "Failed to load model"; setModelsError(message); @@ -227,5 +234,6 @@ export function useChatModelRuntime() { refresh, selectModel, ejectModel, + loadingModel, }; } From 08c3c80d3153a34d654d9b58160bdb7520fe80c0 Mon Sep 17 00:00:00 2001 From: samit Date: Fri, 20 Feb 2026 17:25:03 -0800 Subject: [PATCH 15/23] added vram fit indicator to models in chat --- .../assistant-ui/model-selector/pickers.tsx | 120 +++++++++++++++--- 1 file changed, 105 insertions(+), 15 deletions(-) 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 b90cdd0e3b..e023e36e2b 100644 --- a/studio/frontend/src/components/assistant-ui/model-selector/pickers.tsx +++ b/studio/frontend/src/components/assistant-ui/model-selector/pickers.tsx @@ -1,7 +1,14 @@ import { Input } from "@/components/ui/input"; import { Spinner } from "@/components/ui/spinner"; -import { useDebouncedValue, useHfModelSearch, useInfiniteScroll } from "@/hooks"; +import { + Tooltip, + TooltipContent, + TooltipTrigger, +} from "@/components/ui/tooltip"; +import { useDebouncedValue, useGpuInfo, useHfModelSearch, useInfiniteScroll } from "@/hooks"; import { cn, formatCompact } from "@/lib/utils"; +import type { VramFitStatus } from "@/lib/vram"; +import { checkVramFit, estimateLoadingVram } from "@/lib/vram"; import { Search01Icon } from "@hugeicons/core-free-icons"; import { HugeiconsIcon } from "@hugeicons/react"; import { useMemo, useState, type ReactNode } from "react"; @@ -28,27 +35,77 @@ function ModelRow({ meta, selected, onClick, + vramStatus, + vramEst, + gpuGb, }: { label: string; meta?: string; selected?: boolean; onClick: () => void; + vramStatus?: VramFitStatus | null; + vramEst?: number; + gpuGb?: number; }) { - return ( + const exceeds = vramStatus === "exceeds"; + const showVramTooltip = + vramEst != null && vramEst > 0 && gpuGb != null && gpuGb > 0; + const vramTooltipText = + showVramTooltip && vramStatus + ? exceeds + ? `Needs ~${vramEst}GB VRAM (GPU: ${gpuGb}GB)` + : vramStatus === "tight" + ? `~${vramEst}GB VRAM (tight fit on ${gpuGb}GB)` + : `~${vramEst}GB VRAM` + : null; + + const content = ( ); + + if (vramTooltipText) { + return ( + + {content} + + {label} + {vramTooltipText} + + + ); + } + return content; } export function HubModelPicker({ @@ -60,6 +117,7 @@ export function HubModelPicker({ value?: string; onSelect: (id: string, meta: ModelSelectorChangeMeta) => void; }) { + const gpu = useGpuInfo(); const [query, setQuery] = useState(""); const debouncedQuery = useDebouncedValue(query); const { results, isLoading, isLoadingMore, fetchMore } = useHfModelSearch( @@ -94,6 +152,30 @@ export function HubModelPicker({ [results], ); + const vramMap = useMemo(() => { + const map = new Map< + string, + { est: number; status: VramFitStatus | null; detail: string | null } + >(); + for (const r of results) { + const detail = r.totalParams + ? formatCompact(r.totalParams) + : r.downloads != null + ? `↓${formatCompact(r.downloads)}` + : null; + if (r.totalParams) { + const est = estimateLoadingVram(r.totalParams, "qlora"); + const status = gpu.available + ? checkVramFit(est, gpu.memoryTotalGb) + : null; + map.set(r.id, { est, status, detail }); + } else { + map.set(r.id, { est: 0, status: null, detail }); + } + } + return map; + }, [results, gpu]); + const { scrollRef, sentinelRef } = useInfiniteScroll(fetchMore, results.length); return ( @@ -144,15 +226,23 @@ export function HubModelPicker({ No matching models.
) : ( - hfIds.map((id) => ( - onSelect(id, { source: "hub", isLora: false })} - /> - )) + hfIds.map((id) => { + const vram = vramMap.get(id); + return ( + + onSelect(id, { source: "hub", isLora: false }) + } + vramStatus={vram?.status ?? null} + vramEst={vram?.est} + gpuGb={gpu.available ? gpu.memoryTotalGb : undefined} + /> + ); + }) )}
{isLoadingMore ? ( From f6ebeb1d42b8886cce9c0a8eb6dbf0e2d712a031 Mon Sep 17 00:00:00 2001 From: Manan17 Date: Sat, 21 Feb 2026 01:57:05 +0000 Subject: [PATCH 16/23] Mapping proper tokenizer for VLMs --- studio/backend/core/inference/inference.py | 7 +++++++ 1 file changed, 7 insertions(+) diff --git a/studio/backend/core/inference/inference.py b/studio/backend/core/inference/inference.py index 11408399c0..62b0b031ee 100644 --- a/studio/backend/core/inference/inference.py +++ b/studio/backend/core/inference/inference.py @@ -574,6 +574,8 @@ class InferenceBackend: model_info = self.models[self.active_model_name] is_vision = model_info.get("is_vision", False) tokenizer = model_info.get("tokenizer") or model_info.get("processor") + # Unwrap processor → raw tokenizer for VLMs on the text path + tokenizer = getattr(tokenizer, "tokenizer", tokenizer) top_k = self._normalize_top_k(top_k) if is_vision and image: @@ -769,7 +771,11 @@ class InferenceBackend: model_info = self.models[self.active_model_name] model = model_info["model"] + # For VLMs the stored "tokenizer" is actually the processor. + # Unwrap to get the real tokenizer so TextIteratorStreamer's + # skip_prompt / skip_special_tokens work correctly. tokenizer = model_info["tokenizer"] + tokenizer = getattr(tokenizer, "tokenizer", tokenizer) try: inputs = tokenizer(prompt, return_tensors="pt").to(model.device) @@ -879,6 +885,7 @@ class InferenceBackend: chat_template_info = self.models[self.active_model_name].get("chat_template_info", {}) tokenizer = self.models[self.active_model_name]["tokenizer"] + tokenizer = getattr(tokenizer, "tokenizer", tokenizer) chat_messages = [] From c051e3d532f7b3b37ba3a302efef03c7090e4bf1 Mon Sep 17 00:00:00 2001 From: Roland Tannous Date: Sat, 21 Feb 2026 04:40:29 +0000 Subject: [PATCH 17/23] fix: load proper vision processor from base model when FastVisionModel returns raw tokenizer, add tokenize=False to vision chat template --- studio/backend/core/inference/inference.py | 126 +++++++++++++-------- 1 file changed, 80 insertions(+), 46 deletions(-) diff --git a/studio/backend/core/inference/inference.py b/studio/backend/core/inference/inference.py index 62b0b031ee..6423e7a128 100644 --- a/studio/backend/core/inference/inference.py +++ b/studio/backend/core/inference/inference.py @@ -107,6 +107,23 @@ class InferenceBackend: # Apply inference optimization FastVisionModel.for_inference(model) + # FastVisionModel may return a raw tokenizer (e.g. GemmaTokenizerFast) + # instead of a proper Processor for some models (e.g. Gemma-3). + # In that case, load the real processor from the base model. + from transformers import ProcessorMixin + if not (isinstance(processor, ProcessorMixin) or hasattr(processor, "image_processor")): + processor_source = config.base_model if config.is_lora else config.identifier + logger.warning( + f"FastVisionModel returned {type(processor).__name__} (no image_processor) " + f"for '{model_name}' — loading proper processor from '{processor_source}'" + ) + from transformers import AutoProcessor + processor = AutoProcessor.from_pretrained( + processor_source, + token=hf_token if hf_token and hf_token.strip() else None, + ) + logger.info(f"Loaded {type(processor).__name__} from {processor_source}") + self.models[model_name]["model"] = model self.models[model_name]["tokenizer"] = processor self.models[model_name]["processor"] = processor @@ -580,55 +597,72 @@ class InferenceBackend: if is_vision and image: # Vision model generation (only when an image is actually provided) - yield from self._generate_vision_response( - messages, system_prompt, image, - temperature, top_p, top_k, min_p, max_new_tokens, repetition_penalty, - cancel_event=cancel_event, + # Check that the stored processor can actually handle images. + # FastVisionModel may return a raw tokenizer (e.g. GemmaTokenizerFast) + # instead of a proper ProcessorMixin for some models (e.g. Gemma-3). + from transformers import ProcessorMixin + processor = model_info.get("processor") + has_image_processing = ( + processor is not None + and (isinstance(processor, ProcessorMixin) or hasattr(processor, "image_processor")) ) - else: - # Text model: Use training pipeline approach - # Messages are already in ChatML format from eval.py - - # Step 1: Apply get_chat_template if model is in mapper - try: - from utils.datasets import MODEL_TO_TEMPLATE_MAPPER, get_tokenizer_chat_template - - model_name_lower = self.active_model_name.lower() - - # Check if model has a registered template - if model_name_lower in MODEL_TO_TEMPLATE_MAPPER: - template_name = MODEL_TO_TEMPLATE_MAPPER[model_name_lower] - logger.info(f"Applying chat template '{template_name}' for {self.active_model_name}") - - # This modifies the tokenizer with the correct template - tokenizer = get_chat_template( - tokenizer, - chat_template=template_name, - ) - else: - logger.info(f"No registered template for {self.active_model_name}, using tokenizer default") - except Exception as e: - logger.warning(f"Could not apply get_chat_template: {e}") - - # Step 2: Format with tokenizer.apply_chat_template() - try: - formatted_prompt = tokenizer.apply_chat_template( - messages, - tokenize=False, - add_generation_prompt=True + if has_image_processing: + yield from self._generate_vision_response( + messages, system_prompt, image, + temperature, top_p, top_k, min_p, max_new_tokens, repetition_penalty, + cancel_event=cancel_event, + ) + return + else: + logger.warning( + f"Model '{self.active_model_name}' is marked as vision but its processor " + f"({type(processor).__name__}) has no image_processor — " + f"falling back to text-only generation (image will be ignored)." ) - logger.debug(f"Formatted prompt: {formatted_prompt[:200]}...") - except Exception as e: - logger.error(f"Error applying chat template: {e}") - # Fallback to manual formatting - formatted_prompt = self.format_chat_prompt(messages, system_prompt) - # Step 3: Generate - yield from self.generate_stream( - formatted_prompt, temperature, top_p, top_k, min_p, max_new_tokens, repetition_penalty, - cancel_event=cancel_event, - _adapter_state=_adapter_state, + # Text path: Use training pipeline approach + # Messages are already in ChatML format from eval.py + + # Step 1: Apply get_chat_template if model is in mapper + try: + from utils.datasets import MODEL_TO_TEMPLATE_MAPPER, get_tokenizer_chat_template + + model_name_lower = self.active_model_name.lower() + + # Check if model has a registered template + if model_name_lower in MODEL_TO_TEMPLATE_MAPPER: + template_name = MODEL_TO_TEMPLATE_MAPPER[model_name_lower] + logger.info(f"Applying chat template '{template_name}' for {self.active_model_name}") + + # This modifies the tokenizer with the correct template + tokenizer = get_chat_template( + tokenizer, + chat_template=template_name, + ) + else: + logger.info(f"No registered template for {self.active_model_name}, using tokenizer default") + except Exception as e: + logger.warning(f"Could not apply get_chat_template: {e}") + + # Step 2: Format with tokenizer.apply_chat_template() + try: + formatted_prompt = tokenizer.apply_chat_template( + messages, + tokenize=False, + add_generation_prompt=True ) + logger.debug(f"Formatted prompt: {formatted_prompt[:200]}...") + except Exception as e: + logger.error(f"Error applying chat template: {e}") + # Fallback to manual formatting + formatted_prompt = self.format_chat_prompt(messages, system_prompt) + + # Step 3: Generate + yield from self.generate_stream( + formatted_prompt, temperature, top_p, top_k, min_p, max_new_tokens, repetition_penalty, + cancel_event=cancel_event, + _adapter_state=_adapter_state, + ) def _generate_vision_response(self, messages, system_prompt, image, temperature, top_p, top_k, min_p, max_new_tokens, @@ -663,7 +697,7 @@ class InferenceBackend: } ] - input_text = processor.apply_chat_template(vision_messages, add_generation_prompt=True) + input_text = processor.apply_chat_template(vision_messages, add_generation_prompt=True, tokenize=False) inputs = processor( image, input_text, From 97f40bdc58c55bf8be66a58c53cf09dc7a2aa284 Mon Sep 17 00:00:00 2001 From: samit Date: Fri, 20 Feb 2026 22:14:27 -0800 Subject: [PATCH 18/23] Added dictate and add attachments feature --- .../src/features/chat/api/chat-adapter.ts | 9 + .../src/features/chat/shared-composer.tsx | 246 +++++++++++++++++- 2 files changed, 243 insertions(+), 12 deletions(-) diff --git a/studio/frontend/src/features/chat/api/chat-adapter.ts b/studio/frontend/src/features/chat/api/chat-adapter.ts index 8e9e7df0e0..af8f10156f 100644 --- a/studio/frontend/src/features/chat/api/chat-adapter.ts +++ b/studio/frontend/src/features/chat/api/chat-adapter.ts @@ -65,6 +65,15 @@ function findLatestUserImageBase64(messages: RunMessages): string | undefined { continue; } + // Image in message.content (e.g. compare view appends content with image parts) + for (const part of message.content ?? []) { + if (part.type === "image" && "image" in part) { + const encoded = extractImageBase64(part.image); + if (encoded) return encoded; + } + } + + // Image in message.attachments (e.g. chat composer) if ("attachments" in message && (message.attachments?.length ?? 0) > 0) { for (const attachment of message.attachments ?? []) { for (const part of attachment.content ?? []) { diff --git a/studio/frontend/src/features/chat/shared-composer.tsx b/studio/frontend/src/features/chat/shared-composer.tsx index f849d3f5b7..5ddcb3947f 100644 --- a/studio/frontend/src/features/chat/shared-composer.tsx +++ b/studio/frontend/src/features/chat/shared-composer.tsx @@ -1,25 +1,103 @@ import { TooltipIconButton } from "@/components/assistant-ui/tooltip-icon-button"; import { Button } from "@/components/ui/button"; import { useAui } from "@assistant-ui/react"; -import { ArrowUpIcon, SquareIcon } from "lucide-react"; +import { ArrowUpIcon, MicIcon, PlusIcon, SquareIcon, XIcon } from "lucide-react"; import { type KeyboardEvent, type MutableRefObject, type ReactElement, type ReactNode, createContext, + useCallback, useContext, useEffect, useRef, useState, } from "react"; +export type CompareMessagePart = + | { type: "text"; text: string } + | { type: "image"; image: string }; + export interface CompareHandle { - append: (content: { type: "text"; text: string }[]) => void; + append: (content: CompareMessagePart[]) => void; cancel: () => void; isRunning: () => boolean; } +const IMAGE_ACCEPT = "image/jpeg,image/png,image/webp,image/gif"; +const MAX_IMAGE_SIZE = 20 * 1024 * 1024; + +function fileToBase64DataURL(file: File): Promise { + return new Promise((resolve, reject) => { + const reader = new FileReader(); + reader.onload = () => resolve(reader.result as string); + reader.onerror = () => reject(new Error("Failed to read image file")); + reader.readAsDataURL(file); + }); +} + +function useDictation( + textareaRef: React.RefObject, + setText: (value: string | ((prev: string) => string)) => void, +) { + const [isDictating, setIsDictating] = useState(false); + const recognitionRef = useRef(null); + + const start = useCallback(() => { + const SpeechRecognitionAPI = + typeof window !== "undefined" && + (window.SpeechRecognition ?? (window as unknown as { webkitSpeechRecognition?: typeof SpeechRecognition }).webkitSpeechRecognition); + if (!SpeechRecognitionAPI) { + return; + } + const recognition = new SpeechRecognitionAPI() as SpeechRecognition; + recognition.continuous = true; + recognition.interimResults = true; + recognition.lang = "en-US"; + recognition.onresult = (event: SpeechRecognitionEvent) => { + const last = event.resultIndex; + const result = event.results[last]; + if (!result?.isFinal) return; + const transcript = result[0]?.transcript?.trim(); + if (transcript) { + setText((prev) => (prev ? `${prev} ${transcript}` : transcript)); + } + }; + recognition.onerror = () => { + setIsDictating(false); + }; + recognition.onend = () => { + setIsDictating(false); + }; + recognition.start(); + recognitionRef.current = recognition; + setIsDictating(true); + }, [setText]); + + const stop = useCallback(() => { + if (recognitionRef.current) { + recognitionRef.current.stop(); + recognitionRef.current = null; + } + setIsDictating(false); + }, []); + + useEffect(() => { + return () => { + if (recognitionRef.current) { + recognitionRef.current.abort(); + } + }; + }, []); + + const supported = + typeof window !== "undefined" && + !!(window.SpeechRecognition ?? (window as unknown as { webkitSpeechRecognition?: unknown }).webkitSpeechRecognition); + + return { isDictating, start, stop, supported }; +} + export type CompareHandles = MutableRefObject>; const CompareHandlesContext = createContext(null); @@ -66,6 +144,37 @@ export function RegisterCompareHandle({ return null; } +type PendingImage = { id: string; file: File }; + +function PendingImageThumb({ + file, + onRemove, +}: { + file: File; + onRemove: () => void; +}): ReactElement { + const [src, setSrc] = useState(null); + useEffect(() => { + const url = URL.createObjectURL(file); + setSrc(url); + return () => URL.revokeObjectURL(url); + }, [file]); + if (!src) return
; + return ( +
+ {file.name} + +
+ ); +} + export function SharedComposer({ handlesRef, }: { @@ -73,7 +182,15 @@ export function SharedComposer({ }): ReactElement { const [text, setText] = useState(""); const [running, setRunning] = useState(false); + const [pendingImages, setPendingImages] = useState([]); + const [dragging, setDragging] = useState(false); const textareaRef = useRef(null); + const fileInputRef = useRef(null); + + const { isDictating, start: startDictation, stop: stopDictation, supported: dictationSupported } = useDictation( + textareaRef, + setText, + ); useEffect(() => { const id = setInterval(() => { @@ -84,23 +201,50 @@ export function SharedComposer({ return () => clearInterval(id); }, [handlesRef]); - function send() { - const msg = text.trim(); - if (!msg) { - return; + const addFiles = useCallback((files: FileList | null) => { + if (!files?.length) return; + const next: PendingImage[] = []; + for (let i = 0; i < files.length; i++) { + const file = files[i]; + if (!file?.type.match(/^image\/(jpeg|png|webp|gif)$/i)) continue; + if (file.size > MAX_IMAGE_SIZE) continue; + next.push({ id: crypto.randomUUID(), file }); } + setPendingImages((prev) => [...prev, ...next]); + }, []); + + const removePendingImage = useCallback((id: string) => { + setPendingImages((prev) => prev.filter((p) => p.id !== id)); + }, []); + + async function send() { + const msg = text.trim(); + if (!msg && pendingImages.length === 0) return; + + const content: CompareMessagePart[] = []; + for (const { file } of pendingImages) { + try { + const image = await fileToBase64DataURL(file); + content.push({ type: "image", image }); + } catch { + // skip failed image + } + } + if (msg) { + content.push({ type: "text", text: msg }); + } + if (content.length === 0) return; - const content: { type: "text"; text: string }[] = [ - { type: "text", text: msg }, - ]; for (const handle of Object.values(handlesRef.current)) { handle.append(content); } setText(""); + setPendingImages([]); textareaRef.current?.focus(); } function stop() { + if (isDictating) stopDictation(); for (const handle of Object.values(handlesRef.current)) { handle.cancel(); } @@ -115,8 +259,33 @@ export function SharedComposer({ } } + const canSend = (text.trim().length > 0 || pendingImages.length > 0) && !running; + return ( -
+
{ + e.preventDefault(); + setDragging(true); + }} + onDragLeave={() => setDragging(false)} + onDrop={(e) => { + e.preventDefault(); + setDragging(false); + addFiles(e.dataTransfer.files); + }} + > + {pendingImages.length > 0 && ( +
+ {pendingImages.map(({ id, file }) => ( + removePendingImage(id)} + /> + ))} +
+ )}