fix flash attn

2026-03-05 11:01:51 +08:00
parent 8ac4293d61
commit c0da835b5a
1 changed files with 4 additions and 1 deletions
@@ -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,