Custom VAE loading support for Lumina 2

CHANGELOG.md

@@ -5,7 +5,7 @@
 - **Feature**  
   - Support for Python 3.13  
   - TeaCache support for Lumina 2  
-  - Custom UNet loading support for Lumina 2  
+  - Custom UNet and VAE loading support for Lumina 2  
 
 - **Changes**  
   - Increase the medvram mode threshold from 8GB to 12GB  
@@ -44,6 +44,7 @@
   - VAE Tiling with non-default tile sizes  
   - Lumina 2 with IPEX  
   - Nunchaku updated repo  
+  - Double loading of models with custom UNets  
 
 ## Update for 2025-06-02
 

modules/model_lumina.py

@@ -22,7 +22,7 @@ def load_lumina(_checkpoint_info, diffusers_load_config={}):
 
 def load_lumina2(checkpoint_info, diffusers_load_config={}):
     from modules import shared, devices, sd_models, model_quant
-    transformer, text_encoder = None, None
+    transformer, text_encoder, vae = None, None, None
     repo_id = sd_models.path_to_repo(checkpoint_info.name)
     if os.path.isdir(checkpoint_info.filename) and not repo_exists(repo_id):
         repo_id = checkpoint_info.filename
@@ -51,6 +51,21 @@ def load_lumina2(checkpoint_info, diffusers_load_config={}):
             if debug:
                 errors.display(e, 'Lumina2 UNet:')
 
+    if shared.opts.sd_vae != 'Default' and shared.opts.sd_vae != 'Automatic':
+        try:
+            debug(f'Load model: type=Lumina2 vae="{shared.opts.sd_vae}"')
+            from modules import sd_vae
+            # vae = sd_vae.load_vae_diffusers(None, sd_vae.vae_dict[shared.opts.sd_vae], 'override')
+            vae_file = sd_vae.vae_dict[shared.opts.sd_vae]
+            if os.path.exists(vae_file):
+                vae_config = os.path.join('configs', 'flux', 'vae', 'config.json')
+                vae = diffusers.AutoencoderKL.from_single_file(vae_file, config=vae_config, **diffusers_load_config)
+        except Exception as e:
+            shared.log.error(f"Load model: type=Lumina2 failed to load VAE: {e}")
+            shared.opts.sd_vae = 'Default'
+            if debug:
+                errors.display(e, 'Lumina2 VAE:')
+
     if transformer is None:
         transformer = diffusers.Lumina2Transformer2DModel.from_pretrained(
             repo_id,
@@ -70,6 +85,8 @@ def load_lumina2(checkpoint_info, diffusers_load_config={}):
     )
 
     load_config, quant_config = model_quant.get_dit_args(diffusers_load_config, allow_quant=False)
+    if vae is not None:
+        load_config['vae'] = vae
     pipe = diffusers.Lumina2Pipeline.from_pretrained(
         repo_id,
         cache_dir=shared.opts.diffusers_dir,

modules/sd_models.py

@@ -591,7 +591,7 @@ def load_diffuser(checkpoint_info=None, already_loaded_state_dict=None, timer=No
         if "Kandinsky" in sd_model.__class__.__name__: # need a special case
             sd_model.scheduler.name = 'DDIM'
 
-        if model_type not in ['Stable Cascade']: # need a special-case
+        if hasattr(sd_model, "unet") and model_type not in ['Stable Cascade']: # others calls load_diffuser again
             sd_unet.load_unet(sd_model)
 
         add_noise_pred_to_diffusers_callback(sd_model)