experiments with repa-e

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

CHANGELOG.md

@@ -7,6 +7,7 @@
 - Reorganization of **Reference Models** into *Base, Quantized, Distilled and Community* sections for easier navigation
 - New **Tencent HunyuanImage 2.1** model capable of generating 2K images natively
 - New **offline mode** to use previously downloaded models without internet connection  
+- New SOTA model loader using **Run:ai**  
 - Fixes, fixes, fixes... too many to list here!
 
 ### Details for 2025-10-26
@@ -31,7 +32,7 @@
   - **model load**: SOTA method using nVidia's [Run:ai streamer](https://github.com/run-ai/runai-model-streamer)  
     enable in *settings -> model options -> runai streamer*  
     applies to *diffusers, transformers and sdnq* loaders  
-    *experimental* but shows significant model load speedups  
+    *experimental* but shows significant model load speedups, 20-40% depending on model and hardware  
 - **Backend**
   - switch to `torch==2.9` for *ipex, rocm and openvino*  
   - switch to `rocm==7.0` for nightlies  

installer.py

@@ -613,7 +613,7 @@ def check_diffusers():
     t_start = time.time()
     if args.skip_all:
         return
-    sha = '7536f647e4144c7acaf9e140893ff7edb85bf9a3' # diffusers commit hash
+    sha = '6d1a6486024192951ce696e8f4cf79a39509182f' # diffusers commit hash
     # if args.use_rocm or args.use_zluda or args.use_directml:
     #     sha = '043ab2520f6a19fce78e6e060a68dbc947edb9f9' # lock diffusers versions for now
     pkg = pkg_resources.working_set.by_key.get('diffusers', None)

modules/processing.py

@@ -163,6 +163,9 @@ def process_images(p: StableDiffusionProcessing) -> Processed:
             if p.override_settings.get('sd_vae', None) == 'TAESD':
                 p.vae_type = 'Tiny'
                 p.override_settings.pop('sd_vae', None)
+            if p.override_settings.get('sd_vae', None) == 'REPA-E':
+                p.vae_type = 'Repa'
+                p.override_settings.pop('sd_vae', None)
         if p.override_settings.get('Hires upscaler', None) is not None:
             p.enable_hr = True
         if len(p.override_settings.keys()) > 0:

modules/processing_info.py

@@ -78,6 +78,8 @@ def create_infotext(p: StableDiffusionProcessing, all_prompts=None, all_seeds=No
         args["VAE"] = (None if not shared.opts.add_model_name_to_info or sd_vae.loaded_vae_file is None else os.path.splitext(os.path.basename(sd_vae.loaded_vae_file))[0])
     elif p.vae_type == 'Tiny':
         args["VAE"] = 'TAESD'
+    elif p.vae_type == 'REPA-E':
+        args["VAE"] = 'REPA-E'
     elif p.vae_type == 'Remote':
         args["VAE"] = 'Remote'
     if grid is None and (p.n_iter > 1 or p.batch_size > 1) and index >= 0:

modules/processing_vae.py

@@ -275,15 +275,6 @@ def vae_decode(latents, model, output_type='np', vae_type='Full', width=None, he
     if latents is None or not torch.is_tensor(latents): # already decoded
         return latents
 
-    if vae_type == 'Remote':
-        jobid = shared.state.begin('Remote VAE')
-        from modules.sd_vae_remote import remote_decode
-        tensors = remote_decode(latents=latents, width=width, height=height)
-        shared.state.end(jobid)
-        if tensors is not None and len(tensors) > 0:
-            return vae_postprocess(tensors, model, output_type)
-
-    jobid = shared.state.begin('VAE Decode')
     if latents.shape[0] == 0:
         shared.log.error(f'VAE nothing to decode: {latents.shape}')
         return []
@@ -293,6 +284,21 @@ def vae_decode(latents, model, output_type='np', vae_type='Full', width=None, he
         shared.log.error('VAE not found in model')
         return []
 
+    if vae_type == 'Remote':
+        jobid = shared.state.begin('Remote VAE')
+        from modules.sd_vae_remote import remote_decode
+        tensors = remote_decode(latents=latents, width=width, height=height)
+        shared.state.end(jobid)
+        if tensors is not None and len(tensors) > 0:
+            return vae_postprocess(tensors, model, output_type)
+    if vae_type == 'Repa':
+        from modules.sd_vae_repa import repa_load
+        vae = repa_load(latents)
+        vae_type = 'Full'
+        if vae is not None:
+            model.vae = vae
+
+    jobid = shared.state.begin('VAE Decode')
     if hasattr(model, '_unpack_latents') and hasattr(model, 'transformer_spatial_patch_size') and frames is not None: # LTX
         latent_num_frames = (frames - 1) // model.vae_temporal_compression_ratio + 1
         latents = model._unpack_latents(latents.unsqueeze(0), latent_num_frames, height // 32, width // 32, model.transformer_spatial_patch_size, model.transformer_temporal_patch_size) # pylint: disable=protected-access

modules/sd_vae_repa.py

@@ -0,0 +1,32 @@
+import diffusers
+from modules import shared
+
+
+models = {
+    'sd': { 'repo_id': 'REPA-E/e2e-sdvae-hf', 'cls': 'AutoencoderKL' },
+    'sdxl': { 'repo_id': 'REPA-E/e2e-sdvae-hf', 'cls': 'AutoencoderKL' },
+    'sd3': { 'repo_id': 'REPA-E/e2e-sd3.5-vae', 'cls': 'AutoencoderKL' },
+    'f1': { 'repo_id': 'REPA-E/e2e-flux-vae', 'cls': 'AutoencoderKL' },
+    'qwen': { 'repo_id': 'REPA-E/e2e-qwenimage-vae', 'cls': 'AutoencoderKLQwenImage' },
+}
+loaded_cls = None
+loaded_vae = None
+
+
+def repa_load(latents):
+    global loaded_cls, loaded_vae # pylint: disable=global-statement
+    config = models.get(shared.sd_model_type, None)
+    if config is None:
+        shared.log.error(f'Decode: type="repa" model={shared.sd_model_type} not supported')
+        return latents
+
+    cls = getattr(diffusers, config['cls'])
+    if (cls != loaded_cls) or (loaded_vae is None):
+        shared.log.info(f'RePA VAE load: {config["repo_id"]} cls={config["cls"]}')
+        loaded_vae = cls.from_pretrained(
+            config['repo_id'],
+            torch_dtype=latents.dtype,
+            cache_dir=shared.opts.hfcache_dir,
+            )
+        loaded_cls = cls
+    return loaded_vae