From 9448bf1caa8ad5bd704ae2610bc0b2738c9ca51f Mon Sep 17 00:00:00 2001 From: oobabooga <112222186+oobabooga@users.noreply.github.com> Date: Tue, 2 Dec 2025 14:22:51 -0800 Subject: [PATCH] Image generation: add torchao quantization (supports torch.compile) --- modules/image_models.py | 31 +++++++++++++++++-- modules/shared.py | 2 +- modules/ui_image_generation.py | 4 +-- requirements/full/requirements.txt | 1 + requirements/full/requirements_amd.txt | 1 + requirements/full/requirements_amd_noavx2.txt | 1 + .../full/requirements_apple_intel.txt | 1 + .../full/requirements_apple_silicon.txt | 1 + requirements/full/requirements_cpu_only.txt | 1 + .../full/requirements_cpu_only_noavx2.txt | 1 + requirements/full/requirements_noavx2.txt | 1 + requirements/full/requirements_nowheels.txt | 1 + 12 files changed, 40 insertions(+), 6 deletions(-) diff --git a/modules/image_models.py b/modules/image_models.py index f2e0276a..2ef1e730 100644 --- a/modules/image_models.py +++ b/modules/image_models.py @@ -10,13 +10,14 @@ def get_quantization_config(quant_method): Get the appropriate quantization config based on the selected method. Args: - quant_method: One of 'none', 'bnb-8bit', 'bnb-4bit' + quant_method: One of 'none', 'bnb-8bit', 'bnb-4bit', + 'torchao-int8wo', 'torchao-fp4', 'torchao-float8wo' Returns: PipelineQuantizationConfig or None """ import torch - from diffusers import BitsAndBytesConfig, QuantoConfig + from diffusers import BitsAndBytesConfig, TorchAoConfig from diffusers.quantizers import PipelineQuantizationConfig if quant_method == 'none' or not quant_method: @@ -45,6 +46,30 @@ def get_quantization_config(quant_method): } ) + # torchao int8 weight-only + elif quant_method == 'torchao-int8wo': + return PipelineQuantizationConfig( + quant_mapping={ + "transformer": TorchAoConfig("int8wo") + } + ) + + # torchao fp4 (e2m1) + elif quant_method == 'torchao-fp4': + return PipelineQuantizationConfig( + quant_mapping={ + "transformer": TorchAoConfig("fp4_e2m1") + } + ) + + # torchao float8 weight-only + elif quant_method == 'torchao-float8wo': + return PipelineQuantizationConfig( + quant_mapping={ + "transformer": TorchAoConfig("float8wo") + } + ) + else: logger.warning(f"Unknown quantization method: {quant_method}. Loading without quantization.") return None @@ -76,7 +101,7 @@ def load_image_model(model_name, dtype='bfloat16', attn_backend='sdpa', cpu_offl attn_backend: 'sdpa', 'flash_attention_2', or 'flash_attention_3' cpu_offload: Enable CPU offloading for low VRAM compile_model: Compile the model for faster inference (slow first run) - quant_method: Quantization method - 'none', 'bnb-8bit', 'bnb-4bit' + quant_method: 'none', 'bnb-8bit', 'bnb-4bit', or torchao options (int8wo, fp4, float8wo) """ import torch from diffusers import DiffusionPipeline diff --git a/modules/shared.py b/modules/shared.py index a96cd70c..4a0fd986 100644 --- a/modules/shared.py +++ b/modules/shared.py @@ -60,7 +60,7 @@ group.add_argument('--image-attn-backend', type=str, default=None, choices=['sdp group.add_argument('--image-cpu-offload', action='store_true', help='Enable CPU offloading for image model.') group.add_argument('--image-compile', action='store_true', help='Compile the image model for faster inference.') group.add_argument('--image-quant', type=str, default=None, - choices=['none', 'bnb-8bit', 'bnb-4bit'], + choices=['none', 'bnb-8bit', 'bnb-4bit', 'torchao-int8wo', 'torchao-fp4', 'torchao-float8wo'], help='Quantization method for image model.') # Model loader diff --git a/modules/ui_image_generation.py b/modules/ui_image_generation.py index 8ec879e9..1b0e58f3 100644 --- a/modules/ui_image_generation.py +++ b/modules/ui_image_generation.py @@ -473,9 +473,9 @@ def create_ui(): with gr.Column(): shared.gradio['image_quant'] = gr.Dropdown( label='Quantization', - choices=['none', 'bnb-8bit', 'bnb-4bit'], + choices=['none', 'bnb-8bit', 'bnb-4bit', 'torchao-int8wo', 'torchao-fp4', 'torchao-float8wo'], value=shared.settings['image_quant'], - info='Quantization method for reduced VRAM usage. Quanto supports lower precisions (2-bit, 4-bit, 8-bit).' + info='BnB: bitsandbytes quantization. torchao: int8wo, fp4, float8wo.' ) shared.gradio['image_dtype'] = gr.Dropdown( diff --git a/requirements/full/requirements.txt b/requirements/full/requirements.txt index a031bbab..4e02d76f 100644 --- a/requirements/full/requirements.txt +++ b/requirements/full/requirements.txt @@ -25,6 +25,7 @@ safetensors==0.6.* scipy sentencepiece tensorboard +torchao==0.14.* transformers==4.57.* triton-windows==3.5.1.post21; platform_system == "Windows" tqdm diff --git a/requirements/full/requirements_amd.txt b/requirements/full/requirements_amd.txt index ee8d67ac..c0fbd9ab 100644 --- a/requirements/full/requirements_amd.txt +++ b/requirements/full/requirements_amd.txt @@ -23,6 +23,7 @@ safetensors==0.6.* scipy sentencepiece tensorboard +torchao==0.14.* transformers==4.57.* triton-windows==3.5.1.post21; platform_system == "Windows" tqdm diff --git a/requirements/full/requirements_amd_noavx2.txt b/requirements/full/requirements_amd_noavx2.txt index 7b36b151..b330646a 100644 --- a/requirements/full/requirements_amd_noavx2.txt +++ b/requirements/full/requirements_amd_noavx2.txt @@ -23,6 +23,7 @@ safetensors==0.6.* scipy sentencepiece tensorboard +torchao==0.14.* transformers==4.57.* triton-windows==3.5.1.post21; platform_system == "Windows" tqdm diff --git a/requirements/full/requirements_apple_intel.txt b/requirements/full/requirements_apple_intel.txt index 4f72d5ac..c2c64337 100644 --- a/requirements/full/requirements_apple_intel.txt +++ b/requirements/full/requirements_apple_intel.txt @@ -23,6 +23,7 @@ safetensors==0.6.* scipy sentencepiece tensorboard +torchao==0.14.* transformers==4.57.* triton-windows==3.5.1.post21; platform_system == "Windows" tqdm diff --git a/requirements/full/requirements_apple_silicon.txt b/requirements/full/requirements_apple_silicon.txt index 7942b9f0..10889215 100644 --- a/requirements/full/requirements_apple_silicon.txt +++ b/requirements/full/requirements_apple_silicon.txt @@ -23,6 +23,7 @@ safetensors==0.6.* scipy sentencepiece tensorboard +torchao==0.14.* transformers==4.57.* triton-windows==3.5.1.post21; platform_system == "Windows" tqdm diff --git a/requirements/full/requirements_cpu_only.txt b/requirements/full/requirements_cpu_only.txt index 96013c35..ffcd6473 100644 --- a/requirements/full/requirements_cpu_only.txt +++ b/requirements/full/requirements_cpu_only.txt @@ -23,6 +23,7 @@ safetensors==0.6.* scipy sentencepiece tensorboard +torchao==0.14.* transformers==4.57.* triton-windows==3.5.1.post21; platform_system == "Windows" tqdm diff --git a/requirements/full/requirements_cpu_only_noavx2.txt b/requirements/full/requirements_cpu_only_noavx2.txt index 0813f06c..daa2444f 100644 --- a/requirements/full/requirements_cpu_only_noavx2.txt +++ b/requirements/full/requirements_cpu_only_noavx2.txt @@ -23,6 +23,7 @@ safetensors==0.6.* scipy sentencepiece tensorboard +torchao==0.14.* transformers==4.57.* triton-windows==3.5.1.post21; platform_system == "Windows" tqdm diff --git a/requirements/full/requirements_noavx2.txt b/requirements/full/requirements_noavx2.txt index ea7edf6c..a43a7724 100644 --- a/requirements/full/requirements_noavx2.txt +++ b/requirements/full/requirements_noavx2.txt @@ -25,6 +25,7 @@ safetensors==0.6.* scipy sentencepiece tensorboard +torchao==0.14.* transformers==4.57.* triton-windows==3.5.1.post21; platform_system == "Windows" tqdm diff --git a/requirements/full/requirements_nowheels.txt b/requirements/full/requirements_nowheels.txt index 15247d72..ac24f0cb 100644 --- a/requirements/full/requirements_nowheels.txt +++ b/requirements/full/requirements_nowheels.txt @@ -23,6 +23,7 @@ safetensors==0.6.* scipy sentencepiece tensorboard +torchao==0.14.* transformers==4.57.* triton-windows==3.5.1.post21; platform_system == "Windows" tqdm