Merge branch 'main' of https://github.com/PKU-YuanGroup/Helios

2026-03-06 17:33:50 +08:00
parent a1cc04ae1a a7f32912f5
commit 9d5ac12e92
1 changed files with 5 additions and 3 deletions
@@ -31,8 +31,10 @@ pipe = HeliosPyramidPipeline.from_pretrained(
 )

 pipe.to("cuda")
-# pipe.transformer.set_attention_backend("flash_hub")
-pipe.transformer.set_attention_backend("_flash_3_hub")
+try:
+    pipe.transformer.set_attention_backend("_flash_3_hub")
+except Exception:
+    pipe.transformer.set_attention_backend("flash_hub")

 # @spaces.GPU(duration=1500)
 # def compile_transformer():
@@ -260,4 +262,4 @@ with gr.Blocks(css=CSS, title="Helios Video Generation", theme=gr.themes.Soft())

 if __name__ == "__main__":
    # demo.launch(share=True, allowed_paths=["./examples"])
-    demo.launch(share=True)
+    demo.launch(share=True)