#!/bin/bash
# Setup script for AUTOMATIC1111 Stable Diffusion WebUI
# Target: Ubuntu 22.04 LTS + NVIDIA RTX 2000 Ada (16GB VRAM)

set -e

SD_DIR="$HOME/stable-diffusion-webui"

echo "=== AUTOMATIC1111 Stable Diffusion WebUI Setup ==="
echo "=== Target: Ubuntu 22.04 LTS + NVIDIA RTX 2000 Ada ==="
echo ""

# 1. Install system dependencies that Ubuntu 22.04 doesn't ship by default
echo "Installing system dependencies..."
sudo apt-get update
sudo apt-get install -y \
    git \
    wget \
    python3-venv \
    python3-pip \
    libgl1 \
    libglib2.0-0 \
    libsm6 \
    libxrender1 \
    libxext6 \
    libffi-dev \
    libssl-dev

# 2. Check NVIDIA driver
if ! command -v nvidia-smi &>/dev/null; then
    echo ""
    echo "ERROR: nvidia-smi not found. Install NVIDIA drivers first:"
    echo ""
    echo "  sudo apt install nvidia-driver-535"
    echo "  sudo reboot"
    echo ""
    echo "Or use Ubuntu's driver manager:"
    echo "  sudo ubuntu-drivers autoinstall"
    echo "  sudo reboot"
    exit 1
fi
echo ""
echo "GPU detected:"
nvidia-smi --query-gpu=name,driver_version,memory.total --format=csv,noheader
echo ""

# 3. Check CUDA availability
if ! command -v nvcc &>/dev/null; then
    echo "NOTE: nvcc not found — A1111 will download its own CUDA via PyTorch, which is fine."
    echo "      (No need to install CUDA toolkit separately.)"
    echo ""
fi

# 4. Clone A1111
if [ -d "$SD_DIR" ]; then
    echo "Directory $SD_DIR already exists, pulling latest..."
    cd "$SD_DIR" && git pull
else
    echo "Cloning AUTOMATIC1111..."
    git clone https://github.com/AUTOMATIC1111/stable-diffusion-webui.git "$SD_DIR"
fi
cd "$SD_DIR"

# 5. Download a good default model (SD 1.5 — fast, 16GB friendly)
# You can swap this for SDXL later if you want higher quality
MODEL_DIR="$SD_DIR/models/Stable-diffusion"
mkdir -p "$MODEL_DIR"

if [ ! -f "$MODEL_DIR/v1-5-pruned-emaonly.safetensors" ]; then
    echo "Downloading Stable Diffusion 1.5 model (~4GB)..."
    wget -q --show-progress -O "$MODEL_DIR/v1-5-pruned-emaonly.safetensors" \
        "https://huggingface.co/stable-diffusion-v1-5/stable-diffusion-v1-5/resolve/main/v1-5-pruned-emaonly.safetensors"
else
    echo "SD 1.5 model already downloaded."
fi

echo ""
echo "=== Setup complete ==="
echo ""
echo "To start the WebUI with API enabled, run:"
echo "  cd $SD_DIR"
echo "  ./webui.sh --api --listen --xformers --no-half-vae"
echo ""
echo "First launch will take several minutes (installs PyTorch, xformers, etc.)"
echo ""
echo "The API will be available at http://localhost:7860"
echo "Test it with: curl http://localhost:7860/sdapi/v1/sd-models"
echo ""
echo "Optional: To run as a systemd service, run:"
echo "  chmod +x setup-sd-service.sh && sudo ./setup-sd-service.sh"