FROM python:3.11-slim

# ------------------------------------------------
# 1. システムパッケージ & jemalloc の導入
# ------------------------------------------------
RUN apt-get update && \
    apt-get install -y --no-install-recommends \
        git \
        wget \
        gcc \
        g++ \
        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

# ------------------------------------------------
# 2. 基本環境変数 & メモリアロケータ最適化
# ------------------------------------------------
ENV TRANSFORMERS_CACHE=/tmp/hf_cache
ENV HF_HOME=/tmp/hf_cache
ENV PYTHONUNBUFFERED=1
ENV PYTHONDONTWRITEBYTECODE=1

ENV LD_PRELOAD=/usr/lib/x86_64-linux-gnu/libjemalloc.so.2
ENV MALLOC_CONF="background_thread:true,metadata_thp:auto,dirty_decay_ms:20000,muzzy_decay_ms:20000"

# 💡 チューニング：全力でCPUを回すため WAIT_POLICY をデフォルト（ACTIVE側）に戻します
ENV KMP_BLOCKTIME=1
ENV DNNL_DEFAULT_FPMATH_MODE=BF16
ENV ONEDNN_PRIMITIVE_CACHE_CAPACITY=512

WORKDIR /home/hf
USER hf

# ------------------------------------------------
# 3. ComfyUI の展開と、デグレードを起こさない pip インストール
# ------------------------------------------------
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"

# 💡 修正：最初にCPU専用リポジトリから torch 関連をピンポイントで導入。
# その後、通常のPyPIインデックスから残りの依存関係（requirements.txt）を安全に解決します。
RUN pip install --no-cache-dir --upgrade pip && \
    pip install --no-cache-dir torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cpu && \
    pip install --no-cache-dir -r requirements.txt

# ------------------------------------------------
# 4. モデルのダウンロード
# ------------------------------------------------
RUN wget -nv -c https://huggingface.co/circlestone-labs/Anima/resolve/main/split_files/diffusion_models/anima-base-v1.0.safetensors -P ./models/diffusion_models/ && \
    wget -nv -c https://huggingface.co/circlestone-labs/Anima/resolve/main/split_files/text_encoders/qwen_3_06b_base.safetensors -P ./models/text_encoders/ && \
    wget -nv -c https://huggingface.co/circlestone-labs/Anima/resolve/main/split_files/vae/qwen_image_vae.safetensors -P ./models/vae/

RUN mkdir temp && chmod 777 temp

# ------------------------------------------------
# 5. 起動コマンド（スレッド緊縛・CPU全力駆動版）
# ------------------------------------------------
CMD quota_v2="/sys/fs/cgroup/cpu.max"; \
    quota_v1="/sys/fs/cgroup/cpu/cpu.cfs_quota_us"; \
    period_v1="/sys/fs/cgroup/cpu/cpu.cfs_period_us"; \
    COMPUTE_CORES=""; \
    if [ -f "$quota_v2" ]; then \
        read -r q p < "$quota_v2"; \
        if [ "$q" != "max" ] && [ "$p" -gt 0 ]; then \
            COMPUTE_CORES=$(( (q + p - 1) / p )); \
        fi; \
    elif [ -f "$quota_v1" ] && [ -f "$period_v1" ]; then \
        q=$(cat "$quota_v1"); p=$(cat "$period_v1"); \
        if [ "$q" -gt 0 ] && [ "$p" -gt 0 ]; then \
            COMPUTE_CORES=$(( (q + p - 1) / p )); \
        fi; \
    fi; \
    if [ -z "$COMPUTE_CORES" ] || [ "$COMPUTE_CORES" -le 0 ]; then \
        COMPUTE_CORES=$(nproc); \
    fi; \
    export OMP_NUM_THREADS=$COMPUTE_CORES && \
    export MKL_NUM_THREADS=$COMPUTE_CORES && \
    export TORCH_NUM_THREADS=$COMPUTE_CORES && \
    export OMP_PROC_BIND=CLOSE && \
    export OMP_PLACES=CORES && \
    export OMP_SCHEDULE=STATIC && \
    export KMP_AFFINITY=granularity=fine,compact,1,0 && \
    echo "=== 🚀 Max Dynamic CPU Threading Engaged: $COMPUTE_CORES ===" && \
    python main.py \
        --cpu \
        --listen 0.0.0.0 \
        --port 7860 \
        --output-directory ./temp \
        --bf16-unet \
        --bf16-vae \
        --bf16-text-enc \
        --use-pytorch-cross-attention \
        --preview-method none