FROM nvidia/cuda:12.1.0-runtime-ubuntu22.04

# Wave tank surgery on DeepSeek-V4-Pro model
RUN apt-get update && apt-get install -y \
    python3.10 python3-pip git \
    && rm -rf /var/lib/apt/lists/*

WORKDIR /app

RUN mkdir -p /app/.cache/huggingface && chown -R 1000:1000 /app/.cache

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

COPY wave_tank_surgeon.py .
COPY space_runner.py .

ENV HF_TOKEN=${HF_TOKEN}
ENV HOME=/app
ENV HF_HOME=/app/.cache/huggingface
ENV HUGGINGFACE_HUB_CACHE=/app/.cache/huggingface/hub
ENV TRANSFORMERS_CACHE=/app/.cache/huggingface/transformers
ENV TORCH_CUDA_ARCH_LIST="80;89;90"

EXPOSE 7860
ENTRYPOINT ["python3.10", "space_runner.py"]

CMD ["python3.10", "wave_tank_surgeon.py", \
     "--model", "huihui-ai/Llama-3.3-70B-Instruct-abliterated-finetuned", \
     "--ffn-grid", "64", \
     "--attn-grid", "64", \
    "--steps", "750", \
     "--lr", "0.001", \
     "--device", "cuda:0", \
     "--dtype", "bf16", \
     "--push-to-hub", "Jibbalit/DeepSeek-V4-Pro-WaveTank"]