#!/bin/bash

# Docker entrypoint script for Stable Diffusion WebUI with AMD ROCm support

set -e

# Set up environment variables for ROCm
export HSA_OVERRIDE_GFX_VERSION=${HSA_OVERRIDE_GFX_VERSION:-11.0.0}
export HIP_VISIBLE_DEVICES=${HIP_VISIBLE_DEVICES:-0}
export PYTORCH_ROCM_ARCH=${PYTORCH_ROCM_ARCH:-gfx1100}
export HIP_PATH=/opt/rocm
export PATH=/opt/rocm/bin:$PATH
export LD_LIBRARY_PATH=/opt/rocm/lib:/opt/rocm/hip/lib:/opt/rocm/opencl/lib/x86_64:$LD_LIBRARY_PATH

echo "=========================================="
echo "Stable Diffusion WebUI Docker Entrypoint"
echo "=========================================="
echo "ROCm Version: 5.7"
echo "GPU Target: gfx1100 (${HSA_OVERRIDE_GFX_VERSION})"
echo "PyTorch: ROCm 5.7 compatible"
echo "=========================================="

# Check if ROCm is properly installed
if command -v rocminfo &> /dev/null; then
    echo "ROCm is installed:"
    rocminfo | grep -E "Agent|Name" | head -5
else
    echo "Warning: ROCm not found in PATH"
fi

# Check if GPU devices are accessible
if [ -c /dev/kfd ]; then
    echo "AMD KFD device found"
else
    echo "Warning: /dev/kfd not found"
fi

if [ -c /dev/dri/card0 ]; then
    echo "DRM device found"
else
    echo "Warning: /dev/dri/card0 not found"
fi

# Activate virtual environment
if [ -d "/opt/venv" ]; then
    echo "Activating Python virtual environment..."
    source /opt/venv/bin/activate
else
    echo "Warning: Virtual environment not found at /opt/venv"
fi

# Set default command line arguments if not provided
DEFAULT_ARGS="--api --listen --port 7860 --opt-sdp-attention --no-half-vae"
CLI_ARGS=${CLI_ARGS:-$DEFAULT_ARGS}

echo "Starting Stable Diffusion WebUI with arguments: $CLI_ARGS"
echo "Web UI will be available at http://localhost:7860"
echo "=========================================="

# Execute the command
exec python launch.py $CLI_ARGS