Fix sth

app.py

@@ -1,43 +1,13 @@
-import os
-import subprocess
-import sys
-import time
 import tempfile
-import zipfile
-import torch
-
-# ---------------------------------------------------------------------------
-# Install private diffusers fork
-# ---------------------------------------------------------------------------
-_APP_DIR = os.path.dirname(os.path.abspath(__file__))
-ZIP_PATH = os.path.join(_APP_DIR, "helios_diffusers.zip")
-EXTRACT_DIR = os.path.join(_APP_DIR, "_helios_diffusers")
-_PKG_ROOT = os.path.join(EXTRACT_DIR, "diffusers-new-model-addition-helios-helios")
-
-if not os.path.isdir(_PKG_ROOT):
-    print(f"[setup] Extracting {ZIP_PATH}")
-    with zipfile.ZipFile(ZIP_PATH, "r") as zf:
-        zf.extractall(EXTRACT_DIR)
-
-print(f"[setup] Installing diffusers from {_PKG_ROOT}")
-try:
-    subprocess.check_call([sys.executable, "-m", "pip", "install", _PKG_ROOT])
-except subprocess.CalledProcessError as e:
-    print(f"[setup] pip install failed (exit {e.returncode}), falling back to sys.path")
-
-_SRC_DIR = os.path.join(_PKG_ROOT, "src")
-if os.path.isdir(_SRC_DIR):
-    sys.path.insert(0, _SRC_DIR)
+import time
 
 import gradio as gr
 import spaces
-from diffusers import (
-    AutoencoderKLWan,
-    HeliosPyramidPipeline,
-    HeliosDMDScheduler
-)
+import torch
+
+from diffusers import AutoencoderKLWan, HeliosDMDScheduler, HeliosPyramidPipeline
 from diffusers.utils import export_to_video, load_image, load_video
-from aoti import aoti_load_
+
 
 # ---------------------------------------------------------------------------
 # Pre-load model
@@ -47,18 +17,30 @@ MODEL_ID = "BestWishYsh/Helios-Distilled"
 vae = AutoencoderKLWan.from_pretrained(MODEL_ID, subfolder="vae", torch_dtype=torch.float32)
 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
+    MODEL_ID, vae=vae, scheduler=scheduler, torch_dtype=torch.bfloat16, is_distilled=True
 )
 
-# aoti_load_(pipe.transformer, "multimodalart/helios-distilled-transformer", "helios_distilled_transformer.pt2")
-
 pipe.to("cuda")
+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():
+#     with spaces.aoti_capture(pipe.transformer) as call:
+#         pipe("arbitrary example prompt")
+
+#     exported = torch.export.export(
+#         pipe.transformer,
+#         args=call.args,
+#         kwargs=call.kwargs,
+#     )
+#     return spaces.aoti_compile(exported)
+
+# compiled_transformer = compile_transformer()
+# spaces.aoti_apply(compiled_transformer, pipe.transformer)
 
-pipe.transformer.set_attention_backend("_flash_3_hub")
 
 # ---------------------------------------------------------------------------
 # Generation
@@ -113,6 +95,7 @@ def generate_video(
     info = f"Generated in {elapsed:.1f}s · {num_frames} frames · {height}×{width}"
     return tmp.name, info
 
+
 # ---------------------------------------------------------------------------
 # UI Setup
 # ---------------------------------------------------------------------------
@@ -124,17 +107,44 @@ def update_conditional_visibility(mode):
     else:
         return gr.update(visible=False), gr.update(visible=False)
 
+
 CSS = """
-#header { text-align: center; margin-bottom: 0.5em; }
-#header h1 { font-size: 2.2em; margin-bottom: 0; }
+#header { text-align: center; margin-bottom: 1.5em; }
+#header h1 { font-size: 2.2em; margin-bottom: 0.2em; }
+.logo { max-height: 100px; margin: 0 auto 10px auto; display: block; }
+.link-buttons { display: flex; justify-content: center; gap: 15px; margin-top: 10px; }
+.link-buttons a {
+    background-color: #2b3137;
+    color: #ffffff !important;
+    padding: 8px 20px;
+    border-radius: 6px;
+    text-decoration: none;
+    font-weight: 600;
+    font-size: 1em;
+    transition: all 0.2s ease-in-out;
+    box-shadow: 0 2px 4px rgba(0,0,0,0.1);
+}
+.link-buttons a:hover { background-color: #4a535c; transform: translateY(-1px); }
 .contain { max-width: 1350px; margin: 0 auto !important; }
 """
 
-with gr.Blocks(css=CSS, title="Helios Video Generation", theme=gr.themes.Soft()) as demo:
+with gr.Blocks(title="Helios Video Generation") as demo:
     gr.HTML(
         """
+        <div style='display: flex; align-items: center; justify-content: center; width: 100%;'>
+            <img src="https://github.com/PKU-YuanGroup/Helios-Page/blob/main/figures/logo_white.png?raw=true" style='width: 400px; height: auto;' />
+        </div>
         <div id="header">
-            <h1>🎬 Helios 14B distilled</h1>
+            <h1>🎬 Helios 14B Distilled: Real Real-Time Long Video Generation Model</h1>
+            <p style="font-size: 1.1em; color: #666; margin-top: 0.5em; margin-bottom: 1em;">
+                If you like our project, please give us a star ⭐ on GitHub for the latest update.
+            </p>
+            <div class="link-buttons">
+                <a href="https://github.com/PKU-YuanGroup/Helios" target="_blank">💻 Code</a>
+                <a href="https://pku-yuangroup.github.io/Helios-Page" target="_blank">📄 Page</a>
+                <a href="https://www.youtube.com/watch?v=vd_AgHtOUFQ" target="_blank">🎥 Main Feature</a>
+                <a href="https://www.youtube.com/watch?v=1GeIU2Dn7UY" target="_blank">⚡ Inference Speed</a>
+            </div>
         </div>
         """
     )
@@ -161,7 +171,7 @@ with gr.Blocks(css=CSS, title="Helios Video Generation", theme=gr.themes.Soft())
                     "of hard and soft corals in shades of red, orange, and green. The photo captures "
                     "the fish from a slightly elevated angle, emphasizing its lively movements and the "
                     "vivid colors of its surroundings. A close-up shot with dynamic movement."
-                )
+                ),
             )
             with gr.Accordion("Advanced Settings", open=False):
                 with gr.Row():
