Update app.py

app.py

@@ -1,57 +1,91 @@
-import gradio as gr
-import numpy as np
-import random
-import torch
-import spaces
+# PyTorch 2.8 (temporary hack)
 import os
+os.system('pip install --upgrade --pre --extra-index-url https://download.pytorch.org/whl/nightly/cu126 "torch<2.9" spaces')
 
-from PIL import Image
-from diffusers.pipelines.wan.pipeline_wan import WanPipeline
+# Actual demo code
+import spaces
+import torch
+from diffusers import WanPipeline, AutoencoderKLWan
+from diffusers.models.transformers.transformer_wan import WanTransformer3DModel
 from diffusers.utils.export_utils import export_to_video
-
+import gradio as gr
 import tempfile
-from typing import Optional, Tuple, Any
-from groq import Groq
+import numpy as np
+from PIL import Image
+import random
+import gc
+from optimization import optimize_pipeline_
 
-# =========================================================
-# CUDA BACKEND FIX
-# =========================================================
-# Fix for cusolver error
-torch.backends.cuda.preferred_linalg_library("cusolver")
 
 # =========================================================
 # API CONFIGURATION
 # =========================================================
 GROQ_API_KEY = os.getenv("GROQ_API_KEY")
-MODEL_ID = os.getenv("MODEL_ID")
-HF_TOKEN = os.getenv("HF_TOKEN")
+
+# =========================================================
+# MODEL CONFIGURATION
+# =========================================================
+MODEL_ID = "Wan-AI/Wan2.2-T2V-A14B-Diffusers"
+
+LANDSCAPE_WIDTH = 832
+LANDSCAPE_HEIGHT = 480
+MAX_SEED = np.iinfo(np.int32).max
+
+FIXED_FPS = 16
+MIN_FRAMES_MODEL = 8
+MAX_FRAMES_MODEL = 81
+
+MIN_DURATION = round(MIN_FRAMES_MODEL/FIXED_FPS, 1)
+MAX_DURATION = round(MAX_FRAMES_MODEL/FIXED_FPS, 1)
 
 # =========================================================
 # MODEL LOADING
 # =========================================================
-dtype = torch.bfloat16
-device = "cuda" if torch.cuda.is_available() else "cpu"
+vae = AutoencoderKLWan.from_pretrained(
+    "Wan-AI/Wan2.2-T2V-A14B-Diffusers", 
+    subfolder="vae", 
+    torch_dtype=torch.float32
+)
 
-print("Loading pipeline...")
 pipe = WanPipeline.from_pretrained(
     MODEL_ID,
-    torch_dtype=dtype,
-    token=HF_TOKEN
-).to(device)
-
-MAX_SEED = np.iinfo(np.int32).max
-FIXED_FPS = 24
-MIN_FRAMES = 9  # Must be (4k + 1) format: 9, 13, 17, 21, 25, ...
-MAX_FRAMES = 81
+    transformer=WanTransformer3DModel.from_pretrained(
+        'linoyts/Wan2.2-T2V-A14B-Diffusers-BF16',
+        subfolder='transformer',
+        torch_dtype=torch.bfloat16,
+        device_map='cuda',
+    ),
+    transformer_2=WanTransformer3DModel.from_pretrained(
+        'linoyts/Wan2.2-T2V-A14B-Diffusers-BF16',
+        subfolder='transformer_2',
+        torch_dtype=torch.bfloat16,
+        device_map='cuda',
+    ),
+    vae=vae,
+    torch_dtype=torch.bfloat16,
+).to('cuda')
+
+for i in range(3): 
+    gc.collect()
+    torch.cuda.synchronize() 
+    torch.cuda.empty_cache()
+
+optimize_pipeline_(
+    pipe,
+    prompt='prompt',
+    height=LANDSCAPE_HEIGHT,
+    width=LANDSCAPE_WIDTH,
+    num_frames=MAX_FRAMES_MODEL,
+)
 
 # =========================================================
 # DEFAULT PROMPTS
 # =========================================================
