From 308bb948d16b90559be8d338fa84a6ff0813fb49 Mon Sep 17 00:00:00 2001
From: Lee Jackson <130007945+Imagineer99@users.noreply.github.com>
Date: Tue, 31 Mar 2026 07:11:20 +0100
Subject: [PATCH] studio: prevent false multimodal warning during model loading
 (#4704)

* studio: gate multimodal incompatibility warning on settled model capabilities

* Also disable Start button during isCheckingVision fallback

When getModelConfig fails and the fallback checkVisionModel is still
in-flight, isLoadingModelDefaults clears before isCheckingVision does.
Without also gating on isCheckingVision the Start button briefly
re-enables with stale capability flags.

Add isCheckingVision to the disabled condition and show "Loading
model..." text while either flag is active.

* Show correct error message for audio dataset incompatibility

The incompatibility warning always said "switch to a vision model"
even when the actual issue was an audio dataset on a non-audio model.
Now shows an audio-specific message when the mismatch is audio.

* Extract isLoadingModel constant for clarity

Pull the combined model-loading condition into a single constant
reused by the settled check, the disabled prop, and the button label.

---------

Co-authored-by: Daniel Han <danielhanchen@users.noreply.github.com>
---
 .../features/studio/sections/training-section.tsx | 15 ++++++++++-----
 1 file changed, 10 insertions(+), 5 deletions(-)

diff --git a/studio/frontend/src/features/studio/sections/training-section.tsx b/studio/frontend/src/features/studio/sections/training-section.tsx
index a43d3ea873..c30b4cbd77 100644
--- a/studio/frontend/src/features/studio/sections/training-section.tsx
+++ b/studio/frontend/src/features/studio/sections/training-section.tsx
@@ -45,9 +45,12 @@ const placeholderData = [
 export function TrainingSection() {
   const store = useTrainingConfigStore();
   const { isStarting, startError, startTrainingRun } = useTrainingActions();
+  const isLoadingModel = store.isLoadingModelDefaults || store.isCheckingVision;
+  const isModelCapabilitiesSettled = !!store.selectedModel && !isLoadingModel;
   const isIncompatible =
-    (!store.isVisionModel && store.isDatasetImage === true) ||
-    (!store.isAudioModel && store.isDatasetAudio === true);
+    isModelCapabilitiesSettled &&
+    ((!store.isVisionModel && store.isDatasetImage === true) ||
+      (!store.isAudioModel && store.isDatasetAudio === true));
   const configValidation = validateTrainingConfig(store);
   const hasMessage = !!(startError || isIncompatible || (!configValidation.ok && configValidation.message));
   const fileInputRef = useRef<HTMLInputElement>(null);
@@ -157,17 +160,19 @@ export function TrainingSection() {
           data-tour="studio-start"
           className="w-full cursor-pointer bg-gradient-to-r from-emerald-500 to-teal-500 text-white hover:from-emerald-600 hover:to-teal-600"
           onClick={() => void startTrainingRun()}
-          disabled={isStarting || isIncompatible || store.isCheckingDataset || !configValidation.ok}
+          disabled={isStarting || isIncompatible || store.isCheckingDataset || isLoadingModel || !configValidation.ok}
         >
           <HugeiconsIcon icon={Rocket01Icon} className="size-4" />
-          {isStarting ? "Starting..." : store.isCheckingDataset ? "Checking dataset..." : "Start Training"}
+          {isStarting ? "Starting..." : isLoadingModel ? "Loading model..." : store.isCheckingDataset ? "Checking dataset..." : "Start Training"}
         </Button>
         {startError && (
           <p className="text-xs text-red-500 leading-relaxed">{startError}</p>
         )}
         {isIncompatible && (
           <p className="text-xs text-red-500 leading-relaxed">
-            Text model is not compatible with a multimodal dataset. Switch to a vision model or choose a text-only dataset.
+            {!store.isAudioModel && store.isDatasetAudio === true
+              ? "This model does not support audio. Switch to an audio-capable model or choose a non-audio dataset."
+              : "Text model is not compatible with a multimodal dataset. Switch to a vision model or choose a text-only dataset."}
           </p>
         )}
         {!configValidation.ok && configValidation.message && !isIncompatible && (