FROM pytorch/pytorch:2.10.0-cuda13.0-cudnn9-runtime

# ------------------------------------------------
# System packages
# ------------------------------------------------
RUN apt-get update && \
    apt-get install -y --no-install-recommends \
    git \
    wget \
    python3-venv \
    python3-setuptools \
    libjemalloc2 && \
    apt-get clean && \
    rm -rf /var/lib/apt/lists/*

RUN useradd -rm -d /home/hf -s /bin/bash -g root -G sudo -u 1001 hf

# ------------------------------------------------
# Environment
# ------------------------------------------------
ENV TRANSFORMERS_CACHE=/tmp/hf_cache
ENV HF_HOME=/tmp/hf_cache

ENV PYTHONUNBUFFERED=1
ENV PYTHONDONTWRITEBYTECODE=1
ENV TOKENIZERS_PARALLELISM=false

# CPU thread optimization
ENV KMP_BLOCKTIME=0
ENV KMP_AFFINITY=granularity=fine,compact,1,0
ENV OMP_WAIT_POLICY=ACTIVE
ENV ONEDNN_PRIMITIVE_CACHE_CAPACITY=1024

# jemalloc as allocator — reduces tensor fragmentation
ENV LD_PRELOAD=/usr/lib/x86_64-linux-gnu/libjemalloc.so.2
ENV PYTHONMALLOC=malloc

WORKDIR /home/hf

# ------------------------------------------------
# entrypoint (sets CPU count + jemalloc at runtime)
# ------------------------------------------------
COPY --chown=root:root entrypoint.sh /usr/local/bin/entrypoint.sh
RUN chmod +x /usr/local/bin/entrypoint.sh

USER hf

# ------------------------------------------------
# ComfyUI
# ------------------------------------------------
RUN git clone https://github.com/comfyanonymous/ComfyUI
WORKDIR /home/hf/ComfyUI

ENV VIRTUAL_ENV=/home/hf/ComfyUI/venv
RUN python3 -m venv $VIRTUAL_ENV
ENV PATH="$VIRTUAL_ENV/bin:$PATH"

RUN pip install --upgrade pip
RUN pip install --no-cache-dir --prefer-binary -r requirements.txt

# ------------------------------------------------
# ComfyUI-GGUF custom node (enables GGUF model loading)
# ------------------------------------------------
RUN git clone https://github.com/city96/ComfyUI-GGUF \
    ./custom_nodes/ComfyUI-GGUF && \
    pip install --no-cache-dir -r ./custom_nodes/ComfyUI-GGUF/requirements.txt

# ------------------------------------------------
# Anima Preview3 Q5_K_M GGUF (~1.59 GB vs 4.18 GB BF16)
# ------------------------------------------------
RUN wget -nv -c \
    https://huggingface.co/Bedovyy/Anima-GGUF/resolve/main/anima-preview3-base-Q5_K_M.gguf \
    -P ./models/diffusion_models/

RUN wget -nv -c \
    https://huggingface.co/circlestone-labs/Anima/resolve/main/split_files/text_encoders/qwen_3_06b_base.safetensors \
    -P ./models/text_encoders/

RUN wget -nv -c \
    https://huggingface.co/circlestone-labs/Anima/resolve/main/split_files/vae/qwen_image_vae.safetensors \
    -P ./models/vae/

# Turbo LoRA (14-step distilled — much faster on CPU)
RUN wget -nv -c \
    https://huggingface.co/Einhorn/Anima-Preview2-Turbo-LoRA/resolve/main/anima_turbo_distill_preview2_14step.safetensors \
    -P ./models/loras/

RUN mkdir -p temp && chmod 777 temp

# ------------------------------------------------
# Launch ComfyUI (CPU optimized)
# ComfyUI's web UI is served on port 7860 (HF Spaces default)
# The built-in REST API is available on the same port:
#   POST /prompt      — submit a generation job
#   GET  /history     — poll job status
#   GET  /view        — fetch output image
# ------------------------------------------------
ENTRYPOINT ["/usr/local/bin/entrypoint.sh"]
CMD ["python", "main.py", \
    "--cpu", \
    "--listen", "0.0.0.0", \
    "--port", "7860", \
    "--output-directory", "./temp", \
    "--bf16-vae", \
    "--bf16-text-enc", \
    "--cache-none", \
    "--use-pytorch-cross-attention", \
    "--preview-method", "none"]