[Enhancement] Extend quantization support to InfuseNet via --infusenet_quant (NF4,FP8,INT8,INT4)

pipelines/pipeline_infu_flux.py

@@ -26,7 +26,13 @@
 from insightface.app import FaceAnalysis
 from insightface.utils import face_align
 from PIL import Image
-from optimum.quanto import freeze, qint8, quantize
+from optimum.quanto import freeze, qfloat8, qint4, qint8, quantize
+try:
+    import bitsandbytes as bnb
+    from bitsandbytes.nn import Linear4bit
+    BNB_AVAILABLE = True
+except ImportError:
+    BNB_AVAILABLE = False
 from transformers import T5EncoderModel
 
 from .pipeline_flux_infusenet import FluxInfuseNetPipeline
@@ -133,6 +139,7 @@ def __init__(
             infu_flux_version='v1.0',
             model_version='aes_stage2',
             quantize_8bit=False,
+            quantize_infusenet=None,
             cpu_offload=False,
         ):
 
@@ -150,7 +157,43 @@ def __init__(
             infusenet_path = os.path.join(infu_model_path, 'InfuseNetModel')
             self.infusenet = FluxControlNetModel.from_pretrained(infusenet_path, torch_dtype=torch.bfloat16)
             insightface_root_path = './models/InfiniteYou/supports/insightface'
-        if quantize_8bit:
+        # Quantize InfuseNet (independent of FLUX quantization)
+        # quantize_infusenet options:
+        #   'nf4'  : bitsandbytes NF4 4-bit (true peak VRAM reduction, requires pip install bitsandbytes)
+        #   'fp8'  : optimum.quanto FP8 weight-only
+        #   'int4' : optimum.quanto INT4 weight-only
+        #   'int8' : optimum.quanto INT8 weight-only
+        if quantize_infusenet == 'nf4':
+            if not BNB_AVAILABLE:
+                raise ImportError(
+                    '--infusenet_quant nf4 requires bitsandbytes. '
+                    'Install with: pip install bitsandbytes'
+                )
+            print('[InfuseNet] Applying NF4 quantization via bitsandbytes...')
+            from diffusers import BitsAndBytesConfig as DiffusersBnBConfig
+            bnb_config = DiffusersBnBConfig(
+                load_in_4bit=True,
+                bnb_4bit_quant_type='nf4',
+                bnb_4bit_compute_dtype=torch.bfloat16,
+                bnb_4bit_use_double_quant=True,
+            )
+            # Reload InfuseNet with BnB quantization config
+            infusenet_path = os.path.join(infu_model_path, 'InfuseNetModel')
+            self.infusenet = FluxControlNetModel.from_pretrained(
+                infusenet_path,
+                quantization_config=bnb_config,
+                torch_dtype=torch.bfloat16,
+            )
+        elif quantize_infusenet == 'fp8':
+            print('[InfuseNet] Applying FP8 quantization via optimum.quanto...')
+            quantize(self.infusenet, weights=qfloat8)
+            freeze(self.infusenet)
+        elif quantize_infusenet == 'int4':
+            print('[InfuseNet] Applying INT4 quantization via optimum.quanto...')
+            quantize(self.infusenet, weights=qint4)
+            freeze(self.infusenet)
+        elif quantize_infusenet == 'int8' or quantize_8bit:
+            print('[InfuseNet] Applying INT8 quantization via optimum.quanto...')
             quantize(self.infusenet, weights=qint8)
             freeze(self.infusenet)
         try:

test.py

@@ -44,6 +44,14 @@ def main():
     # Memory reduction options
     parser.add_argument('--quantize_8bit', action='store_true')
     parser.add_argument('--cpu_offload', action='store_true')
+    parser.add_argument('--infusenet_quant', default=None, choices=['nf4', 'fp8', 'int4', 'int8'],
+                        help="""Quantize InfuseNet independently of FLUX: nf4 | fp8 | int4 | int8.
+Approximate peak VRAM savings vs bf16 full (~43GB):
+  int8 : ~32GB  (same as --quantize_8bit but InfuseNet only)
+  fp8  : ~32GB  (slightly lower precision loss than int8)
+  int4 : ~26GB  (most aggressive, may affect ID similarity slightly)
+Combine with --quantize_8bit and --cpu_offload for maximum savings.
+Requires: pip install optimum-quanto""")
     args = parser.parse_args()
 
     # Check arguments
@@ -63,6 +71,7 @@ def main():
         infu_flux_version=args.infu_flux_version,
         model_version=args.model_version,
         quantize_8bit=args.quantize_8bit,
+        quantize_infusenet=args.infusenet_quant,
         cpu_offload=args.cpu_offload,
     )
     # Load LoRAs (optional)