merge: nightly into feature/canvas-lab

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",

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
 

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
@@ -574,59 +591,78 @@ 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:
-            # Vision model generation
-            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,
+        if is_vision and image:
+            # Vision model generation (only when an image is actually provided)
+            # 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,
-                        self.active_model_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,
@@ -635,6 +671,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 = ""
@@ -658,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,
@@ -668,7 +707,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 +715,7 @@ class InferenceBackend:
             import threading
 
             streamer = TextIteratorStreamer(
-                processor.tokenizer,
+                raw_tokenizer,
                 skip_prompt=True,
                 skip_special_tokens=True,
                 timeout=0.2,
@@ -766,7 +805,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)
@@ -876,6 +919,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 = []
 
@@ -1116,24 +1160,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'<s>|</s>', '', 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"):

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

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

studio/backend/requirements/studio.txt

@@ -11,3 +11,4 @@ pyjwt
 easydict
 addict
 gradio>=4.0.0
+huggingface-hub==0.36.0

studio/frontend/src/components/assistant-ui/model-selector/pickers.tsx

@@ -1,7 +1,20 @@
 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,
+  useRecommendedModelVram,
+} 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 +41,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 = (
     <button
       type="button"
       onClick={onClick}
       className={cn(
         "flex w-full items-center justify-between gap-2 rounded-md px-2.5 py-1.5 text-left text-sm transition-colors hover:bg-accent",
         selected && "bg-accent/60",
+        exceeds && "opacity-50",
       )}
     >
-      <span className="min-w-0 flex-1 truncate">{label}</span>
-      {meta ? (
-        <span className="shrink-0 text-[10px] text-muted-foreground">{meta}</span>
-      ) : null}
+      <span
+        className={cn(
+          "min-w-0 flex-1 truncate",
+          exceeds && "line-through decoration-muted-foreground/50",
+        )}
+      >
+        {label}
+      </span>
+      <span className="flex items-center gap-1.5 shrink-0">
+        {vramStatus === "exceeds" && (
+          <span className="text-[9px] font-medium text-red-400">OOM</span>
+        )}
+        {vramStatus === "tight" && (
+          <span className="text-[9px] font-medium text-amber-400">TIGHT</span>
+        )}
+        {vramStatus === "fits" && (
+          <span className="text-[9px] font-medium text-emerald-500/90">FIT</span>
+        )}
+        {meta ? (
+          <span className="text-[10px] text-muted-foreground">{meta}</span>
+        ) : null}
+      </span>
     </button>
   );
