From 7cc8beb2e6be0dd1fbb8a6fc21feb13bb1163e1c Mon Sep 17 00:00:00 2001
From: Dariton4000 <184837391+Dariton4000@users.noreply.github.com>
Date: Sat, 23 May 2026 20:48:55 +0200
Subject: [PATCH] Studio: add VLM image-size control for training

  Studio vision fine-tuning had no explicit way to cap image resolution, so
  users could not trade visual detail against context and memory use from the
  training UI, YAML config, or API payload. :) Add a nullable `vision_image_size`
  setting that keeps the current model default when unset and applies a
  max-side resize when provided.

  - Add `vision_image_size` to the training request model, route payload, backend
    training config, and frontend API/types plumbing.
  - Validate the value server-side as either null or an integer in the supported
    256-2048 range.
  - Surface an Image Size selector for vision LoRA training with Default plus
    common preset sizes.
  - Include the value in training start payloads only for image-dataset vision
    models, and serialize it into vision-aware YAML configs.
  - Map backend model defaults back into the training store and reset the value
    when reapplying model defaults.
  - Pass the resize through the Torch trainer via `UnslothVisionDataCollator`
    using max-dimension semantics.
  - Apply the same max-dimension resize in the MLX VLM path before mlx-vlm's
    internal collation, preserving aspect ratio and avoiding upscaling.
  - Add backend validation coverage and MLX resize-size tests for the new
    behavior.
---
 studio/backend/core/training/trainer.py       | 12 +++-
 studio/backend/core/training/training.py      |  1 +
 studio/backend/core/training/worker.py        | 62 +++++++++++++++++--
 studio/backend/models/training.py             | 21 +++++++
 studio/backend/routes/training.py             |  1 +
 .../tests/test_mlx_training_worker_config.py  |  9 +++
 .../tests/test_studio_train_validation.py     | 24 +++++++
 studio/frontend/src/config/training.ts        |  1 +
 .../studio/sections/params-section.tsx        | 47 ++++++++++++++
 .../src/features/training/api/mappers.ts      |  4 ++
 .../src/features/training/api/models-api.ts   |  1 +
 .../features/training/lib/model-defaults.ts   |  8 +++
 .../src/features/training/lib/yaml-config.ts  | 42 +++++++------
 .../training/stores/training-config-store.ts  |  6 +-
 .../src/features/training/types/api.ts        |  1 +
 .../src/features/training/types/config.ts     |  2 +
 16 files changed, 216 insertions(+), 26 deletions(-)

diff --git a/studio/backend/core/training/trainer.py b/studio/backend/core/training/trainer.py
index b128fb5338..86ec3314ab 100644
--- a/studio/backend/core/training/trainer.py
+++ b/studio/backend/core/training/trainer.py
@@ -3123,7 +3123,17 @@ class UnslothTrainer:
                 from unsloth.trainer import UnslothVisionDataCollator
 
                 FastVisionModel.for_training(self.model)
-                data_collator = UnslothVisionDataCollator(self.model, self.tokenizer)
+                vision_image_size = training_args.get("vision_image_size")
+                if vision_image_size is None:
+                    data_collator = UnslothVisionDataCollator(self.model, self.tokenizer)
+                else:
+                    logger.info(f"Vision image resize: {vision_image_size} (max dimension)\n")
+                    data_collator = UnslothVisionDataCollator(
+                        self.model,
+                        self.tokenizer,
+                        resize = vision_image_size,
+                        resize_dimension = "max",
+                    )
                 logger.info("Vision data collator configured\n")
 
             # ========== TRAINING CONFIGURATION ==========
diff --git a/studio/backend/core/training/training.py b/studio/backend/core/training/training.py
index d2c2316d45..0af3349c6f 100644
--- a/studio/backend/core/training/training.py
+++ b/studio/backend/core/training/training.py
@@ -193,6 +193,7 @@ class TrainingBackend:
             "hf_token": kwargs.get("hf_token", ""),
             "load_in_4bit": kwargs.get("load_in_4bit", True),
             "max_seq_length": kwargs.get("max_seq_length", 2048),
+            "vision_image_size": kwargs.get("vision_image_size"),
             "hf_dataset": kwargs.get("hf_dataset", ""),
             "local_datasets": kwargs.get("local_datasets"),
             "local_eval_datasets": kwargs.get("local_eval_datasets"),
