#!/bin/bash
# sd-webui - start sd-server with a model picked from rofi
# Copyright (C) 2026 Danilo M. <danix@danix.xyz>
#
# This program is free software; you can redistribute it and/or modify
# it under the terms of the GNU General Public License version 2 as
# published by the Free Software Foundation.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License along
# with this program; if not, write to the Free Software Foundation, Inc.,
# 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA.

# Start sd-server (stable-diffusion.cpp) with a model picked from rofi and open its web UI.
# Runs in the foreground: close the terminal (or Ctrl-C) to stop it and free VRAM.

M=/data/LLM-models/SD
QWEN=(--llm $M/Qwen3VL-8B-Instruct-Q8_0.gguf
      --llm_vision $M/mmproj-Qwen3VL-8B-Instruct-F16.gguf
      --vae $M/vae/qwen_image_2.1_vae_bf16.safetensors)
SDXL=(--cfg-scale 7.0 --steps 25)   # same defaults as imggen

# Pick the model (sd-server loads one per process, so choose at start).
case $(printf '%s\n' \
         "Qwen Image 2.1 (base, 20 steps, text/edit/RGBA)" \
         "Qwen Image 2.1 turbo (6 steps)" \
         "Z-Image Turbo (8 steps, photoreal)" \
         "FLUX.2 klein 4B (4 steps, edit)" \
         "Krea 2 Turbo (8 steps, photoreal)" \
         "Wan2.2 TI2V 5B (video, text/image)" \
         "SDXL base" \
         "RealVisXL V5" \
         "Pony Diffusion V6 XL" \
       | rofi -dmenu -i -theme ~/.config/rofi/udt/list.rasi -p '󰋩  SD model:' \
           -theme-str 'entry { placeholder: "Search models"; }') in
  "Qwen Image 2.1 (base"*) MODEL=(--diffusion-model $M/qwen_image_2.1-Q8_0.gguf "${QWEN[@]}"
                                 --cfg-scale 6.0 --cache-mode cache-dit) ;;
  "Qwen Image 2.1 turbo"*) MODEL=(--diffusion-model $M/Qwen-Image-2.1-viggle-turbo-v0.3-6step-Q8_0.gguf "${QWEN[@]}"
                                 --cfg-scale 1.0 --steps 6) ;;
  Z-Image*)  MODEL=(--diffusion-model $M/z_image_turbo-Q8_0.gguf
                    --llm $M/Qwen3-4B-Instruct-2507-Q8_0.gguf --vae $M/vae/flux1-ae.safetensors
                    --cfg-scale 1.0 --steps 8 --diffusion-fa) ;;
  FLUX.2*)   MODEL=(--diffusion-model $M/flux-2-klein-4b-Q8_0.gguf
                    --llm $M/Qwen3-4B-Q8_0.gguf --vae $M/vae/flux2-vae.safetensors
                    --cfg-scale 1.0 --steps 4 --diffusion-fa) ;;
  # Q6_K: same output as Q8_0 at seed 42, ~20% faster (both still offload on 12G).
  Krea*)     MODEL=(--diffusion-model $M/Krea-2-Turbo-Q6_K.gguf
                    --llm $M/Qwen3VL-4B-Instruct-Q8_0.gguf --vae $M/vae/wan_2.1_vae.safetensors
                    --cfg-scale 1.0 --steps 8 --diffusion-fa) ;;
  # 1280x704 is the model's native 720p. Frames must be 4k+1 at 24 fps: 49 = 2 s, 121 = 5 s.
  # Settings follow upstream Wan2.2 5B defaults (cfg 5, shift 5); drop to 832x480 if too slow.
  Wan2.2*)   MODEL=(--diffusion-model $M/Wan2.2-TI2V-5B-Q8_0.gguf
                    --t5xxl $M/umt5-xxl-encoder-Q8_0.gguf --vae $M/vae/wan2.2_vae.safetensors
                    --cfg-scale 5.0 --steps 30 --flow-shift 5.0 --diffusion-fa
                    -W 1280 -H 704 --video-frames 49 --temporal-tiling) ;;
  # imggen's diffusers weights, symlinked into the layout sd.cpp loads as a directory
  "SDXL base")  MODEL=(-m $M/sdxl-links/sdxl "${SDXL[@]}") ;;
  RealVisXL*)   MODEL=(-m $M/sdxl-links/realvis "${SDXL[@]}") ;;
  Pony*)        MODEL=(-m $M/sdxl-links/pony "${SDXL[@]}" --clip-skip 2) ;;
  *) exit 0 ;;
esac

PORT=${PORT:-7860}
URL=http://127.0.0.1:$PORT/

# Open the browser once the server answers (model loading takes a while).
(
  until curl -sf -o /dev/null "$URL"; do sleep 1; done
  notify-send -a sd-server -i applications-graphics "Stable Diffusion server started" "$URL"
  xdg-open "$URL" >/dev/null 2>&1
) &
trap 'kill $! 2>/dev/null' EXIT

# Log to a file too: the terminal loses the output when the server crashes.
# stdbuf -oL: line-buffer stdout so the last lines before a crash reach the log.
LOG=${XDG_CACHE_HOME:-$HOME/.cache}/sd-server.log
# MODEL goes after the shared defaults so a model can override them (e.g. Wan's -W/-H).
stdbuf -oL -eL sd-server --listen-port "$PORT" \
  --lora-model-dir $M/loras \
  --sampling-method euler -W 1024 -H 1024 --seed -1 \
  "${MODEL[@]}" \
  --offload-to-cpu -v "$@" 2>&1 | tee "$LOG"