@@ -183,7 +193,18 @@ with gr.Blocks(css=CSS, title="Helios Video Generation", theme=gr.themes.Soft())
     mode.change(fn=update_conditional_visibility, inputs=[mode], outputs=[image_input, video_input])
     generate_btn.click(
         fn=generate_video,
-        inputs=[mode, prompt, image_input, video_input, height, width, num_frames, num_inference_steps, seed, is_amplify_first_chunk],
+        inputs=[
+            mode,
+            prompt,
+            image_input,
+            video_input,
+            height,
+            width,
+            num_frames,
+            num_inference_steps,
+            seed,
+            is_amplify_first_chunk,
+        ],
         outputs=[video_output, info_output],
     )
 
@@ -200,6 +221,8 @@ with gr.Blocks(css=CSS, title="Helios Video Generation", theme=gr.themes.Soft())
                 "of hard and soft corals in shades of red, orange, and green. The photo captures "
                 "the fish from a slightly elevated angle, emphasizing its lively movements and the "
                 "vivid colors of its surroundings. A close-up shot with dynamic movement.",
+                None,
+                None,
             ],
             [
                 "Text-to-Video",
@@ -211,6 +234,8 @@ with gr.Blocks(css=CSS, title="Helios Video Generation", theme=gr.themes.Soft())
                 "smile as if he found the answer to the mystery of life, the lighting is very cinematic "
                 "with the golden light and the Parisian streets and city in the background, depth of "
                 "field, cinematic 35mm film.",
+                None,
+                None,
             ],
             [
                 "Text-to-Video",
@@ -221,11 +246,25 @@ with gr.Blocks(css=CSS, title="Helios Video Generation", theme=gr.themes.Soft())
                 "of the Amalfi Coast Italy, several distant people are seen walking and enjoying vistas "
                 "on patios of the dramatic ocean views, the warm glow of the afternoon sun creates a "
                 "magical and romantic feeling to the scene, the view is stunning captured with beautiful photography.",
+                None,
+                None,
+            ],
+            [
+                "Image-to-Video",
+                "A towering emerald wave surges forward, its crest curling with raw power and energy. Sunlight glints off the translucent water, illuminating the intricate textures and deep green hues within the wave’s body. A thick spray erupts from the breaking crest, casting a misty veil that dances above the churning surface. As the perspective widens, the immense scale of the wave becomes apparent, revealing the restless expanse of the ocean stretching beyond. The scene captures the ocean’s untamed beauty and relentless force, with every droplet and ripple shimmering in the light. The dynamic motion and vivid colors evoke both awe and respect for nature’s might.",
+                "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/helios/wave.jpg",
+                None,
+            ],
+            [
+                "Video-to-Video",
+                "A bright yellow Lamborghini Huracn Tecnica speeds along a curving mountain road, surrounded by lush green trees under a partly cloudy sky. The car's sleek design and vibrant color stand out against the natural backdrop, emphasizing its dynamic movement. The road curves gently, with a guardrail visible on one side, adding depth to the scene. The motion blur captures the sense of speed and energy, creating a thrilling and exhilarating atmosphere. A front-facing shot from a slightly elevated angle, highlighting the car's aggressive stance and the surrounding greenery.",
+                None,
+                "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/helios/car.mp4",
             ],
         ],
-        inputs=[mode, prompt],
+        inputs=[mode, prompt, image_input, video_input],
         label="Example Prompts",
     )
 
 if __name__ == "__main__":
-    demo.launch()
+    demo.launch(share=True, css=CSS, theme=gr.themes.Soft())