Update app.py

app.py

@@ -38,6 +38,7 @@ import torch
 from diffusers import (
     AutoencoderKLWan,
     HeliosPyramidPipeline,
+    HeliosDMDScheduler
 )
 from diffusers.utils import export_to_video, load_image, load_video
 
@@ -47,8 +48,16 @@ from diffusers.utils import export_to_video, load_image, load_video
 MODEL_ID = "BestWishYsh/Helios-Distilled"
 
 vae = AutoencoderKLWan.from_pretrained(MODEL_ID, subfolder="vae", torch_dtype=torch.float32)
-pipe = HeliosPyramidPipeline.from_pretrained(MODEL_ID, vae=vae, torch_dtype=torch.bfloat16)
+scheduler = HeliosDMDScheduler.from_pretrained(MODEL_ID, subfolder="scheduler")
+pipe = HeliosPyramidPipeline.from_pretrained(
+    MODEL_ID, 
+    vae=vae, 
+    scheduler=scheduler,
+    torch_dtype=torch.bfloat16,
+    is_distilled=True
+)
 pipe.to("cuda")
+pipe.transformer.set_attention_backend("_flash_3_hub")
 
 DEFAULT_NEGATIVE_PROMPT = (
     "Bright tones, overexposed, static, blurred details, subtitles, style, works, "
@@ -66,7 +75,6 @@ DEFAULT_NEGATIVE_PROMPT = (
 def generate_video(
     mode: str,
     prompt: str,
-    negative_prompt: str,
     image_input,
     video_input,
     height: int,
@@ -85,7 +93,6 @@ def generate_video(
 
     kwargs = {
         "prompt": prompt,
-        "negative_prompt": negative_prompt,
         "height": int(height),
         "width": int(width),
         "num_frames": int(num_frames),
@@ -157,11 +164,6 @@ with gr.Blocks(css=CSS, title="Helios Video Generation", theme=gr.themes.Soft())
                     "Medium shot focusing on the train window and the rushing scenery beyond."
                 ),
             )
-            negative_prompt = gr.Textbox(
-                label="Negative Prompt",
-                lines=2,
-                value=DEFAULT_NEGATIVE_PROMPT,
-            )
 
             with gr.Accordion("Conditional Inputs (I2V / V2V)", open=False):
                 image_input = gr.Image(label="Image (for I2V)", type="filepath")
@@ -172,7 +174,7 @@ with gr.Blocks(css=CSS, title="Helios Video Generation", theme=gr.themes.Soft())
                     height = gr.Slider(128, 768, value=384, step=16, label="Height")
                     width = gr.Slider(128, 1280, value=640, step=16, label="Width")
                 with gr.Row():
-                    num_frames = gr.Slider(9, 240, value=99, step=1, label="Num Frames")
+                    num_frames = gr.Slider(9, 240, value=240, step=1, label="Num Frames")
                     num_inference_steps = gr.Slider(
                         1, 10, value=2, step=1, label="Steps (per pyramid stage)"
                     )
@@ -195,7 +197,6 @@ with gr.Blocks(css=CSS, title="Helios Video Generation", theme=gr.themes.Soft())
         inputs=[
             mode,
             prompt,
-            negative_prompt,
             image_input,
             video_input,
             height,