diff --git a/Dockerfile.spark b/Dockerfile.spark index 33cbf2d..4a101c1 100644 --- a/Dockerfile.spark +++ b/Dockerfile.spark @@ -4,12 +4,10 @@ FROM ghcr.io/aeon-7/comfyui-aeon-spark:slim WORKDIR /opt/h3-blackwell-runtime COPY . . -# Build the upstream GB10 Blackwell SageAttention3 package; the base image supplies CUDA 13 and Torch. -RUN git clone --depth 1 https://github.com/thu-ml/SageAttention.git /tmp/SageAttention \ - && cd /tmp/SageAttention/sageattention3_blackwell \ - && sed -i 's/cc_major, cc_minor = torch.cuda.get_device_capability()/cc_major, cc_minor = (12, 1)/' setup.py \ - && TORCH_CUDA_ARCH_LIST=12.1a NVCC_APPEND_FLAGS='-gencode=arch=compute_121a,code=sm_121a' python -m pip install --no-cache-dir --no-build-isolation . \ - && rm -rf /tmp/SageAttention +# Prebuilt against the base image's CUDA 13 / Torch ABI for GB10 (sm_121a). +COPY wheels/sageattn3-*.whl /tmp/wheels/ +RUN python -m pip install --no-cache-dir --no-deps /tmp/wheels/sageattn3-*.whl \ + && rm -rf /tmp/wheels RUN python -m pip install --no-cache-dir --no-deps -e . \ && python -c "import comfy_kitchen, torch; from sageattn3 import sageattn3_blackwell; print(torch.__version__, torch.version.cuda)" diff --git a/src/h3_blackwell_runtime/conditioning.py b/src/h3_blackwell_runtime/conditioning.py index cc7cb05..b5f050e 100644 --- a/src/h3_blackwell_runtime/conditioning.py +++ b/src/h3_blackwell_runtime/conditioning.py @@ -11,7 +11,11 @@ class H3PromptTokenizer: def __init__(self, tokenizer_dir: str | Path): from transformers import AutoTokenizer - self.tokenizer = AutoTokenizer.from_pretrained(str(Path(tokenizer_dir)), local_files_only=True) + self.tokenizer = AutoTokenizer.from_pretrained( + str(Path(tokenizer_dir)), + local_files_only=True, + fix_mistral_regex=True, + ) def __call__(self, prompt: str, *, device: torch.device | str = "cuda") -> torch.Tensor: if not prompt: diff --git a/wheels/sageattn3-1.0.0-cp312-cp312-linux_aarch64.whl b/wheels/sageattn3-1.0.0-cp312-cp312-linux_aarch64.whl new file mode 100644 index 0000000..8ef87d6 Binary files /dev/null and b/wheels/sageattn3-1.0.0-cp312-cp312-linux_aarch64.whl differ