From c0da835b5afdc27ddd9fb5339a8afd2cf4f6d691 Mon Sep 17 00:00:00 2001 From: Shenghai Yuan Date: Thu, 5 Mar 2026 11:01:51 +0800 Subject: [PATCH] fix flash attn --- infer_helios.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/infer_helios.py b/infer_helios.py index 23c730b..4c06c4d 100644 --- a/infer_helios.py +++ b/infer_helios.py @@ -214,7 +214,10 @@ def main(): transformer = replace_rmsnorm_with_fp32(transformer) transformer = replace_all_norms_with_flash_norms(transformer) replace_rope_with_flash_rope() - transformer.set_attention_backend("_flash_3_hub") + try: + transformer.set_attention_backend("_flash_3_hub") + except Exception: + transformer.set_attention_backend("flash_hub") vae = AutoencoderKLWan.from_pretrained( args.base_model_path,