From 8da1780f6e468688c40aa970f57416921514cd65 Mon Sep 17 00:00:00 2001
From: Paul Bailey <psbailey31@gmail.com>
Date: Sun, 22 Feb 2026 23:58:58 +0000
Subject: [PATCH] Update installer to download DreamShaper 8 model and sync
 tool checkpoint name

---
 install.sh                    | 14 ++++++++++++++
 src/ollama_cli/tools/media.py |  2 +-
 2 files changed, 15 insertions(+), 1 deletion(-)

diff --git a/install.sh b/install.sh
index bcb3c7a..2e33b9e 100644
--- a/install.sh
+++ b/install.sh
@@ -79,6 +79,20 @@ if [[ $REPLY =~ ^[Yy]$ ]]; then
         pip3 install -r requirements.txt
         echo "✅ ComfyUI installed."
         
+        # 5. Download Image Generation Model
+        echo ""
+        echo "[*] Downloading Image Generation Model (DreamShaper 8)..."
+        MODEL_DIR="models/checkpoints"
+        mkdir -p "$MODEL_DIR"
+        MODEL_PATH="$MODEL_DIR/dreamshaper_8.safetensors"
+        if [ ! -f "$MODEL_PATH" ]; then
+            echo "    This may take a few minutes (approx. 2GB)..."
+            curl -L "https://huggingface.co/Lykon/DreamShaper/resolve/main/DreamShaper_8_pruned.safetensors" -o "$MODEL_PATH"
+            echo "✅ Model downloaded."
+        else
+            echo "✅ Model already exists."
+        fi
+
         echo ""
         echo "[!] To use Image Generation, you must start ComfyUI:"
         echo "    cd $INSTALL_DIR"
diff --git a/src/ollama_cli/tools/media.py b/src/ollama_cli/tools/media.py
index 6f1563d..e3fd2f4 100644
--- a/src/ollama_cli/tools/media.py
+++ b/src/ollama_cli/tools/media.py
@@ -30,7 +30,7 @@ def generate_image(prompt: str) -> str:
         # Based on the user's previous generate_image.py, but using the configured URL
         workflow = {
             "1": {
-                "inputs": {"ckpt_name": "dreamshaper_8.safetensors"},
+                "inputs": {"ckpt_name": "dreamshaper_8_pruned.safetensors"},
                 "class_type": "CheckpointLoaderSimple"
             },
             "2": {