-default_prompt = "A beautiful sunset over the ocean with gentle waves."
-default_negative_prompt = "low quality, worst quality, blurry, distorted, deformed, ugly, bad anatomy"
+default_prompt_t2v = "Two anthropomorphic cats in comfy boxing gear and bright gloves fight intensely on a spotlighted stage."
+default_negative_prompt = "色调艳丽, 过曝, 静态, 细节模糊不清, 字幕, 风格, 作品, 画作, 画面, 静止, 整体发灰, 最差质量, 低质量, JPEG压缩残留, 丑陋的, 残缺的, 多余的手指, 画得不好的手部, 画得不好的脸部, 畸形的, 毁容的, 形态畸形的肢体, 手指融合, 静止不动的画面, 杂乱的背景, 三条腿, 背景人很多, 倒着走"
 
 # =========================================================
-# PROMPT ENHANCEMENT SYSTEM PROMPT
+# PROMPT ENHANCEMENT
 # =========================================================
 ENHANCE_SYSTEM_PROMPT = """You are a professional video prompt engineer. Your task is to enhance user prompts for AI video generation.
 
@@ -69,29 +103,21 @@ Enhanced: "A fluffy orange tabby cat playfully batting at a dangling yarn ball,
 """
 
 
-# =========================================================
-# PROMPT ENHANCEMENT FUNCTION
-# =========================================================
 def enhance_prompt(prompt: str) -> str:
     """Enhance the user prompt using Groq LLM API."""
     if not GROQ_API_KEY:
         return prompt
     
     try:
+        from groq import Groq
         client = Groq(api_key=GROQ_API_KEY)
         
         enhanced_text = ""
         completion = client.chat.completions.create(
             model="meta-llama/llama-4-scout-17b-16e-instruct",
             messages=[
-                {
-                    "role": "system",
-                    "content": ENHANCE_SYSTEM_PROMPT
-                },
-                {
-                    "role": "user",
-                    "content": f"Enhance this video generation prompt: {prompt}"
-                }
+                {"role": "system", "content": ENHANCE_SYSTEM_PROMPT},
+                {"role": "user", "content": f"Enhance this video generation prompt: {prompt}"}
             ],
             temperature=0.7,
             max_completion_tokens=512,
@@ -112,95 +138,80 @@ def enhance_prompt(prompt: str) -> str:
 
 
 # =========================================================
-# HELPER FUNCTIONS
-# =========================================================
-def get_num_frames(duration_seconds: float) -> int:
-    """
-    Calculate number of frames based on duration.
-    num_frames - 1 must be divisible by 4, so num_frames must be 4k + 1
-    Valid values: 5, 9, 13, 17, 21, 25, 29, 33, 37, 41, 45, 49, 53, 57, 61, 65, 69, 73, 77, 81
-    """
-    raw_frames = int(duration_seconds * FIXED_FPS)
-    # Round to nearest valid frame count (4k + 1)
-    k = round((raw_frames - 1) / 4)
-    num_frames = 4 * k + 1
-    # Clamp to valid range
-    num_frames = max(MIN_FRAMES, min(MAX_FRAMES, num_frames))
-    return num_frames
-
-
+# GENERATION FUNCTIONS
 # =========================================================
-# MAIN GENERATION FUNCTION
-# =========================================================
-@spaces.GPU
+def get_duration(
+    prompt,
+    negative_prompt,
+    enhance_prompt_option,
+    duration_seconds,
+    guidance_scale,
+    guidance_scale_2,
+    steps,
+    seed,
+    randomize_seed,
+    progress,
+):
+    return steps * 15
+
+
+@spaces.GPU(duration=get_duration)
 def generate_video(
-    prompt: str,
-    negative_prompt: str,
-    enhance_prompt_option: bool,
-    duration_seconds: float,
-    guidance_scale: float,
-    num_inference_steps: int,
-    height: int,
-    width: int,
-    seed: int,
-    randomize_seed: bool,
-    progress: gr.Progress = gr.Progress(track_tqdm=True),
-) -> Tuple[str, int, str, str]:
-    """Generate video from text prompt."""
-    
-    if not prompt.strip():
-        raise gr.Error("Please enter a prompt.")
-    
+    prompt,
+    negative_prompt=default_negative_prompt,
+    enhance_prompt_option=False,
+    duration_seconds=MAX_DURATION,
+    guidance_scale=1,
+    guidance_scale_2=3,
+    steps=4,
+    seed=42,
+    randomize_seed=False,
+    progress=gr.Progress(track_tqdm=True),
+):
     # Enhance prompt if option is enabled
     final_prompt = prompt
     if enhance_prompt_option:
         final_prompt = enhance_prompt(prompt)
         print(f"Enhanced Prompt: {final_prompt}")
     