+
+  if (vramTooltipText) {
+    return (
+      <Tooltip>
+        <TooltipTrigger asChild>{content}</TooltipTrigger>
+        <TooltipContent side="left" className="max-w-xs break-all">
+          {label}
+          <span className="block text-[10px] mt-1">{vramTooltipText}</span>
+        </TooltipContent>
+      </Tooltip>
+    );
+  }
+  return content;
 }
 
 export function HubModelPicker({
@@ -60,6 +123,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(
@@ -71,6 +135,9 @@ export function HubModelPicker({
     [models, value],
   );
 
+  const { paramCountById: recommendedParamCountById } =
+    useRecommendedModelVram(recommendedIds);
+
   const showHfSection = debouncedQuery.trim().length > 0;
   const recommendedSet = useMemo(() => new Set(recommendedIds), [recommendedIds]);
 
@@ -94,6 +161,49 @@ 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 recommendedVramMap = useMemo(() => {
+    const map = new Map<
+      string,
+      { est: number; status: VramFitStatus | null; detail: string | null }
+    >();
+    for (const id of recommendedIds) {
+      const totalParams = recommendedParamCountById.get(id);
+      if (totalParams) {
+        const est = estimateLoadingVram(totalParams, "qlora");
+        const status = gpu.available
+          ? checkVramFit(est, gpu.memoryTotalGb)
+          : null;
+        const detail = formatCompact(totalParams);
+        map.set(id, { est, status, detail });
+      }
+    }
+    return map;
+  }, [recommendedIds, recommendedParamCountById, gpu]);
+
   const { scrollRef, sentinelRef } = useInfiniteScroll(fetchMore, results.length);
 
   return (
@@ -124,14 +234,23 @@ export function HubModelPicker({
                   No default models.
                 </div>
               ) : (
-                recommendedIds.map((id) => (
-                  <ModelRow
-                    key={id}
-                    label={id}
-                    selected={value === id}
-                    onClick={() => onSelect(id, { source: "hub", isLora: false })}
-                  />
-                ))
+                recommendedIds.map((id) => {
+                  const vram = recommendedVramMap.get(id);
+                  return (
+                    <ModelRow
+                      key={id}
+                      label={id}
+                      meta={vram?.detail ?? undefined}
+                      selected={value === id}
+                      onClick={() =>
+                        onSelect(id, { source: "hub", isLora: false })
+                      }
+                      vramStatus={vram?.status ?? null}
+                      vramEst={vram?.est}
+                      gpuGb={gpu.available ? gpu.memoryTotalGb : undefined}
+                    />
+                  );
+                })
               )}
             </>
           ) : null}
@@ -144,15 +263,23 @@ export function HubModelPicker({
                   No matching models.
                 </div>
               ) : (
-                hfIds.map((id) => (
-                  <ModelRow
-                    key={id}
-                    label={id}
-                    meta={metricsById.get(id)}
-                    selected={value === id}
-                    onClick={() => onSelect(id, { source: "hub", isLora: false })}
-                  />
-                ))
+                hfIds.map((id) => {
+                  const vram = vramMap.get(id);
+                  return (
+                    <ModelRow
+                      key={id}
+                      label={id}
+                      meta={metricsById.get(id)}
+                      selected={value === id}
+                      onClick={() =>
+                        onSelect(id, { source: "hub", isLora: false })
+                      }
+                      vramStatus={vram?.status ?? null}
+                      vramEst={vram?.est}
+                      gpuGb={gpu.available ? gpu.memoryTotalGb : undefined}
+                    />
+                  );
+                })
               )}
               <div ref={sentinelRef} className="h-px" />
               {isLoadingMore ? (

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<ReturnType<typeof setTimeout> | 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 (
+    <TooltipIconButton tooltip="Copy" onClick={handleCopy}>
+      {copied ? <CheckIcon /> : <CopyIcon />}
+    </TooltipIconButton>
+  );
+};
+
 const AssistantActionBar: FC = () => {
   return (
     <ActionBarPrimitive.Root
@@ -283,16 +310,7 @@ const AssistantActionBar: FC = () => {
       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"
     >
-      <ActionBarPrimitive.Copy asChild={true}>
-        <TooltipIconButton tooltip="Copy">
-          <AuiIf condition={({ message }) => message.isCopied}>
-            <CheckIcon />
-          </AuiIf>
-          <AuiIf condition={({ message }) => !message.isCopied}>
-            <CopyIcon />
-          </AuiIf>
-        </TooltipIconButton>
-      </ActionBarPrimitive.Copy>
+      <CopyButton />
       <ActionBarPrimitive.Reload asChild={true}>
         <TooltipIconButton tooltip="Refresh">
           <RefreshCwIcon />
@@ -352,16 +370,7 @@ const UserActionBar: FC = () => {
       autohide="not-last"
       className="aui-user-action-bar-root flex items-center"
     >
-      <ActionBarPrimitive.Copy asChild={true}>
-        <TooltipIconButton tooltip="Copy">
-          <AuiIf condition={({ message }) => message.isCopied}>
-            <CheckIcon />
-          </AuiIf>
-          <AuiIf condition={({ message }) => !message.isCopied}>
-            <CopyIcon />
-          </AuiIf>
-        </TooltipIconButton>
-      </ActionBarPrimitive.Copy>
+      <CopyButton />
       <ActionBarPrimitive.Edit asChild={true}>
         <TooltipIconButton tooltip="Edit" className="aui-user-action-edit">
           <PencilIcon />

studio/frontend/src/components/navbar.tsx

@@ -59,7 +59,7 @@ export function Navbar() {
   };
 
   return (
-    <header className="top-0 z-40 h-16 w-full">
+    <header className="relative top-0 z-40 h-16 w-full">
       <div className="mx-auto flex h-full max-w-7xl items-center justify-between px-4 sm:px-6">
         {/* Left: logo */}
         <div

studio/frontend/src/config/training.ts

@@ -73,10 +73,26 @@ 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 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,

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 ?? []) {
@@ -159,7 +168,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",

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 ? (
+              <div
+                className="flex items-center gap-1.5 text-muted-foreground"
+                title={`Loading ${loadingModel.displayName}. This may include downloading.`}
+              >
+                <Spinner className="size-3.5 shrink-0" />
+                <span className="text-xs">
+                  Downloading model…
+                </span>
+              </div>
+            ) : null}
           </div>
           {modelsError && (
             <div className="ml-2 text-xs text-destructive truncate max-w-[28rem]">

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<void> {
           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: "Loading 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,
   };
 }

studio/frontend/src/features/chat/shared-composer.tsx

@@ -1,25 +1,102 @@
 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<string> {
+  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(
+  setText: (value: string | ((prev: string) => string)) => void,
+) {
+  const [isDictating, setIsDictating] = useState(false);
+  const recognitionRef = useRef<SpeechRecognition | null>(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<Record<string, CompareHandle>>;
 
 const CompareHandlesContext = createContext<CompareHandles | null>(null);
@@ -66,6 +143,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<string | null>(null);
+  useEffect(() => {
+    const url = URL.createObjectURL(file);
+    setSrc(url);
+    return () => URL.revokeObjectURL(url);
+  }, [file]);
+  if (!src) return <div className="size-14 animate-pulse rounded-[14px] bg-muted" />;
+  return (
+    <div className="relative size-14 shrink-0 overflow-hidden rounded-[14px] border border-foreground/20 bg-muted">
+      <img src={src} alt={file.name} className="h-full w-full object-cover" />
+      <button
+        type="button"
+        onClick={onRemove}
+        className="absolute top-1 right-1 flex size-5 items-center justify-center rounded-full bg-white text-muted-foreground shadow-sm hover:bg-destructive hover:text-destructive-foreground"
+        aria-label="Remove attachment"
+      >
+        <XIcon className="size-3" />
+      </button>
+    </div>
+  );
+}
+
 export function SharedComposer({
   handlesRef,
 }: {
@@ -73,7 +181,14 @@ export function SharedComposer({
 }): ReactElement {
   const [text, setText] = useState("");
   const [running, setRunning] = useState(false);
+  const [pendingImages, setPendingImages] = useState<PendingImage[]>([]);
+  const [dragging, setDragging] = useState(false);
   const textareaRef = useRef<HTMLTextAreaElement>(null);
+  const fileInputRef = useRef<HTMLInputElement>(null);
+
+  const { isDictating, start: startDictation, stop: stopDictation, supported: dictationSupported } = useDictation(
+    setText,
+  );
 
   useEffect(() => {
     const id = setInterval(() => {
@@ -84,23 +199,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 +257,33 @@ export function SharedComposer({
     }
   }
 
+  const canSend = (text.trim().length > 0 || pendingImages.length > 0) && !running;
+
   return (
-    <div className="shadow-border ring-1 ring-border relative flex w-full flex-col rounded-2xl bg-background px-1 pt-2 transition-shadow">
+    <div
+      className={`shadow-border ring-1 ring-border relative flex w-full flex-col rounded-2xl bg-background px-1 pt-2 transition-shadow outline-none ${dragging ? "ring-ring bg-accent/50" : ""}`}
+      onDragOver={(e) => {
+        e.preventDefault();
+        setDragging(true);
+      }}
+      onDragLeave={() => setDragging(false)}
+      onDrop={(e) => {
+        e.preventDefault();
+        setDragging(false);
+        addFiles(e.dataTransfer.files);
+      }}
+    >
+      {pendingImages.length > 0 && (
+        <div className="mb-2 flex w-full flex-row flex-wrap items-center gap-2 px-1.5 pt-0.5 pb-1">
+          {pendingImages.map(({ id, file }) => (
+            <PendingImageThumb
+              key={id}
+              file={file}
+              onRemove={() => removePendingImage(id)}
+            />
+          ))}
+        </div>
+      )}
       <textarea
         ref={textareaRef}
         value={text}
@@ -126,30 +293,85 @@ export function SharedComposer({
         className="mb-1 max-h-32 min-h-14 w-full resize-none bg-transparent px-4 pt-2 pb-3 text-sm outline-none placeholder:text-muted-foreground"
         rows={1}
       />
-      <div className="relative mx-2 mb-2 flex items-center justify-end">
-        {running ? (
-          <Button
-            type="button"
-            variant="default"
-            size="icon"
-            className="size-8 rounded-full"
-            onClick={stop}
-          >
-            <SquareIcon className="size-3 fill-current" />
-          </Button>
-        ) : (
+      <div className="relative mx-2 mb-2 flex items-center justify-between">
+        <div className="flex items-center gap-1">
+          <input
+            ref={fileInputRef}
+            type="file"
+            accept={IMAGE_ACCEPT}
+            multiple
+            className="hidden"
+            onChange={(e) => {
+              addFiles(e.target.files);
+              e.target.value = "";
+            }}
+          />
           <TooltipIconButton
-            tooltip="Send message"
+            tooltip="Add attachment"
             side="bottom"
-            variant="default"
+            variant="ghost"
             size="icon"
-            className="size-8 rounded-full"
-            onClick={send}
-            disabled={!text.trim()}
+            className="size-8 rounded-full text-muted-foreground hover:bg-muted-foreground/15"
+            onClick={() => fileInputRef.current?.click()}
+            aria-label="Add attachment"
           >
-            <ArrowUpIcon className="size-4" />
+            <PlusIcon className="size-5 stroke-[1.5px]" />
           </TooltipIconButton>
-        )}
+        </div>
+        <div className="flex items-center gap-1">
+          {dictationSupported && (
+            <>
+              {!isDictating ? (
+                <TooltipIconButton
+                  tooltip="Dictate"
+                  side="bottom"
+                  variant="ghost"
+                  size="icon"
+                  className="size-8 rounded-full text-muted-foreground hover:bg-muted-foreground/15"
+                  onClick={startDictation}
+                  aria-label="Dictate"
+                >
+                  <MicIcon className="size-4" />
+                </TooltipIconButton>
+              ) : (
+                <TooltipIconButton
+                  tooltip="Stop dictation"
+                  side="bottom"
+                  variant="ghost"
+                  size="icon"
+                  className="size-8 rounded-full text-destructive"
+                  onClick={stopDictation}
+                  aria-label="Stop dictation"
+                >
+                  <SquareIcon className="size-3 animate-pulse fill-current" />
+                </TooltipIconButton>
+              )}
+            </>
+          )}
+          {running ? (
+            <Button
+              type="button"
+              variant="default"
+              size="icon"
+              className="size-8 rounded-full"
+              onClick={stop}
+            >
+              <SquareIcon className="size-3 fill-current" />
+            </Button>
+          ) : (
+            <TooltipIconButton
+              tooltip="Send message"
+              side="bottom"
+              variant="default"
+              size="icon"
+              className="size-8 rounded-full"
+              onClick={send}
+              disabled={!canSend}
+            >
+              <ArrowUpIcon className="size-4" />
+            </TooltipIconButton>
+          )}
+        </div>
       </div>
     </div>
   );

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({
                 <Spinner className="size-5 text-primary" />
               </div>
               <p className="text-sm text-muted-foreground font-medium">
-                Loading preview...
+                {isHfDataset ? "Fetching dataset preview from Hugging Face..." : "Loading preview..."}
               </p>
+              {isHfDataset && (
+                <p className="text-xs text-muted-foreground/60">
+                  This may take a moment for large datasets
+                </p>
+              )}
             </div>
           )}
 

studio/frontend/src/features/studio/sections/dataset-section.tsx

@@ -114,7 +114,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]);
@@ -164,7 +164,20 @@ export function DatasetSection() {
               </TooltipContent>
             </Tooltip>
           </span>
-          <div ref={comboboxAnchorRef}>
+          <div
+            ref={comboboxAnchorRef}
+            onKeyDown={(event) => {
+              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);
+              }
+            }}
+          >
             <Combobox
               items={resultIds}
               filteredItems={resultIds}
@@ -175,10 +188,7 @@ export function DatasetSection() {
               itemToStringValue={(id) => id}
               autoHighlight={true}
             >
-              <ComboboxInput
-                placeholder="Search datasets..."
-                className="w-full"
-              >
+              <ComboboxInput placeholder="Search datasets..." className="w-full">
                 <InputGroupAddon>
                   <HugeiconsIcon icon={Search01Icon} className="size-4" />
                 </InputGroupAddon>

studio/frontend/src/features/studio/sections/model-section.tsx

@@ -165,7 +165,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]);
@@ -380,7 +380,20 @@ export function ModelSection() {
               </TooltipContent>
             </Tooltip>
           </span>
-          <div ref={comboboxAnchorRef}>
+          <div
+            ref={comboboxAnchorRef}
+            onKeyDown={(event) => {
+              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);
+              }
+            }}
+          >
             <Combobox
               items={resultIds}
               filteredItems={resultIds}

studio/frontend/src/features/studio/sections/params-section.tsx

@@ -20,7 +20,12 @@ import {
   TooltipContent,
   TooltipTrigger,
 } from "@/components/ui/tooltip";
-import { CONTEXT_LENGTHS, TARGET_MODULES } from "@/config/training";
+import {
+  CONTEXT_LENGTHS,
+  LR_SCHEDULER_OPTIONS,
+  OPTIMIZER_OPTIONS,
+  TARGET_MODULES,
+} from "@/config/training";
 import { useTrainingConfigStore } from "@/features/training";
 import type { GradientCheckpointing } from "@/types/training";
 import {
@@ -508,6 +513,78 @@ export function ParamsSection(): ReactElement {
                 value="optimization"
                 className="mt-3 flex flex-col gap-3"
               >
+                <Row
+                  label="Optimizer"
+                  tooltip={
+                    <>
+                      Optimization algorithm. 8-bit variants reduce memory usage.
+                      Fused is recommended for vision models.{" "}
+                      <a
+                        href="https://unsloth.ai/docs/get-started/fine-tuning-llms-guide/lora-hyperparameters-guide"
+                        target="_blank"
+                        rel="noopener noreferrer"
+                        className="text-primary underline"
+                      >
+                        Read more
+                      </a>
+                    </>
+                  }
+                >
+                  <Select
+                    value={store.optimizerType}
+                    onValueChange={(v) => store.setOptimizerType(v)}
+                  >
+                    <SelectTrigger className="w-48">
+                      <SelectValue />
+                    </SelectTrigger>
+                    <SelectContent>
+                      {OPTIMIZER_OPTIONS.map((opt) => (
+                        <SelectItem
+                          key={opt.value}
+                          value={opt.value}
+                        >
+                          {opt.label}
+                        </SelectItem>
+                      ))}
+                    </SelectContent>
+                  </Select>
+                </Row>
+                <Row
+                  label="LR scheduler"
+                  tooltip={
+                    <>
+                      How the learning rate changes over training. Linear decays
+                      steadily; cosine decays in a curve.{" "}
+                      <a
+                        href="https://unsloth.ai/docs/get-started/fine-tuning-llms-guide/lora-hyperparameters-guide"
+                        target="_blank"
+                        rel="noopener noreferrer"
+                        className="text-primary underline"
+                      >
+                        Read more
+                      </a>
+                    </>
+                  }
+                >
+                  <Select
+                    value={store.lrSchedulerType}
+                    onValueChange={(v) => store.setLrSchedulerType(v)}
+                  >
+                    <SelectTrigger className="w-48">
+                      <SelectValue />
+                    </SelectTrigger>
+                    <SelectContent>
+                      {LR_SCHEDULER_OPTIONS.map((opt) => (
+                        <SelectItem
+                          key={opt.value}
+                          value={opt.value}
+                        >
+                          {opt.label}
+                        </SelectItem>
+                      ))}
+                    </SelectContent>
+                  </Select>
+                </Row>
                 <SliderRow
                   label="Batch Size"
                   tooltip={

studio/frontend/src/features/studio/sections/progress-section-lib.ts

@@ -2,6 +2,8 @@ import type { TrainingPhase } from "@/features/training";
 
 export const phaseLabel: Record<TrainingPhase, string> = {
   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<TrainingPhase, string> = {
 
 export const phaseColors: Record<TrainingPhase, string> = {
   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:

studio/frontend/src/features/studio/sections/progress-section.tsx

@@ -35,6 +35,7 @@ import { Link, useNavigate } from "@tanstack/react-router";
 import { useShallow } from "zustand/react/shallow";
 import { useGpuUtilization } from "@/hooks";
 import { setTrainingCompareHandoff } from "@/features/chat";
+import { OPTIMIZER_OPTIONS } from "@/config/training";
 import { formatDuration, formatNumber, phaseColors, phaseLabel } from "./progress-section-lib";
 
 export function ProgressSection(): ReactElement {
@@ -71,6 +72,7 @@ export function ProgressSection(): ReactElement {
       maxSteps: state.maxSteps,
       contextLength: state.contextLength,
       warmupSteps: state.warmupSteps,
+      optimizerType: state.optimizerType,
       loraRank: state.loraRank,
       loraAlpha: state.loraAlpha,
       loraDropout: state.loraDropout,
@@ -126,6 +128,10 @@ export function ProgressSection(): ReactElement {
     ? runtime.currentGradNorm
     : lastNonZeroValue(runtime.gradNormHistory) ?? runtime.currentGradNorm;
 
+  const optimizerLabel =
+    OPTIMIZER_OPTIONS.find((o) => 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],

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";

studio/frontend/src/features/training/api/mappers.ts

@@ -40,8 +40,8 @@ export function buildTrainingStartPayload(
     weight_decay: config.weightDecay,
     random_seed: config.randomSeed,
     packing: config.packing,
-    optim: "adamw_8bit",
-    lr_scheduler_type: "linear",
+    optim: config.optimizerType,
+    lr_scheduler_type: config.lrSchedulerType,
     use_lora: adapterMethod,
     lora_r: config.loraRank,
     lora_alpha: config.loraAlpha,

studio/frontend/src/features/training/api/models-api.ts

@@ -9,6 +9,8 @@ interface BackendTrainingDefaults {
   max_seq_length?: number;
   num_epochs?: number;
   learning_rate?: number | string;
+  optim?: string;
+  lr_scheduler_type?: string;
   batch_size?: number;
   gradient_accumulation_steps?: number;
   warmup_steps?: number;

studio/frontend/src/features/training/lib/model-defaults.ts

@@ -7,6 +7,8 @@ type ModelDefaultsPatch = Partial<
     | "epochs"
     | "contextLength"
     | "learningRate"
+    | "optimizerType"
+    | "lrSchedulerType"
     | "loraRank"
     | "loraAlpha"
     | "loraDropout"
@@ -86,6 +88,12 @@ export function mapBackendModelConfigToTrainingPatch(
   const learningRate = toNumber(training?.learning_rate);
   if (learningRate !== undefined) patch.learningRate = learningRate;
 
+  const optim = toStringValue(training?.optim);
+  if (optim !== undefined) patch.optimizerType = optim;
+
+  const lrSchedulerType = toStringValue(training?.lr_scheduler_type);
+  if (lrSchedulerType !== undefined) patch.lrSchedulerType = lrSchedulerType;
+
   const batchSize = toNumber(training?.batch_size);
   if (batchSize !== undefined) patch.batchSize = batchSize;
 

studio/frontend/src/features/training/stores/training-config-store.ts

@@ -305,6 +305,8 @@ export const useTrainingConfigStore = create<TrainingConfigStore>()(
         setEpochs: (epochs) => set({ epochs }),
         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 }),
@@ -346,7 +348,7 @@ export const useTrainingConfigStore = create<TrainingConfigStore>()(
     },
     {
       name: "unsloth_training_config_v1",
-      version: 3,
+      version: 5,
       migrate: (persisted, version) => {
         const s = persisted as Record<string, unknown>;
         if (version < 2 && s.datasetSubset == null && s.datasetConfig != null) {
@@ -356,6 +358,12 @@ export const useTrainingConfigStore = create<TrainingConfigStore>()(
         if (version < 3 && s.modelDefaultsAppliedFor == null) {
           s.modelDefaultsAppliedFor = null;
         }
+        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,

studio/frontend/src/features/training/types/config.ts

@@ -30,6 +30,8 @@ export interface TrainingConfigState {
   epochs: number;
   contextLength: number;
   learningRate: number;
+  optimizerType: string;
+  lrSchedulerType: string;
   loraRank: number;
   loraAlpha: number;
   loraDropout: number;
@@ -85,6 +87,8 @@ export interface TrainingConfigActions {
   setEpochs: (epochs: number) => void;
   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;

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"

studio/frontend/src/hooks/index.ts

@@ -3,6 +3,7 @@ export { useGpuInfo } from "./use-gpu-info";
 export { useGpuUtilization } from "./use-gpu-utilization";
 export { useHardwareInfo } from "./use-hardware-info";
 export { useHfModelSearch } from "./use-hf-model-search";
+export { useRecommendedModelVram } from "./use-recommended-model-vram";
 export { useHfDatasetSearch } from "./use-hf-dataset-search";
 export { useHfDatasetSplits } from "./use-hf-dataset-splits";
 export { useHfTokenValidation } from "./use-hf-token-validation";

studio/frontend/src/hooks/use-recommended-model-vram.ts

@@ -0,0 +1,57 @@
+import { modelInfo } from "@huggingface/hub";
+import { useEffect, useState } from "react";
+
+/**
+ * Fetches Hugging Face model info (safetensors total param count) for a list of
+ * model IDs. Used to show VRAM fit (FIT / TIGHT / OOM) for recommended/default
+ * models in the chat model dropdown.
+ */
+export function useRecommendedModelVram(ids: string[]) {
+  const [paramCountById, setParamCountById] = useState<
+    Map<string, number>
+  >(new Map());
+  const [isLoading, setIsLoading] = useState(false);
+
+  const stableKey = [...ids].filter(Boolean).sort().join(",");
+
+  useEffect(() => {
+    const stableIds = stableKey ? stableKey.split(",") : [];
+    if (stableIds.length === 0) {
+      setParamCountById(new Map());
+      setIsLoading(false);
+      return;
+    }
+    let canceled = false;
+    void (async () => {
+      setIsLoading(true);
+      const next = new Map<string, number>();
+      await Promise.all(
+        stableIds.map(async (id) => {
+          if (canceled) return;
+          try {
+            const info = await modelInfo({
+              name: id,
+              additionalFields: ["safetensors"],
+            });
+            const raw = info as { safetensors?: { total?: number } };
+            const total = raw.safetensors?.total;
+            if (typeof total === "number" && total > 0) {
+              next.set(id, total);
+            }
+          } catch {
+            // Model not on HF or no safetensors; skip
+          }
+        }),
+      );
+      if (!canceled) {
+        setParamCountById(next);
+        setIsLoading(false);
+      }
+    })();
+    return () => {
+      canceled = true;
+    };
+  }, [stableKey]);
+
+  return { paramCountById, isLoading };
+}

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;
+}

studio/frontend/src/speech-recognition.d.ts

@@ -0,0 +1,49 @@
+/**
+ * Minimal Web Speech API (Speech Recognition) types for browsers that support it.
+ * Full types: @types/dom-speech-recognition
+ */
+interface SpeechRecognitionResultList {
+  readonly length: number;
+  item(index: number): SpeechRecognitionResult;
+  [index: number]: SpeechRecognitionResult;
+}
+
+interface SpeechRecognitionResult {
+  readonly length: number;
+  readonly isFinal: boolean;
+  item(index: number): SpeechRecognitionAlternative;
+  [index: number]: SpeechRecognitionAlternative;
+}
+
+interface SpeechRecognitionAlternative {
+  readonly transcript: string;
+  readonly confidence: number;
+}
+
+interface SpeechRecognitionEvent extends Event {
+  readonly resultIndex: number;
+  readonly results: SpeechRecognitionResultList;
+}
+
+interface SpeechRecognition extends EventTarget {
+  continuous: boolean;
+  interimResults: boolean;
+  lang: string;
+  onresult: ((event: SpeechRecognitionEvent) => void) | null;
+  onerror: ((event: Event) => void) | null;
+  onend: (() => void) | null;
+  start(): void;
+  stop(): void;
+  abort(): void;
+}
+
+interface SpeechRecognitionConstructor {
+  new (): SpeechRecognition;
+}
+
+interface Window {
+  SpeechRecognition?: SpeechRecognitionConstructor;
+  webkitSpeechRecognition?: SpeechRecognitionConstructor;
+}
+
+declare var SpeechRecognition: SpeechRecognitionConstructor | undefined;