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

# ------------------------------------------------
# System packages
# ------------------------------------------------
# [修正1] python3-distutils は Ubuntu 24.04 / Python 3.12 で廃止済み。
#         python3-setuptools で代替。
# [最適化] libjemalloc2: 大テンソルの確保/解放時の断片化を抑制。
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

# ------------------------------------------------
# Global performance / memory optimization
# ------------------------------------------------

# [修正2] PYTORCH_ENABLE_MPS_FALLBACK は Apple Silicon 専用。Linux では削除。
# [修正3] HF_HOME/TRANSFORMERS_CACHE を /tmp 直下 → /tmp/hf_cache に変更。
ENV TRANSFORMERS_CACHE=/tmp/hf_cache
ENV HF_HOME=/tmp/hf_cache

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

# ------------------------------------------------
# CPU 最適化環境変数 (静的に設定できるもの)
# ------------------------------------------------

# KMP_BLOCKTIME=0: アイドルスレッドが即スリープ → 推論ループ間のウェイクアップ遅延ゼロ
ENV KMP_BLOCKTIME=0

# KMP_AFFINITY: スレッドを物理コアに密集配置し HT を活用 (スレッド数固定なし)
ENV KMP_AFFINITY=granularity=fine,compact,1,0

# OMP_WAIT_POLICY=ACTIVE: OpenMP スレッドをビジーウェイト → 短タスク間の再起動コスト削減
ENV OMP_WAIT_POLICY=ACTIVE

# oneDNN プリミティブキャッシュ拡大 → 同一 shape 演算の再コンパイルを削減
ENV ONEDNN_PRIMITIVE_CACHE_CAPACITY=1024

# jemalloc をアロケータとして使用 → 断片化抑制
ENV LD_PRELOAD=/usr/lib/x86_64-linux-gnu/libjemalloc.so.2

# [修正6] Python 内部アロケータ (pymalloc) をバイパスし、小オブジェクトも
#         jemalloc に直結する。pymalloc はメモリを自身のヒープに溜め込み
#         LD_PRELOAD 経由の jemalloc に戻さないため、RSS が肥大化する。
ENV PYTHONMALLOC=malloc

# ------------------------------------------------
# スレッド数は entrypoint.sh で nproc を使って動的設定。
# ENV に静的に書くと Docker ビルド時の値に固定されてしまい、
# 実行環境の CPU 数と一致しない場合に 50% 止まりになる。
# ------------------------------------------------

WORKDIR /home/hf

# ------------------------------------------------
# entrypoint.sh を配置
# ------------------------------------------------
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

# ------------------------------------------------
# ERNIE Image models
# ------------------------------------------------
RUN wget -nv -c \
https://huggingface.co/Comfy-Org/ERNIE-Image/resolve/main/diffusion_models/ernie-image-turbo.safetensors \
-P ./models/diffusion_models/

RUN wget -nv -c \
https://huggingface.co/Comfy-Org/ERNIE-Image/resolve/main/diffusion_models/ernie-image.safetensors \
-P ./models/diffusion_models/

RUN wget -nv -c \
https://huggingface.co/Comfy-Org/ERNIE-Image/resolve/main/text_encoders/ernie-image-prompt-enhancer.safetensors \
-P ./models/text_encoders/

RUN wget -nv -c \
https://huggingface.co/Comfy-Org/ERNIE-Image/resolve/main/text_encoders/ministral-3-3b.safetensors \
-P ./models/text_encoders/

RUN wget -nv -c \
https://huggingface.co/Comfy-Org/ERNIE-Image/resolve/main/vae/flux2-vae.safetensors \
-P ./models/vae/

RUN mkdir temp && chmod 777 temp

# ------------------------------------------------
# AMP (自動混合精度) 設定
# ------------------------------------------------
# torch.amp.autocast("cpu") は CPU AMP API であり、
# コードレベルの context manager のため Dockerfile から直接制御できない。
# CPU AMP でサポートされる低精度型は BF16 のみ (FP16 は CPU では非対応)。
# ComfyUI の --bf16-* フラグは内部で torch.amp.autocast("cpu", dtype=bfloat16)
# 相当の処理を有効化する。これが Dockerfile レベルで設定できる唯一の手段。
#
# --force-fp32 はアキュムレーション (累積演算) を FP32 に固定し、
# BF16 演算の精度劣化を最小限に抑える (AMP の標準パターン)。

# ------------------------------------------------
# Launch ComfyUI (CPU optimized)
# ------------------------------------------------
# [修正4] --lowvram: GPU VRAM 管理フラグ。CPU モードで削除。
# [修正5] --use-split-cross-attention: GPU 専用。
#         --use-pytorch-cross-attention (CPU 対応) に置換。
# [最適化] --preview-method none: CPU での VAE デコードを省いて本推論に集中。
ENTRYPOINT ["/usr/local/bin/entrypoint.sh"]
CMD ["python", "main.py", \
     "--cpu", \
     "--listen", "0.0.0.0", \
     "--port", "7860", \
     "--output-directory", "./temp", \
     "--bf16-unet", \
     "--bf16-vae", \
     "--bf16-text-enc", \
     "--cache-none", \
     "--use-pytorch-cross-attention", \
     "--preview-method", "none"]