-    # Handle seed
+    num_frames = np.clip(int(round(duration_seconds * FIXED_FPS)), MIN_FRAMES_MODEL, MAX_FRAMES_MODEL)
     current_seed = random.randint(0, MAX_SEED) if randomize_seed else int(seed)
-    generator = torch.Generator(device=device).manual_seed(current_seed)
-    
-    # Calculate frames (must be 4k + 1)
-    num_frames = get_num_frames(duration_seconds)
-    actual_duration = num_frames / FIXED_FPS
-    
-    print(f"Generating video with {num_frames} frames ({actual_duration:.2f}s)")
-    
-    # Generate video
-    output_frames = pipe(
+
+    output_frames_list = pipe(
         prompt=final_prompt,
         negative_prompt=negative_prompt,
-        height=height,
-        width=width,
+        height=480,
+        width=832,
         num_frames=num_frames,
-        guidance_scale=guidance_scale,
-        num_inference_steps=num_inference_steps,
-        generator=generator,
+        guidance_scale=float(guidance_scale),
+        guidance_scale_2=float(guidance_scale_2),
+        num_inference_steps=int(steps),
+        generator=torch.Generator(device="cuda").manual_seed(current_seed),
     ).frames[0]
-    
-    # Export to video file
+
     with tempfile.NamedTemporaryFile(suffix=".mp4", delete=False) as tmpfile:
         video_path = tmpfile.name
-    export_to_video(output_frames, video_path, fps=FIXED_FPS)
-    
+
+    export_to_video(output_frames_list, video_path, fps=FIXED_FPS)
+
     # Build info log
+    actual_duration = num_frames / FIXED_FPS
     info_log = f"""✅ VIDEO GENERATION COMPLETE!
 {'=' * 50}
 🎬 Video Info:
    • Duration: {actual_duration:.2f} seconds
    • Total Frames: {num_frames}
    • FPS: {FIXED_FPS}
-   • Resolution: {width} x {height}
+   • Resolution: 832 x 480
 {'=' * 50}
 ⚙️ Generation Settings:
-   • Guidance Scale: {guidance_scale}
-   • Inference Steps: {num_inference_steps}
+   • Guidance Scale: {guidance_scale} / {guidance_scale_2}
+   • Inference Steps: {steps}
    • Seed: {current_seed}
    • Prompt Enhanced: {'Yes' if enhance_prompt_option else 'No'}
 {'=' * 50}
 💾 Ready to download!"""
-    
+
     return video_path, current_seed, final_prompt, info_log
 
 
@@ -538,6 +549,20 @@ a:hover {
     gap: 1rem !important;
 }
 
+/* ===== 🎨 Examples 섹션 ===== */
+#examples .gr-sample {
+    border: 3px solid #1F2937 !important;
+    border-radius: 8px !important;
+    box-shadow: 4px 4px 0px #1F2937 !important;
+    background: #FFFFFF !important;
+    transition: all 0.2s ease !important;
+}
+
+#examples .gr-sample:hover {
+    transform: translate(-2px, -2px) !important;
+    box-shadow: 6px 6px 0px #1F2937 !important;
+}
+
 /* ===== 반응형 조정 ===== */
 @media (max-width: 768px) {
     .header-text h1 {
@@ -595,7 +620,7 @@ with gr.Blocks() as demo:
         with gr.Column(scale=1, min_width=320):
             prompt_input = gr.Textbox(
                 label="✏️ Your Prompt",
-                value=default_prompt,
+                value=default_prompt_t2v,
                 placeholder="Describe the video you want to create...",
                 lines=4
             )
@@ -606,15 +631,16 @@ with gr.Blocks() as demo:
                 info="Use AI to automatically enhance your prompt for better results"
             )
             
-            duration_slider = gr.Slider(
+            duration_seconds_input = gr.Slider(
+                minimum=MIN_DURATION,
+                maximum=MAX_DURATION,
+                step=0.1,
+                value=MAX_DURATION,
                 label="⏱️ Duration (seconds)",
-                minimum=0.5,
-                maximum=3.5,
-                step=0.5,
-                value=2.0
+                info=f"Range: {MIN_DURATION}s - {MAX_DURATION}s at {FIXED_FPS}fps"
             )
             
-            generate_btn = gr.Button(
+            generate_button = gr.Button(
                 "🎬 GENERATE VIDEO! 🚀",
                 variant="primary",
                 size="lg",
@@ -625,46 +651,41 @@ with gr.Blocks() as demo:
                 negative_prompt_input = gr.Textbox(
                     label="Negative Prompt",
                     value=default_negative_prompt,
-                    lines=2
-                )
-                guidance_scale_slider = gr.Slider(
-                    label="Guidance Scale",
-                    minimum=1.0,
-                    maximum=15.0,
-                    step=0.5,
-                    value=7.5
-                )
-                num_inference_steps_slider = gr.Slider(
-                    label="Inference Steps",
-                    minimum=10,
-                    maximum=50,
-                    step=1,
-                    value=20
-                )
-                height_slider = gr.Slider(
-                    label="Height",
-                    minimum=256,
-                    maximum=720,
-                    step=16,
-                    value=480
-                )
-                width_slider = gr.Slider(
-                    label="Width",
-                    minimum=256,
-                    maximum=1280,
-                    step=16,
-                    value=832
+                    lines=3
                 )
-                seed_slider = gr.Slider(
+                seed_input = gr.Slider(
                     label="Seed",
                     minimum=0,
                     maximum=MAX_SEED,
                     step=1,
-                    value=0
+                    value=42,
+                    interactive=True
                 )
                 randomize_seed_checkbox = gr.Checkbox(
                     label="Randomize Seed",
-                    value=True
+                    value=True,
+                    interactive=True
+                )
+                steps_slider = gr.Slider(
+                    minimum=1,
+                    maximum=30,
+                    step=1,
+                    value=4,
+                    label="Inference Steps"
+                )
+                guidance_scale_input = gr.Slider(
+                    minimum=0.0,
+                    maximum=10.0,
+                    step=0.5,
+                    value=1,
+                    label="Guidance Scale (High Noise)"
+                )
+                guidance_scale_2_input = gr.Slider(
+                    minimum=0.0,
+                    maximum=10.0,
+                    step=0.5,
+                    value=3,
+                    label="Guidance Scale 2 (Low Noise)"
                 )
             
             with gr.Accordion("📜 Generation Info", open=True):
@@ -682,6 +703,7 @@ with gr.Blocks() as demo:
             video_output = gr.Video(
                 label="🎥 Generated Video",
                 autoplay=True,
+                interactive=False,
                 height=400,
                 elem_classes="video-output"
             )
@@ -699,33 +721,46 @@ with gr.Blocks() as demo:
                 </p>
                 """
             )
-
+    
+    # Examples section
+    gr.Examples(
+        examples=[
+            ["POV selfie video, white cat with sunglasses standing on surfboard, relaxed smile, tropical beach behind. Surfboard tips, cat falls into ocean, camera plunges underwater with bubbles."],
+            ["Two anthropomorphic cats in comfy boxing gear and bright gloves fight intensely on a spotlighted stage."],
+            ["A cinematic shot of a boat sailing on a calm sea at sunset."],
+            ["Drone footage flying over a futuristic city with flying cars."],
+        ],
+        inputs=[prompt_input],
+        outputs=[video_output, seed_input, final_prompt_output, info_log],
+        fn=generate_video,
+        cache_examples="lazy",
+        elem_id="examples"
+    )
+    
     # Define inputs and outputs
-    inputs = [
+    ui_inputs = [
         prompt_input,
         negative_prompt_input,
         enhance_prompt_checkbox,
-        duration_slider,
-        guidance_scale_slider,
-        num_inference_steps_slider,
-        height_slider,
-        width_slider,
-        seed_slider,
-        randomize_seed_checkbox,
+        duration_seconds_input,
+        guidance_scale_input,
+        guidance_scale_2_input,
+        steps_slider,
+        seed_input,
+        randomize_seed_checkbox
     ]
     
-    outputs = [
+    ui_outputs = [
         video_output,
-        seed_slider,
+        seed_input,
         final_prompt_output,
-        info_log,
+        info_log
     ]
-
-    # Generate button click
-    generate_btn.click(
+    
+    generate_button.click(
         fn=generate_video,
-        inputs=inputs,
-        outputs=outputs,
+        inputs=ui_inputs,
+        outputs=ui_outputs
     )
 
 if __name__ == "__main__":