FROM pytorch/pytorch:2.2.2-cuda12.1-cudnn8-runtime

ENV DEBIAN_FRONTEND=noninteractive \
    PIP_DISABLE_PIP_VERSION_CHECK=1 \
    PYTHONUNBUFFERED=1 \
    PATH="/home/user/.local/bin:$PATH" \
    COMMANDLINE_ARGS="--listen --port 7860 --xformers --enable-insecure-extension-access --api --no-half-vae --embeddings-dir /app/stable-diffusion-webui/embeddings" \
    VENV_DIR="/opt/venv" \
    WEBUI_DIR="/app/stable-diffusion-webui" \
    MODELS_DIR="/app/models" \
    HF_HOME="/app/.cache/huggingface"

RUN useradd -m -u 1000 user
USER user
WORKDIR /app

USER root
RUN apt-get update && apt-get install -y --no-install-recommends \
    git git-lfs wget aria2 ffmpeg libgl1 libglib2.0-0 \
  && rm -rf /var/lib/apt/lists/* \
  && git lfs install
USER user

RUN git clone https://github.com/AUTOMATIC1111/stable-diffusion-webui.git ${WEBUI_DIR}

# Optional extensions
# RUN git clone https://github.com/Mikubill/sd-webui-controlnet ${WEBUI_DIR}/extensions/sd-webui-controlnet
RUN git clone https://github.com/bmaltais/kohya_ss ${WEBUI_DIR}/extensions/kohya_ss || true

# Create embeddings directory with correct permissions
USER root
RUN mkdir -p ${WEBUI_DIR}/embeddings \
    /app/log/ti \
    /app/datasets \
    && chown -R user:user /app
USER user

RUN pip install --no-cache-dir --upgrade pip wheel setuptools \
  && pip install --no-cache-dir xformers==0.0.24 huggingface_hub

RUN mkdir -p ${MODELS_DIR}/Stable-diffusion ${MODELS_DIR}/Lora ${MODELS_DIR}/VAE

COPY --chown=user:users ./start.sh /app/start.sh
RUN chmod +x /app/start.sh

EXPOSE 7860
WORKDIR ${WEBUI_DIR}
CMD ["/app/start.sh"]