# Stable Diffusion WebUI with AMD ROCm Support
# Base image: Ubuntu 22.04 for best ROCm compatibility

FROM ubuntu:22.04 AS builder

# Set environment variables
ENV DEBIAN_FRONTEND=noninteractive \
    PYTHON_VERSION=3.10 \
    ROCM_VERSION=5.7 \
    HSA_OVERRIDE_GFX_VERSION=11.0.0 \
    HIP_VISIBLE_DEVICES=0 \
    HIP_PATH=/opt/rocm \
    PATH=/opt/rocm/bin:$PATH \
    LD_LIBRARY_PATH=/opt/rocm/lib:/opt/rocm/hip/lib:/opt/rocm/opencl/lib/x86_64 \
    PYTORCH_ROCM_ARCH=gfx1100

# Update and install system dependencies
RUN apt-get update && apt-get upgrade -y && \
    apt-get install -y --no-install-recommends \
    wget \
    git \
    curl \
    ca-certificates \
    gnupg \
    lsb-release \
    software-properties-common \
    build-essential \
    python3 \
    python3-pip \
    python3-venv \
    libgl1 \
    libglib2.0-0 \
    libsm6 \
    libxrender1 \
    libxext6 \
    libfontconfig1 \
    libfreetype6 \
    libgl1-mesa-glx \
    libglib2.0-0 \
    libsm6 \
    libxrender1 \
    libxext6 \
    libfontconfig1 \
    libfreetype6 \
    && rm -rf /var/lib/apt/lists/*

# Install ROCm 5.7
RUN mkdir -p /tmp/rocm && cd /tmp/rocm && \
    wget https://repo.radeon.com/amdgpu-install/5.7/ubuntu/jammy/amdgpu-install_5.7.50700-1_all.deb && \
    apt-get install -y ./amdgpu-install_5.7.50700-1_all.deb && \
    apt-get update && \
    DEBIAN_FRONTEND=noninteractive apt-get install -y --no-install-recommends \
    rocm-device-libs rocm-hip-runtime rocm-opencl-runtime rocminfo && \
    rm -rf /tmp/rocm

# Verify ROCm installation
RUN rocminfo | grep "Agent" || echo "ROCm installation verification"

# Install Python 3.10 and create virtual environment
RUN apt-get update && apt-get install -y --no-install-recommends \
    python3.10 \
    python3.10-dev \
    python3.10-venv \
    && rm -rf /var/lib/apt/lists/*

# Create and activate virtual environment
RUN python3.10 -m venv /opt/venv
ENV PATH=/opt/venv/bin:$PATH

# Install PyTorch with ROCm support
RUN pip install --upgrade pip && \
    pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/rocm5.7

# Clone the upstream repository
RUN git clone https://github.com/AUTOMATIC1111/stable-diffusion-webui.git /app/webui && \
    cd /app/webui && \
    git checkout v1.7.0

# Install project dependencies
WORKDIR /app
RUN pip install -r /app/webui/requirements.txt

# Install xformers with ROCm support or handle gracefully
RUN pip install xformers --index-url https://download.pytorch.org/whl/rocm5.7 || \
    (echo "xformers with ROCm support failed, installing without ROCm acceleration" && \
     pip install xformers --no-deps && \
     echo "xformers installed but may not have full ROCm acceleration")

# Install additional dependencies that might be needed
RUN apt-get update && apt-get install -y --no-install-recommends \
    libgl1-mesa-glx \
    libglib2.0-0 \
    libsm6 \
    libxrender1 \
    libxext6 \
    libfontconfig1 \
    libfreetype6 \
    && rm -rf /var/lib/apt/lists/*

# Set working directory and expose port
WORKDIR /app/webui
EXPOSE 7860

# Create entrypoint script
COPY docker/docker-entrypoint.sh /docker-entrypoint.sh
RUN chmod +x /docker-entrypoint.sh

ENTRYPOINT ["/docker-entrypoint.sh"]
CMD ["python", "launch.py", "--listen", "--port", "7860", "--opt-sdp-attention", "--no-half-vae"]