FROM nvidia/cuda:11.8.0-cudnn8-runtime-ubuntu22.04

ENV DEBIAN_FRONTEND=noninteractive \
    TZ=America/Los_Angeles

ARG USE_PERSISTENT_DATA

RUN apt-get update && apt-get install -y \
    git \
    make build-essential libssl-dev zlib1g-dev \
    libbz2-dev libreadline-dev libsqlite3-dev wget curl llvm \
    libncursesw5-dev xz-utils tk-dev libxml2-dev libxmlsec1-dev libffi-dev liblzma-dev git-lfs  \
    ffmpeg libsm6 libxext6 cmake libgl1-mesa-glx \
    && rm -rf /var/lib/apt/lists/* \
    && git lfs install

WORKDIR /code

COPY ./requirements.txt /code/requirements.txt

# User
RUN useradd -m -u 1000 user
USER user
ENV HOME=/home/user \
    PATH=/home/user/.local/bin:$PATH

# Pyenv
RUN curl https://pyenv.run | bash
ENV PATH=$HOME/.pyenv/shims:$HOME/.pyenv/bin:$PATH

ARG PYTHON_VERSION=3.9.17
# Python và các thư viện
RUN pyenv install $PYTHON_VERSION && \
    pyenv global $PYTHON_VERSION && \
    pyenv rehash && \
    pip install --no-cache-dir --upgrade pip setuptools wheel && \
    pip install --no-cache-dir --extra-index-url https://download.pytorch.org/whl/cu118 -r /code/requirements.txt && \
    pip install --no-cache-dir datasets huggingface-hub "protobuf<4" "click<8.1"

WORKDIR $HOME/app

# Clone ComfyUI
RUN git clone https://github.com/comfyanonymous/ComfyUI.git . && \
    pip install --no-cache-dir -r requirements.txt

# Tải WAN 2.1
RUN python -c "from huggingface_hub import hf_hub_download; \
    hf_hub_download(repo_id='Comfy-Org/Wan_2.1_ComfyUI_repackaged', filename='wan2.1_i2v_720p_14B_fp8_e4m3fn.safetensors', subfolder='split_files/diffusion_models', local_dir='./models/diffusion_models'); \
    hf_hub_download(repo_id='Comfy-Org/Wan_2.1_ComfyUI_repackaged', filename='umt5_xxl_fp8_e4m3fn_scaled.safetensors', subfolder='split_files/text_encoders', local_dir='./models/text_encoders'); \
    hf_hub_download(repo_id='Comfy-Org/Wan_2.1_ComfyUI_repackaged', filename='wan_2.1_vae.safetensors', subfolder='split_files/vae', local_dir='./models/vae'); \
    hf_hub_download(repo_id='Comfy-Org/Wan_2.1_ComfyUI_repackaged', filename='clip_vision_h.safetensors', subfolder='split_files/clip_vision', local_dir='./models/clip_vision'); \
    hf_hub_download(repo_id='Comfy-Org/Wan_2.1_ComfyUI_repackaged', filename='image_to_video_wan_720p_example.json', subfolder='example%20workflows_Wan2.1', local_dir='./user/default/workflows')"

# Tải workflow từ foxtonai
RUN mkdir -p ./user/default/workflows && \
    wget -P ./user/default/workflows https://foxtonai.com/downloads/comfyui-wan-2.1-i2v-modified-workflow.json
CMD ["python", "main.py", "--listen", "0.0.0.0", "--cpu", "--port", "7860", "--use-split-cross-attention", "--output-directory", "${USE_PERSISTENT_DATA:+/data/}"]