From 6ab65318eefc997617ae71ee6409fb84aa91d704 Mon Sep 17 00:00:00 2001 From: Vladimir Mandic Date: Thu, 12 Sep 2024 17:51:15 -0400 Subject: [PATCH] fix auraflow --- CHANGELOG.md | 6 ++++-- modules/model_auraflow.py | 10 ++++++---- modules/sd_models.py | 6 +++++- 3 files changed, 15 insertions(+), 7 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 5b859333f..be0e5ff63 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -108,8 +108,9 @@ Examples: e.g. if generating 1024px image, reducing layers by 1 will result in preview being 512px set via *settings -> live preview -> taesd decode layers* - **xhinker** prompt parser handle offloaded models -- **control** better handle offloading -- **control** upscale will use resize-to if set to non-zero values over resize-by +- **control** better handle offloading +- **upscale** will use resize-to if set to non-zero values over resize-by + applies to any upscale options, including refine workflow - **networks** add option to choose if mouse-over on network should attempt to fetch additional info option:`extra_networks_fetch` enable/disable in *settings -> networks* - speed up some garbage collection ops @@ -131,6 +132,7 @@ Examples: - **modernui** update - **zluda** update to 3.8.4, thanks @lshqqytiger - update `requirements` +- fix **AuraFlow** - fix handling of model configs if offline config is not available - fix vae decode in backend original - fix model path typos diff --git a/modules/model_auraflow.py b/modules/model_auraflow.py index 344e6558a..6f18bf13c 100644 --- a/modules/model_auraflow.py +++ b/modules/model_auraflow.py @@ -1,15 +1,17 @@ +import os import torch import diffusers +from modules import shared, sd_models, devices -repo_id = 'fal/AuraFlow' +debug = shared.log.trace if os.environ.get('SD_LOAD_DEBUG', None) is not None else lambda *args, **kwargs: None -def load_auraflow(_checkpoint_info, diffusers_load_config={}): - from modules import shared, devices +def load_auraflow(checkpoint_info, diffusers_load_config={}): + repo_id = sd_models.path_to_repo(checkpoint_info.name) if 'torch_dtype' not in diffusers_load_config: diffusers_load_config['torch_dtype'] = torch.float16 - + debug(f'Loading AuraFlow: repo="{repo_id}" config={diffusers_load_config}') pipe = diffusers.AuraFlowPipeline.from_pretrained( repo_id, cache_dir = shared.opts.diffusers_dir, diff --git a/modules/sd_models.py b/modules/sd_models.py index 6958e3ea0..b7474a9fa 100644 --- a/modules/sd_models.py +++ b/modules/sd_models.py @@ -1236,7 +1236,7 @@ def load_diffuser(checkpoint_info=None, already_loaded_state_dict=None, timer=No if model_config is not None: if debug_load: shared.log.debug(f'Model config: path="{model_config}"') - diffusers_load_config['config_files'] = model_config + diffusers_load_config['config'] = model_config if model_type.startswith('Stable Diffusion 3'): from modules.model_sd3 import load_sd3 sd_model = load_sd3(fn=checkpoint_info.path, cache_dir=shared.opts.diffusers_dir, config=diffusers_load_config.get('config', None)) @@ -1566,6 +1566,10 @@ def set_diffusers_attention(pipe): module.set_attn_processor(p.FluxAttnProcessor2_0()) elif module.__class__.__name__ in ['HunyuanDiT2DModel']: module.set_attn_processor(p.HunyuanAttnProcessor2_0()) + elif module.__class__.__name__ in ['AuraFlowTransformer2DModel']: + module.set_attn_processor(p.AuraFlowAttnProcessor2_0()) + elif 'Transformer' in module.__class__.__name__: + pass # unknown transformer so probably dont want to force attention processor else: module.set_attn_processor(attention)