diff --git a/setup.sh b/setup.sh
index a911d9ed7a..a6d227de2b 100755
--- a/setup.sh
+++ b/setup.sh
@@ -170,7 +170,7 @@ if [ "$IS_COLAB" = true ]; then
     curl -sSL "https://raw.githubusercontent.com/unslothai/unsloth-zoo/refs/heads/main/unsloth_zoo/llama_cpp.py" \
         -o "$LLAMA_CPP_DST"
     # Patch: override vision.py with fix from unsloth PR: https://github.com/unslothai/unsloth/pull/4091 until next pypi release
-    VISION_DST="$(pip show unsloth | grep -i '^Location:' | awk '{print $2}')/unsloth/vision.py"
+    VISION_DST="$(pip show unsloth | grep -i '^Location:' | awk '{print $2}')/unsloth/models/vision.py"
     curl -sSL "https://raw.githubusercontent.com/unslothai/unsloth/80e0108a684c882965a02a8ed851e3473c1145ab/unsloth/models/vision.py" \
         -o "$VISION_DST"
     echo "   Installing studio dependencies..."
@@ -193,6 +193,10 @@ else
     LLAMA_CPP_DST="$(pip show unsloth-zoo | grep -i '^Location:' | awk '{print $2}')/unsloth_zoo/llama_cpp.py"
     curl -sSL "https://raw.githubusercontent.com/unslothai/unsloth-zoo/refs/heads/main/unsloth_zoo/llama_cpp.py" \
         -o "$LLAMA_CPP_DST"
+    # Patch: override vision.py with fix from unsloth PR: https://github.com/unslothai/unsloth/pull/4091 until next pypi release
+    VISION_DST="$(pip show unsloth | grep -i '^Location:' | awk '{print $2}')/unsloth/models/vision.py"
+    curl -sSL "https://raw.githubusercontent.com/unslothai/unsloth/80e0108a684c882965a02a8ed851e3473c1145ab/unsloth/models/vision.py" \
+        -o "$VISION_DST"
     echo "   Installing studio dependencies..."
     run_quiet "pip install studio" pip install -r "$SCRIPT_DIR/studio/backend/requirements/studio.txt"
     echo "✅ Python dependencies installed"