From c797e571c9ec5ee9ca3363d0a9bea882bc4794ac Mon Sep 17 00:00:00 2001
From: Daniel Han <danielhanchen@gmail.com>
Date: Sun, 24 May 2026 15:51:15 +0000
Subject: [PATCH] Studio: extend DeepSeek OCR Image Size exclusion to MLX +
 frontend

Round 4 of the parallel-reviewer pass flagged that the Torch trainer
exclusion I added did not have a matching MLX guard, and that the UI
still offered the dropdown for DeepSeek OCR even though the backend
ignores it.

- worker.py: _run_mlx_training now mirrors the Torch exclusion. When
  the model name matches DeepSeek OCR, vision_image_size is forced
  back to None before _adapt_for_mlx_vlm sees it, so dataset images
  pass through unchanged just like the Torch path. Emits a clear
  status line when this happens.
- params-section.tsx: the Image Size Row is now gated on
  showVisionImageSize (showVisionLora && !isDeepseekOcr) instead of
  showVisionLora alone, so DeepSeek OCR users no longer see a control
  that silently has no effect.
- mappers.ts: buildTrainingStartPayload sends null for vision_image_size
  whenever the selected model is DeepSeek OCR, so the backend log line
  about ignoring the value never fires from a UI-driven start.
---
 studio/backend/core/training/worker.py         | 18 +++++++++++++++++-
 .../studio/sections/params-section.tsx         | 10 +++++++++-
 .../src/features/training/api/mappers.ts       |  9 +++++++--
 3 files changed, 33 insertions(+), 4 deletions(-)

diff --git a/studio/backend/core/training/worker.py b/studio/backend/core/training/worker.py
index 6915bedd4b..d0cba0c85b 100644
--- a/studio/backend/core/training/worker.py
+++ b/studio/backend/core/training/worker.py
@@ -1211,7 +1211,23 @@ def _run_mlx_training(event_queue, stop_queue, config):
     is_vlm = bool(is_dataset_image and getattr(model, "_is_vlm_model", False))
     model._is_vlm_model = is_vlm
     vision_image_size = config.get("vision_image_size")
-    if is_vlm and vision_image_size is not None:
+    # Mirror the Torch trainer.py exclusion: DeepSeek OCR's preset is a tuple
+    # (image_size, base_size, crop_mode), so resizing dataset images outside
+    # that preset desyncs the token grid. Skip the resize on MLX too.
+    _model_name_lower = str(config.get("model_name", "")).lower()
+    _is_deepseek_ocr = (
+        "deepseek" in _model_name_lower and "ocr" in _model_name_lower
+    )
+    if is_vlm and vision_image_size is not None and _is_deepseek_ocr:
+        _send(
+            "status",
+            status_message = (
+                "MLX vision image resize ignored for DeepSeek OCR "
+                "(uses fixed Gundam preset)."
+            ),
+        )
+        vision_image_size = None
+    elif is_vlm and vision_image_size is not None:
         vision_image_size = int(vision_image_size)
         _send(
             "status",
diff --git a/studio/frontend/src/features/studio/sections/params-section.tsx b/studio/frontend/src/features/studio/sections/params-section.tsx
index d6489d677c..799ae4e059 100644
--- a/studio/frontend/src/features/studio/sections/params-section.tsx
+++ b/studio/frontend/src/features/studio/sections/params-section.tsx
@@ -133,6 +133,14 @@ export function ParamsSection(): ReactElement {
   const isCpt = store.trainingMethod === "cpt";
   const isRawText = isRawTextDatasetFormat(store.datasetFormat);
   const showVisionLora = store.isVisionModel && store.isDatasetImage === true;
+  // DeepSeek OCR's preset is a coupled tuple, so the backend ignores any
+  // user-selected image size for it. Hide the control rather than offer a
+  // setting that silently has no effect.
+  const _selectedModelLower = (store.selectedModel ?? "").toLowerCase();
+  const isDeepseekOcr =
+    _selectedModelLower.includes("deepseek") &&
+    _selectedModelLower.includes("ocr");
+  const showVisionImageSize = showVisionLora && !isDeepseekOcr;
   const [loraOpen, setLoraOpen] = useState(false);
   const [hyperOpen, setHyperOpen] = useState(false);
   const needsExpandedHeight = isCpt || (isLora && loraOpen) || hyperOpen;
@@ -917,7 +925,7 @@ export function ParamsSection(): ReactElement {
                 </TabsContent>
 
                 <TabsContent value="memory" className="mt-3 flex flex-col gap-3">
-                  {showVisionLora && (
+                  {showVisionImageSize && (
                     <Row
                       label="Image Size"
                       tooltip={
diff --git a/studio/frontend/src/features/training/api/mappers.ts b/studio/frontend/src/features/training/api/mappers.ts
index 75ce8ccdb5..1223888b11 100644
--- a/studio/frontend/src/features/training/api/mappers.ts
+++ b/studio/frontend/src/features/training/api/mappers.ts
@@ -23,7 +23,12 @@ export function buildTrainingStartPayload(
   const isCpt = config.trainingMethod === "cpt";
   const adapterMethod = config.trainingMethod !== "full";
   const isQloraMethod = config.trainingMethod === "qlora";
-  const isFourBitModel = (config.selectedModel ?? "").toLowerCase().includes("4bit");
+  const _selectedModelLower = (config.selectedModel ?? "").toLowerCase();
+  const isFourBitModel = _selectedModelLower.includes("4bit");
+  // DeepSeek OCR ignores user-selected image size; do not send it.
+  const isDeepseekOcr =
+    _selectedModelLower.includes("deepseek") &&
+    _selectedModelLower.includes("ocr");
   const isEmbedding = config.isEmbeddingModel;
   const isRawText = isRawTextDatasetFormat(config.datasetFormat);
   const hfDataset = config.datasetSource === "huggingface" ? config.dataset : null;
@@ -56,7 +61,7 @@ export function buildTrainingStartPayload(
     load_in_4bit: (adapterMethod && isQloraMethod) || (isCpt && isFourBitModel),
     max_seq_length: config.contextLength,
     vision_image_size:
-      config.isVisionModel && config.isDatasetImage === true
+      config.isVisionModel && config.isDatasetImage === true && !isDeepseekOcr
         ? config.visionImageSize
         : null,
     trust_remote_code: config.trustRemoteCode ?? false,