Dockerfile

# hart-cli:thor — containerized HART (mit-han-lab/hart) text-to-image CLI
# for NVIDIA aarch64 CUDA / unified-memory boxes. Tested on Jetson AGX Thor
# (sm_110, CUDA 13). Two stages: compile HART's custom CUDA kernels, then
# install the package + CLI on top.
#
# HART_SRC (build context) must be a plain clone of https://github.com/mit-han-lab/hart —
# the Thor/Blackwell arch-support patch is applied automatically during the build.

# ---- Stage 1: compile HART's custom CUDA kernels -----------------------------
FROM nvcr.io/nvidia/pytorch:25.08-py3 AS kernels

ARG HART_CUDA_ARCH=11.0
ENV TORCH_CUDA_ARCH_LIST=${HART_CUDA_ARCH}

WORKDIR /workspace
COPY hart-src /workspace/repo
COPY patches/0001-thor-blackwell-sm110-arch-support.patch /tmp/thor-arch.patch

# Apply unconditionally: this build always targets a Blackwell-generation
# compute capability by default, and the patch is a no-op on an already-patched
# clone (patch -N skips hunks that are already applied instead of failing).
RUN cd /workspace/repo && patch -p1 -N < /tmp/thor-arch.patch || true

RUN apt-get update && apt-get install -y --no-install-recommends \
        build-essential python3-dev ninja-build \
    && rm -rf /var/lib/apt/lists/* \
    && cd /workspace/repo/hart/kernels && python setup.py install

# ---- Stage 2: install the HART package + CLI ----------------------------------
FROM kernels AS cli

ARG HART_GIT_COMMIT=unknown
ARG HART_CUDA_ARCH=11.0
ARG HART_BUILD_DATE=unknown
LABEL org.opencontainers.image.title="HART (mit-han-lab/hart) — aarch64 CUDA container" \
      org.opencontainers.image.source="https://github.com/mit-han-lab/hart" \
      org.opencontainers.image.revision="${HART_GIT_COMMIT}" \
      org.opencontainers.image.created="${HART_BUILD_DATE}" \
      com.hart.cuda_arch="${HART_CUDA_ARCH}"

ENV HF_HUB_OFFLINE=1 \
    TRANSFORMERS_OFFLINE=1 \
    HF_DATASETS_OFFLINE=1 \
    WANDB_MODE=offline \
    PYTORCH_CUDA_ALLOC_CONF=max_split_size_mb:128

WORKDIR /workspace/repo

# --no-deps: HART's own pinned requirements conflict with the NGC image's
# newer torch/torchvision/etc — those are already correct for this base image.
# Install only the extra runtime deps the CLI actually needs.
RUN pip install --no-cache-dir . --no-deps && \
    pip install --no-cache-dir \
      transformers==4.42.2 tokenizers==0.19.1 sentencepiece==0.2.0 \
      accelerate==0.27.2 einops==0.6.1 timm==0.9.12

# The safety-checker gate in upstream's sample.py depends on a "shield model"
# checkpoint that was never published — this CLI omits it entirely rather
# than shipping a broken reference to an unavailable model. See README.
COPY run_hart_cli.py /usr/local/bin/hart-txt2img
RUN chmod +x /usr/local/bin/hart-txt2img

WORKDIR /workspace
ENV NVIDIA_VISIBLE_DEVICES=all
ENV NVIDIA_DRIVER_CAPABILITIES=compute,utility

ENTRYPOINT ["hart-txt2img"]