fix flux unet override load with bnb

Signed-off-by: Vladimir Mandic <mandic00@live.com>

CHANGELOG.md

@@ -12,11 +12,13 @@
 - **VAE**: tiling granular options in *settings -> variable auto encoder*  
 - **UI**:
   - live preview optimizations and error handling  
-  - live preview high quality for flow models, thanks @Disty0  
+  - live preview high quality output, thanks @Disty0  
   - CSS optimizations when log view is disabled  
 - **Samplers**:
   - add flow shift options and separate dynamic thresholding from dynamic shifting  
   - autodetect matching sigma capabilities  
+- **API**
+  - better default values for generate  
 - **Refactor**:
   - remove all LDM imports if running in native mode  
 - **Fixes**  
@@ -29,7 +31,8 @@
   - xyz grid recover on error  
   - hires batch  
   - sdxl refiner  
-  - kandinsky   
+  - kandinsky  
+  - flux custom unet loader for bnb  
 
 ## Update for 2024-12-24
 

modules/model_flux.py

@@ -212,7 +212,8 @@ def load_transformer(file_path): # triggered by opts.sd_unet change
         _transformer, _text_encoder_2 = load_flux_quanto(file_path)
         if _transformer is not None:
             transformer = _transformer
-    elif quant == 'fp8' or quant == 'fp4' or quant == 'nf4':
+    elif quant == 'fp8' or quant == 'fp4' or quant == 'nf4' or 'Model' in shared.opts.bnb_quantization:
+        print('HERE0')
         _transformer, _text_encoder_2 = load_flux_bnb(file_path, diffusers_load_config)
         if _transformer is not None:
             transformer = _transformer
@@ -222,9 +223,18 @@ def load_transformer(file_path): # triggered by opts.sd_unet change
         if _transformer is not None:
             transformer = _transformer
     else:
-        diffusers_load_config = model_quant.create_bnb_config(diffusers_load_config)
-        diffusers_load_config = model_quant.create_ao_config(diffusers_load_config)
-        transformer = diffusers.FluxTransformer2DModel.from_single_file(file_path, **diffusers_load_config)
+        print('HERE1')
+        quant_args = {}
+        quant_args = model_quant.create_bnb_config(quant_args)
+        if quant_args:
+            model_quant.load_bnb(f'Load model: type=Sana quant={quant_args}')
+        if not quant_args:
+            quant_args = model_quant.create_ao_config(quant_args)
+            if quant_args:
+                model_quant.load_torchao(f'Load model: type=Sana quant={quant_args}')
+        print('HERE2', diffusers_load_config)
+        print('HERE3', quant_args)
+        transformer = diffusers.FluxTransformer2DModel.from_single_file(file_path, **diffusers_load_config, **quant_args)
     if transformer is None:
         shared.log.error('Failed to load UNet model')
         shared.opts.sd_unet = 'None'

modules/processing_class.py

@@ -29,7 +29,7 @@ class StableDiffusionProcessing:
                  seed_resize_from_w: int = -1,
                  batch_size: int = 1,
                  n_iter: int = 1,
-                 steps: int = 50,
+                 steps: int = 20,
                  clip_skip: int = 1,
                  width: int = 1024,
                  height: int = 1024,
@@ -39,7 +39,7 @@ class StableDiffusionProcessing:
                  hr_sampler_name: str = None,
                  eta: float = None,
                  # guidance
-                 cfg_scale: float = 7.0,
+                 cfg_scale: float = 6.0,
                  cfg_end: float = 1,
                  diffusers_guidance_rescale: float = 0.7,
                  pag_scale: float = 0.0,