tsiendragon/qwen-image-finetune/qwen-image-edit-aipc-finetune-06-quantize

# quantize_transformer_nf4.py
# Pre-quantize the Qwen-Image-Edit transformer (DiT) to NF4 and save to disk. Run once per machine.
# Prerequisites: qflux importable (run from project root with src/ on PYTHONPATH, or pip install -e .);
# conda env active + setvars.bat called (bnb's Triton quantize_4bit needs the oneAPI runtime).
# Usage:
#   python quantize_transformer_nf4.py --src <pipeline-dir> --dst <nf4-output-dir>
# After completion, set model.transformer_path: <DST> in your training config. The qflux loader
# auto-detects the saved quantization_config in <DST>/config.json and skips online quantization.
import argparse
from pathlib import Path
import torch
from transformers import BitsAndBytesConfig
from qflux.models.transformer_qwenimage import QwenImageTransformer2DModel
try:
    from qflux.utils.memory_probe import MemoryProbe
except ImportError:
    MemoryProbe = None
def main():
    p = argparse.ArgumentParser(description="Pre-quantize Qwen-Image-Edit transformer to NF4 (one-time).")
    p.add_argument("--src", required=True, help="Source pipeline dir (contains transformer/ subfolder)")
    p.add_argument("--dst", required=True, help="Destination dir for the NF4-quantized transformer")
    p.add_argument("--compute-dtype", default="bfloat16", choices=["bfloat16", "float16"],
                   help="bnb_4bit_compute_dtype (default: bfloat16)")
    p.add_argument("--device", default="auto", choices=["auto", "cpu"],
                   help="Where the load + quantization run. 'auto' (default) lets "
                        "bitsandbytes target the XPU; 'cpu' keeps the whole "
                        "load on the CPU. See this step's Troubleshooting")
    p.add_argument("--probe-interval", type=float, default=0.5,
                   help="Memory sampling interval (seconds); ignored if MemoryProbe unavailable")
    args = p.parse_args()
    weight_dtype = getattr(torch, args.compute_dtype)
    load_kwargs = {"device_map": {"": "cpu"}} if args.device == "cpu" else {}
    bnb_config = BitsAndBytesConfig(load_in_4bit=True, bnb_4bit_quant_type="nf4",
                                    bnb_4bit_compute_dtype=weight_dtype, bnb_4bit_use_double_quant=True)
    def _quantize():
        print(f"[quantize] streaming-load & quantize on {args.device}: {args.src}/transformer")
        transformer = QwenImageTransformer2DModel.from_pretrained(
            args.src, subfolder="transformer",
            quantization_config=bnb_config, torch_dtype=weight_dtype, **load_kwargs)
        Path(args.dst).mkdir(parents=True, exist_ok=True)
        print(f"[quantize] saving NF4 weights to {args.dst}")
        transformer.save_pretrained(args.dst, safe_serialization=True)
        print(f'[quantize] done. Set in config:  model.transformer_path: "{args.dst}"')
    if MemoryProbe is not None:
        with MemoryProbe(interval=args.probe_interval, label="quantize"):
            _quantize()
    else:
        print("[quantize] (qflux.utils.memory_probe not installed; running without probe)")
        _quantize()
if __name__ == "__main__":
    main()

@echo off
REM quantize_xpu.bat <src-pipeline-dir> <dst-nf4-dir> [extra args...]
REM One-time NF4 pre-quantization of the Qwen-Image-Edit transformer (DiT).
setlocal
if "%~1"=="" ( echo ERROR: source pipeline directory required. & echo Usage: quantize_xpu.bat ^<src^> ^<dst^> & exit /b 1 )
if "%~2"=="" ( echo ERROR: destination directory required. & echo Usage: quantize_xpu.bat ^<src^> ^<dst^> & exit /b 1 )
set SRC=%~1
set DST=%~2
REM Forward any further arguments to the Python script (e.g. --device cpu).
shift
shift
set "EXTRA="
:collect_extra
if "%~1"=="" goto collected_extra
set "EXTRA=%EXTRA% %1"
shift
goto collect_extra
:collected_extra
REM NoDefaultCurrentDirectoryInExePath guards against DLL loading from the current directory.
REM Cleared only for setvars.bat; restored below — DO NOT remove the restore line.
set "NoDefaultCurrentDirectoryInExePath="
call "C:\Program Files (x86)\Intel\oneAPI\setvars.bat" --force
set "NoDefaultCurrentDirectoryInExePath=1"
if errorlevel 1 ( echo ERROR: setvars.bat failed & exit /b 1 )
call conda activate qwen-image-edit-xpu
if errorlevel 1 ( echo ERROR: conda activate failed -- expected env name 'qwen-image-edit-xpu' & exit /b 1 )
set QFLUX_DOTENV_LOADED=1
set PYTHONUTF8=1
set SYCL_CACHE_PERSISTENT=1
echo [quantize] Quantizing transformer: %SRC% -^> %DST%
python quantize_transformer_nf4.py --src "%SRC%" --dst "%DST%"%EXTRA%
echo [quantize] Done. Exit code: %ERRORLEVEL%
endlocal