diff --git a/studio/backend/core/training/worker.py b/studio/backend/core/training/worker.py
index f47a6bd599..c367527ca5 100644
--- a/studio/backend/core/training/worker.py
+++ b/studio/backend/core/training/worker.py
@@ -959,7 +959,43 @@ def _activate_transformers_version(model_name: str) -> None:
     activate_transformers_for_subprocess(model_name)
 
 
-def _adapt_for_mlx_vlm(items):
+def _mlx_vlm_max_resized_size(width: int, height: int, target: int) -> tuple[int, int]:
+    if width <= 0 or height <= 0 or target <= 0:
+        return width, height
+    largest_side = max(width, height)
+    if largest_side <= target:
+        return width, height
+    scale = float(target) / float(largest_side)
+    return max(1, int(round(width * scale))), max(1, int(round(height * scale)))
+
+
+def _resize_mlx_vlm_image(image, resize):
+    if resize is None:
+        return image
+    try:
+        from PIL import Image
+        import numpy as np
+    except ImportError:
+        return image
+    if not isinstance(image, Image.Image):
+        return image
+    image = image.convert("RGB")
+    new_size = _mlx_vlm_max_resized_size(*image.size, int(resize))
+    if new_size != image.size:
+        resampling = getattr(Image, "Resampling", Image).LANCZOS
+        image = image.resize(new_size, resampling)
+    # mlx-vlm's internal collator square-resizes PIL images. Return an ndarray
+    # so Studio's max-dimension resize is the final resize, like trainer.py.
+    return np.asarray(image)
+
+
+def _resize_mlx_vlm_images(value, resize):
+    if isinstance(value, list):
+        return [_resize_mlx_vlm_image(image, resize) for image in value]
+    return _resize_mlx_vlm_image(value, resize)
+
+
+def _adapt_for_mlx_vlm(items, resize = None):
     """Adapt GPU-path VLM dataset output for mlx-vlm consumption.
 
     The GPU path embeds PIL images inside messages content as
@@ -979,7 +1015,7 @@ def _adapt_for_mlx_vlm(items):
                     if isinstance(part, dict) and part.get("type") == "image":
                         img = part.get("image")
                         if img is not None:
-                            images.append(img)
+                            images.append(_resize_mlx_vlm_image(img, resize))
                         new_content.append({"type": "image"})
                     else:
                         new_content.append(part)
@@ -990,9 +1026,9 @@ def _adapt_for_mlx_vlm(items):
         if images:
             out["image"] = images[0] if len(images) == 1 else images
         elif "image" in item:
-            out["image"] = item["image"]
+            out["image"] = _resize_mlx_vlm_images(item["image"], resize)
         elif "images" in item:
-            out["images"] = item["images"]
+            out["images"] = _resize_mlx_vlm_images(item["images"], resize)
         adapted.append(out)
     return adapted
 
@@ -1168,6 +1204,13 @@ 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:
+        vision_image_size = int(vision_image_size)
+        _send(
+            "status",
+            status_message = f"MLX vision image resize: {vision_image_size} (max dimension)",
+        )
 
     # ── 2. Apply LoRA / full FT ──
     # Pass gradient_checkpointing as string ("mlx"/"unsloth"/"none"/etc.)
@@ -1302,7 +1345,10 @@ def _run_mlx_training(event_queue, stop_queue, config):
                 progress_callback = _fmt_progress,
             )
             if vlm_info.get("success"):
-                dataset = _adapt_for_mlx_vlm(vlm_info["dataset"])
+                dataset = _adapt_for_mlx_vlm(
+                    vlm_info["dataset"],
+                    resize = vision_image_size,
+                )
             else:
                 errors = vlm_info.get("errors", [])
                 raise ValueError(
@@ -1317,7 +1363,10 @@ def _run_mlx_training(event_queue, stop_queue, config):
                     dataset_name = hf_dataset or "local",
                 )
                 if ev_info.get("success"):
-                    eval_dataset = _adapt_for_mlx_vlm(ev_info["dataset"])
+                    eval_dataset = _adapt_for_mlx_vlm(
+                        ev_info["dataset"],
+                        resize = vision_image_size,
+                    )
 
         elif format_type:
             _send("status", status_message = f"Formatting dataset ({format_type})...")
@@ -2248,6 +2297,7 @@ def run_training_process(
             eval_dataset = eval_dataset,
             eval_steps = eval_steps,
             max_seq_length = config.get("max_seq_length", 2048),
+            vision_image_size = config.get("vision_image_size"),
             optim = config.get("optim", "adamw_8bit"),
             lr_scheduler_type = config.get("lr_scheduler_type", "linear"),
             is_cpt = is_cpt,
diff --git a/studio/backend/models/training.py b/studio/backend/models/training.py
index 7c53b0fee5..9af7c94241 100644
--- a/studio/backend/models/training.py
+++ b/studio/backend/models/training.py
@@ -18,6 +18,9 @@ _MAX_SEQ_LENGTH = 2_000_000
 _MAX_LR_VALUE = 1.0
 _MAX_LORA_R = 16_384
 _MAX_LORA_ALPHA = 32_768
+_MIN_VISION_IMAGE_SIZE = 256
+# 2048 was the most I could get most llms to work at without getting unstable
+_MAX_VISION_IMAGE_SIZE = 2048
 
 
 def _parse_lr(v: Any) -> float:
@@ -58,6 +61,10 @@ class TrainingStartRequest(BaseModel):
     hf_token: Optional[str] = Field(None, description = "HuggingFace token")
     load_in_4bit: bool = Field(True, description = "Load model in 4-bit quantization")
     max_seq_length: int = Field(2048, description = "Maximum sequence length")
+    vision_image_size: Optional[int] = Field(
+        None,
+        description = "Optional maximum image side length for VLM training. Null uses model default.",
+    )
     trust_remote_code: bool = Field(
         False,
         description = "Allow loading models with custom code (e.g. NVIDIA Nemotron). Only enable for repos you trust.",
@@ -159,6 +166,20 @@ class TrainingStartRequest(BaseModel):
             )
         return v
 
+    @field_validator("vision_image_size")
+    @classmethod
+    def _check_vision_image_size(cls, v: Optional[int]) -> Optional[int]:
+        if v is None:
+            return v
+        if isinstance(v, bool) or not isinstance(v, int):
+            raise ValueError("vision_image_size must be an integer or null")
+        if v < _MIN_VISION_IMAGE_SIZE or v > _MAX_VISION_IMAGE_SIZE:
+            raise ValueError(
+                f"vision_image_size must be in [{_MIN_VISION_IMAGE_SIZE}, "
+                f"{_MAX_VISION_IMAGE_SIZE}] (got {v!r})"
+            )
+        return v
+
     @field_validator("warmup_steps")
     @classmethod
     def _check_warmup_steps(cls, v: Optional[int]) -> Optional[int]:
diff --git a/studio/backend/routes/training.py b/studio/backend/routes/training.py
index 6e2413b3e9..41a9e15562 100644
--- a/studio/backend/routes/training.py
+++ b/studio/backend/routes/training.py
@@ -194,6 +194,7 @@ async def start_training(
             "hf_token": request.hf_token or "",
             "load_in_4bit": request.load_in_4bit,
             "max_seq_length": request.max_seq_length,
+            "vision_image_size": request.vision_image_size,
             "hf_dataset": request.hf_dataset or "",
             "local_datasets": request.local_datasets,
             "local_eval_datasets": request.local_eval_datasets,
diff --git a/studio/backend/tests/test_mlx_training_worker_config.py b/studio/backend/tests/test_mlx_training_worker_config.py
index 98c7bdaa55..cf94705033 100644
--- a/studio/backend/tests/test_mlx_training_worker_config.py
+++ b/studio/backend/tests/test_mlx_training_worker_config.py
@@ -66,6 +66,7 @@ def _load_worker_module():
 _worker = _load_worker_module()
 _normalize_mlx_studio_optimizer = _worker._normalize_mlx_studio_optimizer
 _normalize_mlx_studio_scheduler = _worker._normalize_mlx_studio_scheduler
+_mlx_vlm_max_resized_size = _worker._mlx_vlm_max_resized_size
 
 
 def test_mlx_studio_optimizer_aliases_are_explicit():
@@ -82,3 +83,11 @@ def test_mlx_studio_rejects_unknown_optimizer():
 def test_mlx_studio_rejects_unknown_scheduler():
     with pytest.raises(ValueError, match = "Unsupported LR scheduler for MLX training"):
         _normalize_mlx_studio_scheduler("linear_typo")
+
+
+def test_mlx_vlm_resize_uses_max_dimension_like_torch_trainer():
+    assert _mlx_vlm_max_resized_size(1000, 500, 512) == (512, 256)
+    assert _mlx_vlm_max_resized_size(500, 1000, 512) == (256, 512)
+    assert _mlx_vlm_max_resized_size(1000, 1000, 512) == (512, 512)
+    assert _mlx_vlm_max_resized_size(256, 128, 1536) == (256, 128)
+    assert _mlx_vlm_max_resized_size(512, 256, 512) == (512, 256)
diff --git a/studio/backend/tests/test_studio_train_validation.py b/studio/backend/tests/test_studio_train_validation.py
index 7ffa9bb384..27ba89d8ad 100644
--- a/studio/backend/tests/test_studio_train_validation.py
+++ b/studio/backend/tests/test_studio_train_validation.py
@@ -18,6 +18,8 @@ from models.training import (
     _MAX_LORA_ALPHA,
     _MAX_LORA_R,
     _MAX_SEQ_LENGTH,
+    _MAX_VISION_IMAGE_SIZE,
+    _MIN_VISION_IMAGE_SIZE,
 )
 
 
@@ -62,6 +64,28 @@ class TestBatchSizeCap:
             _check_field("batch_size", 0)
 
 
+class TestVisionImageSizeCap:
+    def test_none_accepts_model_default(self):
+        _check_field("vision_image_size", None)
+
+    @pytest.mark.parametrize(
+        "value",
+        [_MIN_VISION_IMAGE_SIZE, 640, 1000, _MAX_VISION_IMAGE_SIZE],
+    )
+    def test_in_range_accepts(self, value):
+        _check_field("vision_image_size", value)
+        assert _MIN_VISION_IMAGE_SIZE == 256
+        assert _MAX_VISION_IMAGE_SIZE == 2048
+
+    @pytest.mark.parametrize(
+        "value",
+        [_MIN_VISION_IMAGE_SIZE - 1, _MAX_VISION_IMAGE_SIZE + 1, 640.5, True],
+    )
+    def test_invalid_rejects(self, value):
+        with pytest.raises(ValidationError):
+            _check_field("vision_image_size", value)
+
+
 class TestLoraRCap:
     def test_at_cap_accepts(self):
         _check_field("lora_r", _MAX_LORA_R)
diff --git a/studio/frontend/src/config/training.ts b/studio/frontend/src/config/training.ts
index e9fe1d679c..c0a60d0de2 100644
--- a/studio/frontend/src/config/training.ts
+++ b/studio/frontend/src/config/training.ts
@@ -108,6 +108,7 @@ export const LR_DEFAULT_CPT = 5e-5;
 export const DEFAULT_HYPERPARAMS = {
   epochs: 3,
   contextLength: 2048,
+  visionImageSize: null as number | null,
   learningRate: LR_DEFAULT_LORA,
   // null = let backend auto-compute (lr/10 per Unsloth CPT recipe). Only used by CPT.
   embeddingLearningRate: null as number | null,
diff --git a/studio/frontend/src/features/studio/sections/params-section.tsx b/studio/frontend/src/features/studio/sections/params-section.tsx
index 19b75db140..56eedfa651 100644
--- a/studio/frontend/src/features/studio/sections/params-section.tsx
+++ b/studio/frontend/src/features/studio/sections/params-section.tsx
@@ -139,6 +139,7 @@ export function ParamsSection(): ReactElement {
   const [ctxInput, setCtxInput] = useState(String(store.contextLength));
   const ctxAnchorRef = useRef<HTMLDivElement>(null);
   const ctxItems = CONTEXT_LENGTHS.map(String);
+  const visionImageSizePresets = [384, 512, 768, 1024, 1536, 2048];
 
   // Keep input in sync when the store value changes externally
   // (e.g. model defaults being applied after model selection).
@@ -915,6 +916,52 @@ export function ParamsSection(): ReactElement {
                 </TabsContent>
 
                 <TabsContent value="memory" className="mt-3 flex flex-col gap-3">
+                  {showVisionLora && (
+                    <Row
+                      label="Image Size"
+                      tooltip={
+                        <>
+                          Resize images by maximum side length. Default uses the
+                          model image size. Larger images use up more context. Does not upscale or change aspect ratio.{" "}
+                          <a
+                            href="https://unsloth.ai/docs/basics/vision-fine-tuning"
+                            target="_blank"
+                            rel="noopener noreferrer"
+                            className="text-primary underline"
+                          >
+                            Read more
+                          </a>
+                        </>
+                      }
+                    >
+                      <Select
+                        value={
+                          store.visionImageSize == null
+                            ? "default"
+                            : String(store.visionImageSize)
+                        }
+                        onValueChange={(value) => {
+                          if (value === "default") {
+                            store.setVisionImageSize(null);
+                            return;
+                          }
+                          store.setVisionImageSize(Number(value));
+                        }}
+                      >
+                        <SelectTrigger className="w-32">
+                          <SelectValue />
+                        </SelectTrigger>
+                        <SelectContent>
+                          <SelectItem value="default">Default</SelectItem>
+                          {visionImageSizePresets.map((size) => (
+                            <SelectItem key={size} value={String(size)}>
+                              {size}
+                            </SelectItem>
+                          ))}
+                        </SelectContent>
+                      </Select>
+                    </Row>
+                  )}
                   <Row
                     label="Grad Checkpoint"
                     tooltip={
diff --git a/studio/frontend/src/features/training/api/mappers.ts b/studio/frontend/src/features/training/api/mappers.ts
index 5d81f0df9c..75ce8ccdb5 100644
--- a/studio/frontend/src/features/training/api/mappers.ts
+++ b/studio/frontend/src/features/training/api/mappers.ts
@@ -55,6 +55,10 @@ export function buildTrainingStartPayload(
     hf_token: config.hfToken.trim() || null,
     load_in_4bit: (adapterMethod && isQloraMethod) || (isCpt && isFourBitModel),
     max_seq_length: config.contextLength,
+    vision_image_size:
+      config.isVisionModel && config.isDatasetImage === true
+        ? config.visionImageSize
+        : null,
     trust_remote_code: config.trustRemoteCode ?? false,
     hf_dataset: hfDataset,
     subset: hfDataset ? config.datasetSubset : null,
diff --git a/studio/frontend/src/features/training/api/models-api.ts b/studio/frontend/src/features/training/api/models-api.ts
index e512b9e28d..8fecfaf7b8 100644
--- a/studio/frontend/src/features/training/api/models-api.ts
+++ b/studio/frontend/src/features/training/api/models-api.ts
@@ -27,6 +27,7 @@ interface BackendTrainingDefaults {
   eval_steps?: number;
   weight_decay?: number;
   random_seed?: number;
+  vision_image_size?: number | string | null;
   packing?: boolean;
   train_on_completions?: boolean;
   gradient_checkpointing?: "none" | "true" | "unsloth";
diff --git a/studio/frontend/src/features/training/lib/model-defaults.ts b/studio/frontend/src/features/training/lib/model-defaults.ts
index 8bc9c4e064..00008f026a 100644
--- a/studio/frontend/src/features/training/lib/model-defaults.ts
+++ b/studio/frontend/src/features/training/lib/model-defaults.ts
@@ -28,6 +28,7 @@ type ModelDefaultsPatch = Partial<
     | "trainOnCompletions"
     | "gradientCheckpointing"
     | "randomSeed"
+    | "visionImageSize"
     | "enableWandb"
     | "wandbProject"
     | "enableTensorboard"
@@ -129,6 +130,13 @@ export function mapBackendModelConfigToTrainingPatch(
   const randomSeed = toNumber(training?.random_seed);
   if (randomSeed !== undefined) patch.randomSeed = randomSeed;
 
+  if (Object.hasOwn(training ?? {}, "vision_image_size")) {
+    const visionImageSize = training?.vision_image_size == null
+      ? null
+      : toNumber(training.vision_image_size);
+    if (visionImageSize !== undefined) patch.visionImageSize = visionImageSize;
+  }
+
   const packing = toBoolean(training?.packing);
   if (packing !== undefined) patch.packing = packing;
 
diff --git a/studio/frontend/src/features/training/lib/yaml-config.ts b/studio/frontend/src/features/training/lib/yaml-config.ts
index d168da5347..add0eb3520 100644
--- a/studio/frontend/src/features/training/lib/yaml-config.ts
+++ b/studio/frontend/src/features/training/lib/yaml-config.ts
@@ -58,25 +58,31 @@ export function serializeConfigToYaml(
     lora.finetune_mlp_modules = state.finetuneMLPModules;
   }
 
+  const training: Record<string, unknown> = {
+    max_seq_length: state.contextLength,
+    num_epochs: state.epochs,
+    learning_rate: state.learningRate,
+    batch_size: state.batchSize,
+    gradient_accumulation_steps: state.gradientAccumulation,
+    warmup_steps: state.warmupSteps,
+    max_steps: state.maxSteps,
+    save_steps: state.saveSteps,
+    eval_steps: state.evalSteps,
+    weight_decay: state.weightDecay,
+    random_seed: state.randomSeed,
+    packing: state.packing,
+    train_on_completions: state.trainOnCompletions,
+    gradient_checkpointing: state.gradientCheckpointing,
+    optim: state.optimizerType,
+    lr_scheduler_type: state.lrSchedulerType,
+  };
+
+  if (includeVisionFields) {
+    training.vision_image_size = state.visionImageSize;
+  }
+
   const config = {
-    training: {
-      max_seq_length: state.contextLength,
-      num_epochs: state.epochs,
-      learning_rate: state.learningRate,
-      batch_size: state.batchSize,
-      gradient_accumulation_steps: state.gradientAccumulation,
-      warmup_steps: state.warmupSteps,
-      max_steps: state.maxSteps,
-      save_steps: state.saveSteps,
-      eval_steps: state.evalSteps,
-      weight_decay: state.weightDecay,
-      random_seed: state.randomSeed,
-      packing: state.packing,
-      train_on_completions: state.trainOnCompletions,
-      gradient_checkpointing: state.gradientCheckpointing,
-      optim: state.optimizerType,
-      lr_scheduler_type: state.lrSchedulerType,
-    },
+    training,
     lora,
   };
 
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 ef16f641f5..112a394a7a 100644
--- a/studio/frontend/src/features/training/stores/training-config-store.ts
+++ b/studio/frontend/src/features/training/stores/training-config-store.ts
@@ -701,6 +701,7 @@ export const useTrainingConfigStore = create<TrainingConfigStore>()(
         }),
         setEpochs: (epochs) => set({ epochs }),
         setContextLength: (contextLength) => set({ contextLength }),
+        setVisionImageSize: (visionImageSize) => set({ visionImageSize }),
         setLearningRate: (learningRate) => {
           _learningRateManuallySet = true;
           set({ learningRate });
@@ -755,7 +756,10 @@ export const useTrainingConfigStore = create<TrainingConfigStore>()(
         resetToModelDefaults: () => {
           const { selectedModel } = get();
           if (!selectedModel) return;
-          set({ modelDefaultsAppliedFor: null });
+          set({
+            modelDefaultsAppliedFor: null,
+            visionImageSize: DEFAULT_HYPERPARAMS.visionImageSize,
+          });
           loadAndApplyModelDefaults(selectedModel);
         },
         applyConfigPatch: (config: BackendModelConfig) => {
diff --git a/studio/frontend/src/features/training/types/api.ts b/studio/frontend/src/features/training/types/api.ts
index 0cb881e634..5c12e71186 100644
--- a/studio/frontend/src/features/training/types/api.ts
+++ b/studio/frontend/src/features/training/types/api.ts
@@ -7,6 +7,7 @@ export interface TrainingStartRequest {
   hf_token: string | null;
   load_in_4bit: boolean;
   max_seq_length: number;
+  vision_image_size?: number | null;
   /** Allow loading models with custom code. Only enable for repos you trust. */
   trust_remote_code?: boolean;
   hf_dataset: string | null;
diff --git a/studio/frontend/src/features/training/types/config.ts b/studio/frontend/src/features/training/types/config.ts
index 5b156316ca..f40f053e8b 100644
--- a/studio/frontend/src/features/training/types/config.ts
+++ b/studio/frontend/src/features/training/types/config.ts
@@ -82,6 +82,7 @@ export interface TrainingConfigState {
   finetuneMLPModules: boolean;
   targetModules: string[];
   maxPositionEmbeddings: number | null;
+  visionImageSize: number | null;
 }
 
 export interface TrainingConfigActions {
@@ -115,6 +116,7 @@ export interface TrainingConfigActions {
   setUploadedEvalFile: (file: string | null) => void;
   setEpochs: (epochs: number) => void;
   setContextLength: (length: number) => void;
+  setVisionImageSize: (size: number | null) => void;
   setLearningRate: (rate: number) => void;
   setEmbeddingLearningRate: (rate: number | null) => void;
   setOptimizerType: (value: string) => void;