# ---------- base image -------------------------------------------------------
FROM nvidia/cuda:12.1.1-cudnn8-runtime-ubuntu22.04
ENV DEBIAN_FRONTEND=noninteractive \
    PYTHONUNBUFFERED=1 \
    HF_HUB_DISABLE_TELEMETRY=1 \
    PIP_NO_CACHE_DIR=1

# ---------- system deps ------------------------------------------------------
RUN apt-get update && \
    apt-get install -y --no-install-recommends \
    git git-lfs ffmpeg wget libgl1 libglib2.0-0 libsm6 libxext6 libxrender1 build-essential cmake python3-dev && \
    rm -rf /var/lib/apt/lists/*

# ---------- python & cuda ----------------------------------------------------
RUN python3 -m pip install --no-cache-dir --upgrade pip

# CUDA 12 wheels (torch 2.2+ ships them)
RUN pip install --no-cache-dir torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu121

# xformers & deepspeed for fp16 + multi-GPU
RUN pip install --no-cache-dir xformers==0.0.26 deepspeed==0.14.2

# ---------- project code -----------------------------------------------------
WORKDIR /workspace

# Install requirements first
COPY requirements.txt .
RUN pip install --no-cache-dir -r requirements.txt

# Clone SadTalker and install its requirements
RUN git lfs install --system && \
    git clone https://github.com/OpenTalker/SadTalker.git && \
    cd SadTalker && \
    git lfs pull && \
    pip install --no-cache-dir -r requirements.txt

# Install additional dependencies for SadTalker
RUN pip install --no-cache-dir \
    yacs==0.1.8 \
    librosa==0.8.0 \
    numba \
    resampy \
    pydub \
    scipy \
    kornia==0.6.8 \
    face_alignment \
    imageio \
    imageio-ffmpeg \
    av \
    dlib

# ---------- Real-ESRGAN install from source ----------------------------------
RUN git clone https://github.com/xinntao/Real-ESRGAN.git && \
    cd Real-ESRGAN && \
    pip install basicsr==1.4.2 facexlib==0.2.5 gfpgan==1.3.8 && \
    pip install -r requirements.txt && \
    python setup.py develop

# ---------- Download SadTalker models manually --------------------------------
RUN mkdir -p /workspace/SadTalker/checkpoints && \
    cd /workspace/SadTalker/checkpoints && \
    wget -q https://github.com/OpenTalker/SadTalker/releases/download/v0.0.2-rc/SadTalker_V0.0.2_256.safetensors && \
    wget -q https://github.com/OpenTalker/SadTalker/releases/download/v0.0.2-rc/SadTalker_V0.0.2_512.safetensors && \
    wget -q https://github.com/OpenTalker/SadTalker/releases/download/v0.0.2-rc/mapping_00109-model.pth.tar && \
    wget -q https://github.com/OpenTalker/SadTalker/releases/download/v0.0.2-rc/mapping_00229-model.pth.tar && \
    wget -q https://github.com/OpenTalker/SadTalker/releases/download/v0.0.2-rc/BFM_Fitting.zip && \
    unzip -q BFM_Fitting.zip && rm BFM_Fitting.zip

# ---------- app --------------------------------------------------------------
# Set working directory to SadTalker
WORKDIR /workspace/SadTalker

# Add SadTalker to Python path
ENV PYTHONPATH="${PYTHONPATH}:/workspace/SadTalker:/workspace/SadTalker/src"

# Copy our app to the SadTalker directory
COPY app.py /workspace/SadTalker/

EXPOSE 7860
CMD ["python", "app.py"]