quantize_xpu.bat <SRC-pipeline-dir> <DST-nf4-transformer-dir>
REM or, after activating conda + setvars.bat manually:
python quantize_transformer_nf4.py --src <SRC> --dst <DST>

model:
  pretrained_model_name_or_path: <SRC>   # still needed for text_encoder, VAE, etc.
  transformer_path: <DST>                # NF4-quantized DiT
  quantize_type: "nf4"

# quantize_text_encoder_nf4.py
# Pre-quantize the Qwen2.5-VL text_encoder to NF4 and save to disk. Run once per machine. Saves
# ~8 GB CPU RAM vs bf16 load — the margin that prevents the ~31.5 GB ceiling crash on 32 GB machines.
# Prerequisites: qflux importable; conda env active + setvars.bat called (oneAPI runtime + bnb XPU).
# Usage:
#   python quantize_text_encoder_nf4.py --src <pipeline-dir> --dst <nf4-text-encoder-dir>
# After completion, set model.text_encoder_path: <DST> in your training config. The qflux loader
# auto-detects quantization_config in <DST>/config.json and skips online quantization.
import argparse
from pathlib import Path
import torch
from transformers import BitsAndBytesConfig, Qwen2_5_VLForConditionalGeneration
try:
    from qflux.utils.memory_probe import MemoryProbe
except ImportError:
    MemoryProbe = None
def main():
    p = argparse.ArgumentParser(description="Pre-quantize Qwen2.5-VL text_encoder to NF4 (one-time).")
    p.add_argument("--src", required=True, help="Source pipeline dir (contains text_encoder/ subfolder)")
    p.add_argument("--dst", required=True, help="Destination dir for the NF4-quantized text_encoder")
    p.add_argument("--compute-dtype", default="bfloat16", choices=["bfloat16", "float16"],
                   help="bnb_4bit_compute_dtype (default: bfloat16)")
    p.add_argument("--device", default="auto", choices=["auto", "cpu"],
                   help="Where the load + quantization run. 'auto' (default) lets "
                        "bitsandbytes target the XPU; 'cpu' keeps the whole "
                        "load on the CPU. See this step's Troubleshooting")
    p.add_argument("--probe-interval", type=float, default=0.5,
                   help="Memory sampling interval (seconds); ignored if MemoryProbe unavailable")
    args = p.parse_args()
    weight_dtype = getattr(torch, args.compute_dtype)
    load_kwargs = {"device_map": {"": "cpu"}} if args.device == "cpu" else {}
    bnb_config = BitsAndBytesConfig(load_in_4bit=True, bnb_4bit_quant_type="nf4",
                                    bnb_4bit_compute_dtype=weight_dtype, bnb_4bit_use_double_quant=True)
    def _quantize():
        print(f"[quantize_te] streaming-load & quantize on {args.device}: {args.src}/text_encoder")
        model = Qwen2_5_VLForConditionalGeneration.from_pretrained(
            args.src, subfolder="text_encoder",
            quantization_config=bnb_config, torch_dtype=weight_dtype, **load_kwargs)
        Path(args.dst).mkdir(parents=True, exist_ok=True)
        print(f"[quantize_te] saving NF4 text_encoder to {args.dst}")
        model.save_pretrained(args.dst, safe_serialization=True)
        print(f'[quantize_te] done. Set in config:  model.text_encoder_path: "{args.dst}"')