add auraflow

modules/model_auraflow.py

@@ -0,0 +1,19 @@
+import torch
+import diffusers
+
+
+repo_id = 'fal/AuraFlow'
+
+
+def load_auraflow(_checkpoint_info, diffusers_load_config={}):
+    from modules import shared, devices
+    if 'torch_dtype' not in diffusers_load_config:
+        diffusers_load_config['torch_dtype'] = torch.float16
+
+    pipe = diffusers.AuraFlowPipeline.from_pretrained(
+        repo_id,
+        cache_dir = shared.opts.diffusers_dir,
+        **diffusers_load_config,
+    )
+    devices.torch_gc()
+    return pipe

modules/sd_models.py

@@ -615,6 +615,8 @@ def detect_pipeline(f: str, op: str = 'model', warning=True, quiet=False):
                 guess = 'Lumina-Next'
             if 'kolors' in f.lower():
                 guess = 'Kolors'
+            if 'auraflow' in f.lower():
+                guess = 'AuraFlow'
             # switch for specific variant
             if guess == 'Stable Diffusion' and 'inpaint' in f.lower():
                 guess = 'Stable Diffusion Inpaint'
@@ -1014,6 +1016,15 @@ def load_diffuser(checkpoint_info=None, already_loaded_state_dict=None, timer=No
                     if debug_load:
                         errors.display(e, 'Load')
                     return
+            elif model_type in ['AuraFlow']: # forced pipeline
+                try:
+                    from modules.model_auraflow import load_auraflow
+                    sd_model = load_auraflow(checkpoint_info, diffusers_load_config)
+                except Exception as e:
+                    shared.log.error(f'Diffusers Failed loading {op}: {checkpoint_info.path} {e}')
+                    if debug_load:
+                        errors.display(e, 'Load')
+                    return
             elif model_type in ['Stable Diffusion 3']:
                 try:
                     from modules.model_sd3 import load_sd3

modules/shared_items.py

@@ -71,9 +71,10 @@ def get_pipelines():
         'Kandinsky 3': getattr(diffusers, 'Kandinsky3Pipeline', None),
         'DeepFloyd IF': getattr(diffusers, 'IFPipeline', None),
         'Custom Diffusers Pipeline': getattr(diffusers, 'DiffusionPipeline', None),
-        'Kolors': getattr(diffusers, 'StableDiffusionXLPipeline', None),
         'InstaFlow': getattr(diffusers, 'StableDiffusionPipeline', None), # dynamically redefined and loaded in sd_models.load_diffuser
         'SegMoE': getattr(diffusers, 'StableDiffusionPipeline', None), # dynamically redefined and loaded in sd_models.load_diffuser
+        'Kolors': getattr(diffusers, 'KolorsPipeline', None),
+        'AuraFlow': getattr(diffusers, 'AuraFlowPipeline', None),
     }
     if hasattr(diffusers, 'OnnxStableDiffusionPipeline'):
         onnx_pipelines = {