pass tests now building wheels

README.md

@@ -10,13 +10,13 @@ Join our user groups on [**Slack**](https://join.slack.com/t/nunchaku/shared_inv
 
 ## News
 
+- **[2025-04-05]** 🚀 **Nunchaku v0.2.0 released!** This release brings **multi-LoRA** and **ControlNet** support with even faster performance. We've also added compatibility for **20-series GPUs** — Nunchaku is now more accessible than ever!
 - **[2025-03-17]** 🚀 Released NVFP4 4-bit [Shuttle-Jaguar](https://huggingface.co/mit-han-lab/svdq-int4-shuttle-jaguar) and FLUX.1-tools and also upgraded the INT4 FLUX.1-tool models. Download and update your models from our [HuggingFace](https://huggingface.co/collections/mit-han-lab/svdquant-67493c2c2e62a1fc6e93f45c) or [ModelScope](https://modelscope.cn/collections/svdquant-468e8f780c2641) collections!
 - **[2025-03-13]** 📦 Separate the ComfyUI node into a [standalone repository](https://github.com/mit-han-lab/ComfyUI-nunchaku) for easier installation and release node v0.1.6! Plus, [4-bit Shuttle-Jaguar](https://huggingface.co/mit-han-lab/svdq-int4-shuttle-jaguar) is now fully supported!
 - **[2025-03-07]** 🚀 **Nunchaku v0.1.4 Released!** We've supported [4-bit text encoder and per-layer CPU offloading](#Low-Memory-Inference), reducing FLUX's minimum memory requirement to just **4 GiB** while maintaining a **2–3× speedup**. This update also fixes various issues related to resolution, LoRA, pin memory, and runtime stability. Check out the release notes for full details!
 - **[2025-02-20]** 🚀 We release the [pre-built wheels](https://huggingface.co/mit-han-lab/nunchaku) to simplify installation! Check [here](#Installation) for the guidance!
 - **[2025-02-20]** 🚀 **Support NVFP4 precision on NVIDIA RTX 5090!** NVFP4 delivers superior image quality compared to INT4, offering **~3× speedup** on the RTX 5090 over BF16. Learn more in our [blog](https://hanlab.mit.edu/blog/svdquant-nvfp4), checkout  [`examples`](./examples) for usage and try [our demo](https://svdquant.mit.edu/flux1-schnell/) online!
 - **[2025-02-18]** 🔥 [**Customized LoRA conversion**](#Customized-LoRA) and [**model quantization**](#Customized-Model-Quantization) instructions are now available! **[ComfyUI](./comfyui)** workflows now support **customized LoRA**, along with **FLUX.1-Tools**!
-- **[2025-02-14]** 🔥 **[LoRA conversion script](nunchaku/convert_lora.py)** is now available! [ComfyUI FLUX.1-tools workflows](./comfyui) is released!
 - **[2025-02-11]** 🎉 **[SVDQuant](http://arxiv.org/abs/2411.05007) has been selected as a ICLR 2025 Spotlight! FLUX.1-tools Gradio demos are now available!** Check [here](#gradio-demos) for the usage details! Our new [depth-to-image demo](https://svdquant.mit.edu/flux1-depth-dev/) is also online—try it out!
 
 
@@ -63,31 +63,21 @@ SVDQuant is a post-training quantization technique for 4-bit weights and activat
 
 ### Wheels
 
-**Note:** For native Windows users, we have released a preliminary wheel to ease the installation. See [here](https://github.com/mit-han-lab/nunchaku/issues/169) for more details!
-
-#### For Windows WSL Users
-
-To install and use WSL (Windows Subsystem for Linux), follow the instructions [here](https://learn.microsoft.com/en-us/windows/wsl/install). You can also install WSL directly by running the following commands in PowerShell:
-```shell
-wsl --install # install the latest WSL
-wsl # launch WSL
-```
-
-#### Prerequisites for all users
+#### Prerequisites
 Before installation, ensure you have [PyTorch>=2.5](https://pytorch.org/) installed. For example, you can use the following command to install PyTorch 2.6:
 
 ```shell
 pip install torch==2.6 torchvision==0.21 torchaudio==2.6
 ```
 
-#### Installing nunchaku
-Once PyTorch is installed, you can directly install `nunchaku` from our [Hugging Face repository](https://huggingface.co/mit-han-lab/nunchaku/tree/main). Be sure to select the appropriate wheel for your Python and PyTorch version. For example, for Python 3.11 and PyTorch 2.6:
+#### Install nunchaku
+Once PyTorch is installed, you can directly install `nunchaku` from our whell repositories [Hugging Face](https://huggingface.co/mit-han-lab/nunchaku/tree/main) or [ModelScope](https://modelscope.cn/models/Lmxyy1999/nunchaku). Be sure to select the appropriate wheel for your Python and PyTorch version. For example, for Python 3.11 and PyTorch 2.6:
 
 ```shell
-pip install https://huggingface.co/mit-han-lab/nunchaku/resolve/main/nunchaku-0.1.4+torch2.6-cp311-cp311-linux_x86_64.whl
+pip install https://huggingface.co/mit-han-lab/nunchaku/resolve/main/nunchaku-0.2.0+torch2.6-cp311-cp311-linux_x86_64.whl
 ```
 
-**Note**: NVFP4 wheels are not currently available because PyTorch has not officially supported CUDA 12.8. To use NVFP4, you will need **Blackwell GPUs (e.g., 50-series GPUs)** and must **build from source**.
+**Note**: If you're using a Blackwell GPU (e.g., 50-series GPUs), install a wheel with PyTorch 2.7. Additionally, use **FP4 models** instead of INT4 models.
 
 ### Build from Source
 
@@ -97,32 +87,38 @@ pip install https://huggingface.co/mit-han-lab/nunchaku/resolve/main/nunchaku-0.
 
 *  For Windows users, please refer to [this issue](https://github.com/mit-han-lab/nunchaku/issues/6) for the instruction. Please upgrade your MSVC compiler to the latest version.
 
-*  We currently support only NVIDIA GPUs with architectures sm_86 (Ampere: RTX 3090, A6000), sm_89 (Ada: RTX 4090), and sm_80 (A100). See [this issue](https://github.com/mit-han-lab/nunchaku/issues/1) for more details.
+*  We currently support only NVIDIA GPUs with architectures sm_75 (Turing: RTX 2080), sm_86 (Ampere: RTX 3090, A6000), sm_89 (Ada: RTX 4090), and sm_80 (A100). See [this issue](https://github.com/mit-han-lab/nunchaku/issues/1) for more details.
 
 
 1. Install dependencies:
-  ```shell
-  conda create -n nunchaku python=3.11
-  conda activate nunchaku
-  pip install torch torchvision torchaudio
-  pip install ninja wheel diffusers transformers accelerate sentencepiece protobuf huggingface_hub
-  pip install peft opencv-python gradio spaces GPUtil  # For gradio demos
-  ```
 
- To enable NVFP4 on Blackwell GPUs (e.g., 50-series GPUs), please install nightly PyTorch with CUDA 12.8. The installation command can be:
+   ```shell
+   conda create -n nunchaku python=3.11
+   conda activate nunchaku
+   pip install torch torchvision torchaudio
+   pip install ninja wheel diffusers transformers accelerate sentencepiece protobuf huggingface_hub
+   
+   # For gradio demos
+   pip install peft opencv-python gradio spaces GPUtil  
+   ```
+
+   To enable NVFP4 on Blackwell GPUs (e.g., 50-series GPUs), please install nightly PyTorch with CUDA 12.8. The installation command can be:
 
-  ```shell
-  pip install --pre torch torchvision torchaudio --index-url https://download.pytorch.org/whl/nightly/cu128
-  ```
+   ```shell
+   pip install --pre torch torchvision torchaudio --index-url https://download.pytorch.org/whl/nightly/cu128
+   ```
 
 2. Install `nunchaku` package:
-    Make sure you have `gcc/g++>=11`. If you don't, you can install it via Conda:
+    Make sure you have `gcc/g++>=11`. If you don't, you can install it via Conda on Linux:
 
     ```shell
     conda install -c conda-forge gxx=11 gcc=11
     ```
 
-    Then build the package from source:
+    For Windows users, you can download and install the lastest [Visual Studio](https://visualstudio.microsoft.com/thank-you-downloading-visual-studio/?sku=Community&channel=Release&version=VS2022&source=VSLandingPage&cid=2030&passive=false).
+    
+    Then build the package from source with
+    
     ```shell
     git clone https://github.com/mit-han-lab/nunchaku.git
     cd nunchaku
@@ -130,51 +126,69 @@ pip install https://huggingface.co/mit-han-lab/nunchaku/resolve/main/nunchaku-0.
     git submodule update
     python setup.py develop
     ```
+    
+    If you are building wheels for distribution, use:
+    
+    ```shell
+    NUNCHAKU_INSTALL_MODE=ALL NUNCHAKU_BUILD_WHEELS=1 python -m build --wheel --no-isolation
+    ```
+    
+    Make sure to set the environment variable `NUNCHAKU_INSTALL_MODE` to `ALL`. Otherwise, the generated wheels will only work on GPUs with the same architecture as the build machine.
+
+### Docker (Coming soon)
 
 **[Optional]** You can verify your installation by running: `python -m nunchaku.test`. This command will download and run our 4-bit FLUX.1-schnell model.
 
 ## Usage Example
 
-In [examples](examples), we provide minimal scripts for running INT4 [FLUX.1](https://github.com/black-forest-labs/flux) and [SANA](https://github.com/NVlabs/Sana) models with Nunchaku. For example, the [script](examples/int4-flux.1-dev.py) for [FLUX.1-dev](https://huggingface.co/black-forest-labs/FLUX.1-dev) is as follows:
+In [examples](examples), we provide minimal scripts for running INT4 [FLUX.1](https://github.com/black-forest-labs/flux) and [SANA](https://github.com/NVlabs/Sana) models with Nunchaku. For example, the [script](examples/flux.1-dev.py) for [FLUX.1-dev](https://huggingface.co/black-forest-labs/FLUX.1-dev) is as follows:
 
 ```python
 import torch
 from diffusers import FluxPipeline
 
 from nunchaku import NunchakuFluxTransformer2dModel
+from nunchaku.utils import get_precision
 
-transformer = NunchakuFluxTransformer2dModel.from_pretrained("mit-han-lab/svdq-int4-flux.1-dev")
+precision = get_precision()  # auto-detect your precision is 'int4' or 'fp4' based on your GPU
+transformer = NunchakuFluxTransformer2dModel.from_pretrained(f"mit-han-lab/svdq-{precision}-flux.1-dev")
 pipeline = FluxPipeline.from_pretrained(
     "black-forest-labs/FLUX.1-dev", transformer=transformer, torch_dtype=torch.bfloat16
 ).to("cuda")
 image = pipeline("A cat holding a sign that says hello world", num_inference_steps=50, guidance_scale=3.5).images[0]
-image.save("flux.1-dev.png")
+image.save(f"flux.1-dev-{precision}.png")
 ```
 
 Specifically, `nunchaku` shares the same APIs as [diffusers](https://github.com/huggingface/diffusers) and can be used in a similar way.
 
-### Low Memory Inference
+### First-Block Cache and Low-Precision Attention
+
+
 
-To further reduce GPU memory usage, you can use our 4-bit T5 encoder along with CPU offloading, requiring a minimum of just 4GiB of memory. The usage is also simple in the diffusers' way. For example, the [script](examples/flux.1-dev-qencoder.py) for FLUX.1-dev is as follows:
+### CPU Offloading
+
+To further reduce GPU memory usage, you can use CPU offloading, requiring a minimum of just 4GiB of memory. The usage is also simple in the diffusers' way. For example, the [script](examples/flux.1-dev-offload.py) for FLUX.1-dev is as follows:
 
 ```python
 import torch
 from diffusers import FluxPipeline
 
-from nunchaku import NunchakuFluxTransformer2dModel, NunchakuT5EncoderModel
+from nunchaku import NunchakuFluxTransformer2dModel
+from nunchaku.utils import get_precision
 
+precision = get_precision()  # auto-detect your precision is 'int4' or 'fp4' based on your GPU
 transformer = NunchakuFluxTransformer2dModel.from_pretrained(
-    "mit-han-lab/svdq-int4-flux.1-dev", offload=True
+    f"mit-han-lab/svdq-{precision}-flux.1-dev", offload=True
 )  # set offload to False if you want to disable offloading
-text_encoder_2 = NunchakuT5EncoderModel.from_pretrained("mit-han-lab/svdq-flux.1-t5")
 pipeline = FluxPipeline.from_pretrained(
-    "black-forest-labs/FLUX.1-dev", text_encoder_2=text_encoder_2, transformer=transformer, torch_dtype=torch.bfloat16
-).to("cuda")
-pipeline.enable_sequential_cpu_offload()  # remove this line if you want to disable the CPU offloading
+    "black-forest-labs/FLUX.1-dev", transformer=transformer, torch_dtype=torch.bfloat16
+)  # no need to set the device here
+pipeline.enable_sequential_cpu_offload()  # diffusers' offloading
 image = pipeline("A cat holding a sign that says hello world", num_inference_steps=50, guidance_scale=3.5).images[0]
-image.save("flux.1-dev.png")
+image.save(f"flux.1-dev-{precision}.png")
 ```
 
+
 ## Customized LoRA
 
 ![lora](./assets/lora.jpg)
@@ -182,7 +196,7 @@ image.save("flux.1-dev.png")
 [SVDQuant](http://arxiv.org/abs/2411.05007) seamlessly integrates with off-the-shelf LoRAs without requiring requantization. You can simply use your LoRA with:
 
 ```python
-transformer.update_lora_params(path_to_your_converted_lora)
+transformer.update_lora_params(path_to_your_lora)
 transformer.set_lora_strength(lora_strength)
 ```
 
@@ -216,7 +230,7 @@ image = pipeline(
 image.save(f"flux.1-dev-ghibsky-{precision}.png")
 ```
 
-**For ComfyUI users, we have implemented a node to convert the LoRA weights on the fly. All you need to do is specify the correct LoRA format. Please refer to [mit-han-lab/ComfyUI-nunchaku](https://github.com/mit-han-lab/ComfyUI-nunchaku) for more details.**
+**For ComfyUI users, you can directly use our LoRA loader. The converted LoRA is deprecated. Please refer to [mit-han-lab/ComfyUI-nunchaku](https://github.com/mit-han-lab/ComfyUI-nunchaku) for more details.**
 
 ## ComfyUI
 
@@ -235,7 +249,7 @@ Please refer to [mit-han-lab/ComfyUI-nunchaku](https://github.com/mit-han-lab/Co
 
 ## Customized Model Quantization
 
-Please refer to [mit-han-lab/deepcompressor](https://github.com/mit-han-lab/deepcompressor/tree/main/examples/diffusion).
+Please refer to [mit-han-lab/deepcompressor](https://github.com/mit-han-lab/deepcompressor/tree/main/examples/diffusion). A simpler workflow is coming soon.
 
 ## Benchmark
 

examples/flux.1-dev-offload.py

@@ -10,7 +10,7 @@ transformer = NunchakuFluxTransformer2dModel.from_pretrained(
 )  # set offload to False if you want to disable offloading
 pipeline = FluxPipeline.from_pretrained(
     "black-forest-labs/FLUX.1-dev", transformer=transformer, torch_dtype=torch.bfloat16
-)
-pipeline.enable_sequential_cpu_offload()  # remove this line if you want to disable the CPU offloading
+)  # no need to set the device here
+pipeline.enable_sequential_cpu_offload()  # diffusers' offloading
 image = pipeline("A cat holding a sign that says hello world", num_inference_steps=50, guidance_scale=3.5).images[0]
 image.save(f"flux.1-dev-{precision}.png")

nunchaku/__version__.py

-__version__ = "0.2.0dev0"
+__version__ = "0.2.0"

nunchaku/models/transformers/transformer_flux.py

@@ -82,9 +82,15 @@ class NunchakuFluxTransformerBlocks(nn.Module):
         image_rotary_emb = image_rotary_emb.to(self.device)
 
         if controlnet_block_samples is not None:
-            controlnet_block_samples = torch.stack(controlnet_block_samples).to(self.device)
-        if controlnet_single_block_samples is not None:
-            controlnet_single_block_samples = torch.stack(controlnet_single_block_samples).to(self.device)
+            controlnet_block_samples = (
+                torch.stack(controlnet_block_samples).to(self.device) if len(controlnet_block_samples) > 0 else None
+            )
+        if controlnet_single_block_samples is not None and len(controlnet_single_block_samples) > 0:
+            controlnet_single_block_samples = (
+                torch.stack(controlnet_single_block_samples).to(self.device)
+                if len(controlnet_single_block_samples) > 0
+                else None
+            )
 
         assert image_rotary_emb.ndim == 6
         assert image_rotary_emb.shape[0] == 1

scripts/build_all_linux_wheels.sh

@@ -7,8 +7,18 @@ cuda_versions=("12.4")
 # Loop through all combinations of Python, Torch, and CUDA versions
 for python_version in "${python_versions[@]}"; do
   for torch_version in "${torch_versions[@]}"; do
+    # Skip building for Python 3.13 and PyTorch 2.5
+    if [[ "$python_version" == "3.13" && "$torch_version" == "2.5" ]]; then
+      echo "Skipping Python 3.13 with PyTorch 2.5"
+      continue
+    fi
     for cuda_version in "${cuda_versions[@]}"; do
       bash scripts/build_linux_wheel.sh "$python_version" "$torch_version" "$cuda_version"
     done
   done
-done
\ No newline at end of file
+done
+
+bash scripts/build_linux_wheel_cu128.sh "3.10" "2.7" "12.8"
+bash scripts/build_linux_wheel_cu128.sh "3.11" "2.7" "12.8"
+bash scripts/build_linux_wheel_cu128.sh "3.12" "2.7" "12.8"
+bash scripts/build_linux_wheel_cu128.sh "3.13" "2.7" "12.8"
\ No newline at end of file

scripts/build_all_windows_wheels.cmd

+@echo off
+setlocal enabledelayedexpansion
+
+REM Define Python and Torch versions
+set "python_versions=3.10 3.11 3.12 3.13"
+set "torch_versions=2.5 2.6"
+set "cuda_version=12.4"
+
+REM Iterate over Python and Torch versions
+for %%P in (%python_versions%) do (
+    for %%T in (%torch_versions%) do (
+        REM Python 3.13 only supports Torch 2.6 and above
+        if not "%%P"=="3.13" (
+            echo Building with Python %%P, Torch %%T, CUDA %cuda_version%...
+            call scripts\build_windows_wheel.cmd %%P %%T %cuda_version%
+        ) else if not "%%T"=="2.5" (
+            echo Building with Python %%P, Torch %%T, CUDA %cuda_version%...
+            call scripts\build_windows_wheel.cmd %%P %%T %cuda_version%
+        )
+    )
+)
+
+call scripts\build_windows_wheel.cmd 3.10 2.7 12.8
+call scripts\build_windows_wheel.cmd 3.11 2.7 12.8
+call scripts\build_windows_wheel.cmd 3.12 2.7 12.8
+call scripts\build_windows_wheel.cmd 3.13 2.7 12.8
+
+echo All builds completed successfully!
+exit /b 0

scripts/build_linux_wheel_cu128.sh

+#!/bin/bash
+# Modified from https://github.com/sgl-project/sglang/blob/main/sgl-kernel/build.sh
+set -ex
+PYTHON_VERSION=$1
+TORCH_VERSION=$2 # has no use for now
+CUDA_VERSION=$3
+MAX_JOBS=${4:-} # optional
+PYTHON_ROOT_PATH=/opt/python/cp${PYTHON_VERSION//.}-cp${PYTHON_VERSION//.}
+
+# Check if TORCH_VERSION is 2.5 or 2.6 and set the corresponding versions for TORCHVISION and TORCHAUDIO
+#if [ "$TORCH_VERSION" == "2.5" ]; then
+#  TORCHVISION_VERSION="0.20"
+#  TORCHAUDIO_VERSION="2.5"
+#  echo "TORCH_VERSION is 2.5, setting TORCHVISION_VERSION to $TORCHVISION_VERSION and TORCHAUDIO_VERSION to $TORCHAUDIO_VERSION"
+#elif [ "$TORCH_VERSION" == "2.6" ]; then
+#  TORCHVISION_VERSION="0.21"
+#  TORCHAUDIO_VERSION="2.6"
+#  echo "TORCH_VERSION is 2.6, setting TORCHVISION_VERSION to $TORCHVISION_VERSION and TORCHAUDIO_VERSION to $TORCHAUDIO_VERSION"
+#else
+#  echo "TORCH_VERSION is not 2.5 or 2.6, no changes to versions."
+#fi
+
+docker run --rm \
+    -v "$(pwd)":/nunchaku \
+    pytorch/manylinux-builder:cuda${CUDA_VERSION} \
+    bash -c "
+    cd /nunchaku && \
+    rm -rf build && \
+    yum install -y devtoolset-11 && \
+    source scl_source enable devtoolset-11 && \
+    gcc --version && g++ --version && \
+    ${PYTHON_ROOT_PATH}/bin/pip install --pre torch torchvision torchaudio --index-url https://download.pytorch.org/whl/nightly/cu128 && \
+    ${PYTHON_ROOT_PATH}/bin/pip install build ninja wheel setuptools && \
+    export NUNCHAKU_INSTALL_MODE=ALL && \
+    export NUNCHAKU_BUILD_WHEELS=1 && \
+    export MAX_JOBS=${MAX_JOBS} && \
+    ${PYTHON_ROOT_PATH}/bin/python -m build --wheel --no-isolation
+    "
\ No newline at end of file

scripts/build_windows_wheel.cmd

+@echo off
+setlocal enabledelayedexpansion
+
+:: get arguments
+set PYTHON_VERSION=%1
+set TORCH_VERSION=%2
+set CUDA_VERSION=%3
+
+set CUDA_SHORT_VERSION=%CUDA_VERSION:.=%
+echo %CUDA_SHORT_VERSION%
+
+:: setup some variables
+if "%TORCH_VERSION%"=="2.5" (
+    set TORCHVISION_VERSION=0.20
+    set TORCHAUDIO_VERSION=2.5
+) else if "%TORCH_VERSION%"=="2.6" (
+    set TORCHVISION_VERSION=0.21
+    set TORCHAUDIO_VERSION=2.6
+) else (
+    echo TORCH_VERSION is not 2.5 or 2.6, no changes to versions.
+)
+echo setting TORCHVISION_VERSION to %TORCHVISION_VERSION% and TORCHAUDIO_VERSION to %TORCHAUDIO_VERSION%
+
+:: conda environment name
+set ENV_NAME=build_env_%PYTHON_VERSION%_%TORCH_VERSION%
+echo Using conda environment: %ENV_NAME%
+
+:: create conda environment
+call conda create -y -n %ENV_NAME% python=%PYTHON_VERSION%
+call conda activate %ENV_NAME%
+
+:: install dependencies
+call pip install ninja setuptools wheel build
+call pip install --no-cache-dir torch==%TORCH_VERSION% torchvision==%TORCHVISION_VERSION% torchaudio==%TORCHAUDIO_VERSION% --index-url "https://download.pytorch.org/whl/cu%CUDA_SHORT_VERSION%/"
+
+:: set environment variables
+set NUNCHAKU_INSTALL_MODE=ALL
+set NUNCHAKU_BUILD_WHEELS=1
+
+:: cd to the parent directory
+cd /d "%~dp0.."
+if exist build rd /s /q build
+
+:: set up Visual Studio compilation environment
+call "C:\Program Files (x86)\Microsoft Visual Studio\2022\BuildTools\Common7\Tools\VsDevCmd.bat" -startdir=none -arch=x64 -host_arch=x64
+set DISTUTILS_USE_SDK=1
+
+:: build wheels
+python -m build --wheel --no-isolation
+
+:: exit conda
+call conda deactivate
+call conda remove -y -n %ENV_NAME% --all
+
+echo Build complete!

scripts/build_windows_wheel_cu128.cmd

+@echo off
+setlocal enabledelayedexpansion
+
+:: get arguments
+set PYTHON_VERSION=%1
+set TORCH_VERSION=%2
+set CUDA_VERSION=%3
+
+set CUDA_SHORT_VERSION=%CUDA_VERSION:.=%
+echo %CUDA_SHORT_VERSION%
+
+:: conda environment name
+set ENV_NAME=build_env_%PYTHON_VERSION%_%TORCH_VERSION%
+echo Using conda environment: %ENV_NAME%
+
+:: create conda environment
+call conda create -y -n %ENV_NAME% python=%PYTHON_VERSION%
+call conda activate %ENV_NAME%
+
+:: install dependencies
+call pip install ninja setuptools wheel build
+call pip install --pre torch torchvision torchaudio --index-url https://download.pytorch.org/whl/nightly/cu128
+
+:: set environment variables
+set NUNCHAKU_INSTALL_MODE=ALL
+set NUNCHAKU_BUILD_WHEELS=1
+
+:: cd to the parent directory
+cd /d "%~dp0.."
+if exist build rd /s /q build
+
+:: set up Visual Studio compilation environment
+call "C:\Program Files (x86)\Microsoft Visual Studio\2022\BuildTools\Common7\Tools\VsDevCmd.bat" -startdir=none -arch=x64 -host_arch=x64
+set DISTUTILS_USE_SDK=1
+
+:: build wheels
+python -m build --wheel --no-isolation
+
+:: exit conda
+call conda deactivate
+call conda remove -y -n %ENV_NAME% --all
+
+echo Build complete!

scripts/build_windows_wheels.ps1

-param (
-    [string]$PYTHON_VERSION,
-    [string]$TORCH_VERSION,
-    [string]$CUDA_VERSION,
-    [string]$MAX_JOBS = ""
-)
-
-# Check if TORCH_VERSION is 2.5 or 2.6 and set the corresponding versions for TORCHVISION and TORCHAUDIO
-if ($TORCH_VERSION -eq "2.5") {
-    $TORCHVISION_VERSION = "0.20"
-    $TORCHAUDIO_VERSION = "2.5"
-    Write-Output "TORCH_VERSION is 2.5, setting TORCHVISION_VERSION to $TORCHVISION_VERSION and TORCHAUDIO_VERSION to $TORCHAUDIO_VERSION"
-}
-elseif ($TORCH_VERSION -eq "2.6") {
-    $TORCHVISION_VERSION = "0.21"
-    $TORCHAUDIO_VERSION = "2.6"
-    Write-Output "TORCH_VERSION is 2.6, setting TORCHVISION_VERSION to $TORCHVISION_VERSION and TORCHAUDIO_VERSION to $TORCHAUDIO_VERSION"
-}
-else {
-    Write-Output "TORCH_VERSION is not 2.5 or 2.6, no changes to versions."
-}
-
-# Conda 环境名称
-$ENV_NAME = "build_env_$PYTHON_VERSION_$TORCH_VERSION"
-
-# 创建 Conda 环境
-conda create -y -n $ENV_NAME python=$PYTHON_VERSION
-conda activate $ENV_NAME
-
-# 安装依赖
-conda install -y ninja setuptools wheel pip
-pip install --no-cache-dir torch==$TORCH_VERSION torchvision==$TORCHVISION_VERSION torchaudio==$TORCHAUDIO_VERSION --index-url "https://download.pytorch.org/whl/cu$($CUDA_VERSION.Substring(0,2))/"
-
-# 设置环境变量
-$env:NUNCHAKU_INSTALL_MODE="ALL"
-$env:NUNCHAKU_BUILD_WHEELS="1"
-$env:MAX_JOBS=$MAX_JOBS
-
-# 进入当前脚本所在目录并构建 wheels
-Set-Location -Path "$PSScriptRoot\.."
-if (Test-Path "build") { Remove-Item -Recurse -Force "build" }
-
-python -m build --wheel --no-isolation
-
-# 退出 Conda 环境
-conda deactivate
-conda remove -y -n $ENV_NAME --all
-Write-Output "Build complete!"

scripts/linux_cleanup.sh

 #!/bin/bash
 set -ex
 
-#docker run --rm \
-#    -v "$(pwd)":/nunchaku \
-#    pytorch/manylinux-builder:cuda12.4 \
-#    bash -c "cd /nunchaku && rm -r *"
-docker run --rm -it \
+docker run --rm \
     -v "$(pwd)":/nunchaku \
     pytorch/manylinux-builder:cuda12.4 \
-    bash
\ No newline at end of file
+    bash -c "cd /nunchaku && rm -rf *"
\ No newline at end of file

tests/flux/test_flux_tools.py

@@ -39,7 +39,7 @@ def test_flux_depth_dev():
         attention_impl="nunchaku-fp16",
         cpu_offload=False,
         cache_threshold=0,
-        expected_lpips=0.103 if get_precision() == "int4" else 0.120,
+        expected_lpips=0.170 if get_precision() == "int4" else 0.120,
     )
 
 
@@ -140,5 +140,5 @@ def test_flux_dev_redux():
         attention_impl="nunchaku-fp16",
         cpu_offload=False,
         cache_threshold=0,
-        expected_lpips=0.187 if get_precision() == "int4" else 0.55,  # redux seems to generate different images on 5090
+        expected_lpips=0.198 if get_precision() == "int4" else 0.55,  # redux seems to generate different images on 5090
     )