Update app.py

app.py

@@ -1,20 +1,13 @@
-# 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')
-
-# Actual demo code
 import spaces
 import torch
-from diffusers import WanPipeline, AutoencoderKLWan
-from diffusers.models.transformers.transformer_wan import WanTransformer3DModel
+from diffusers import WanPipeline
 from diffusers.utils.export_utils import export_to_video
 import gradio as gr
 import tempfile
 import numpy as np
-from PIL import Image
 import random
 import gc
-from optimization import optimize_pipeline_
 
 
 # =========================================================
@@ -23,16 +16,16 @@ from optimization import optimize_pipeline_
 GROQ_API_KEY = os.getenv("GROQ_API_KEY")
 
 # =========================================================
-# MODEL CONFIGURATION
+# MODEL CONFIGURATION - Using lighter 1.3B model
 # =========================================================
-MODEL_ID = "Wan-AI/Wan2.2-T2V-A14B-Diffusers"
+MODEL_ID = "Wan-AI/Wan2.1-T2V-1.3B-Diffusers"
 
 LANDSCAPE_WIDTH = 832
 LANDSCAPE_HEIGHT = 480
 MAX_SEED = np.iinfo(np.int32).max
 
 FIXED_FPS = 16
-MIN_FRAMES_MODEL = 8
+MIN_FRAMES_MODEL = 9  # Must be 4k+1: 5, 9, 13, 17...
 MAX_FRAMES_MODEL = 81
 
 MIN_DURATION = round(MIN_FRAMES_MODEL/FIXED_FPS, 1)
@@ -41,48 +34,25 @@ MAX_DURATION = round(MAX_FRAMES_MODEL/FIXED_FPS, 1)
 # =========================================================
 # MODEL LOADING
 # =========================================================
-vae = AutoencoderKLWan.from_pretrained(
-    "Wan-AI/Wan2.2-T2V-A14B-Diffusers", 
-    subfolder="vae", 
-    torch_dtype=torch.float32
-)
-
+print("Loading Wan 1.3B pipeline...")
 pipe = WanPipeline.from_pretrained(
     MODEL_ID,
-    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')
 
+# Memory cleanup
 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,
-)
+print("Pipeline loaded successfully!")
 
 # =========================================================
 # DEFAULT PROMPTS
 # =========================================================
 default_prompt_t2v = "Two anthropomorphic cats in comfy boxing gear and bright gloves fight intensely on a spotlighted stage."
-default_negative_prompt = "色调艳丽, 过曝, 静态, 细节模糊不清, 字幕, 风格, 作品, 画作, 画面, 静止, 整体发灰, 最差质量, 低质量, JPEG压缩残留, 丑陋的, 残缺的, 多余的手指, 画得不好的手部, 画得不好的脸部, 畸形的, 毁容的, 形态畸形的肢体, 手指融合, 静止不动的画面, 杂乱的背景, 三条腿, 背景人很多, 倒着走"
+default_negative_prompt = "low quality, worst quality, blurry, distorted, deformed, ugly, bad anatomy, static, watermark, text, subtitle, oversaturated, underexposed"
 
 # =========================================================
 # PROMPT ENHANCEMENT
@@ -136,41 +106,34 @@ def enhance_prompt(prompt: str) -> str:
 # =========================================================
 # GENERATION FUNCTIONS
 # =========================================================
-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)
+@spaces.GPU(duration=300)
 def generate_video(
     prompt,
     negative_prompt=default_negative_prompt,
     enhance_prompt_option=False,
     duration_seconds=MAX_DURATION,
-    guidance_scale=1,
-    guidance_scale_2=3,
-    steps=4,
+    guidance_scale=5.0,
+    steps=30,
     seed=42,
     randomize_seed=False,
     progress=gr.Progress(track_tqdm=True),
 ):
+    # Clear memory before generation
+    gc.collect()
+    torch.cuda.empty_cache()
+    
     # Enhance prompt if option is enabled
     final_prompt = prompt
     if enhance_prompt_option:
         final_prompt = enhance_prompt(prompt)
         print(f"Enhanced Prompt: {final_prompt}")
     
-    num_frames = np.clip(int(round(duration_seconds * FIXED_FPS)), MIN_FRAMES_MODEL, MAX_FRAMES_MODEL)
+    # Calculate num_frames - must satisfy (num_frames - 1) % 4 == 0
+    raw_frames = int(round(duration_seconds * FIXED_FPS))
+    k = round((raw_frames - 1) / 4)
+    num_frames = 4 * k + 1
+    num_frames = np.clip(num_frames, MIN_FRAMES_MODEL, MAX_FRAMES_MODEL)
+    
     current_seed = random.randint(0, MAX_SEED) if randomize_seed else int(seed)
 
     output_frames_list = pipe(
@@ -180,7 +143,6 @@ def generate_video(
         width=LANDSCAPE_WIDTH,
         num_frames=num_frames,
         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]
@@ -189,6 +151,10 @@ def generate_video(
         video_path = tmpfile.name
 
     export_to_video(output_frames_list, video_path, fps=FIXED_FPS)
+    
+    # Clear memory after generation
+    gc.collect()
+    torch.cuda.empty_cache()
 
     # Build info log
     actual_duration = num_frames / FIXED_FPS
@@ -201,7 +167,7 @@ def generate_video(
    • Resolution: {LANDSCAPE_WIDTH} x {LANDSCAPE_HEIGHT}
 {'=' * 50}
 ⚙️ Generation Settings:
-   • Guidance Scale: {guidance_scale} / {guidance_scale_2}
+   • Guidance Scale: {guidance_scale}
    • Inference Steps: {steps}
    • Seed: {current_seed}
    • Prompt Enhanced: {'Yes' if enhance_prompt_option else 'No'}
@@ -446,16 +412,6 @@ button.secondary:active,
     font-size: 1.1rem !important;
 }
 
-/* ===== 🎨 이미지 출력 영역 ===== */
-.gr-image,
-.image-container {
-    border: 4px solid #1F2937 !important;
-    border-radius: 8px !important;
-    box-shadow: 8px 8px 0px #1F2937 !important;
-    overflow: hidden !important;
-    background: #FFFFFF !important;
-}
-
 /* ===== 🎨 라벨 스타일 ===== */
 label,
 .gr-input-label,
@@ -482,14 +438,6 @@ input[type="range"] {
     accent-color: #3B82F6 !important;
 }
 
-/* ===== 🎨 정보 텍스트 ===== */
-.gr-info,
-.info {
-    color: #6B7280 !important;
-    font-family: 'Comic Neue', cursive !important;
-    font-size: 0.9rem !important;
-}
-
 /* ===== 🎨 프로그레스 바 ===== */
 .progress-bar,
 .gr-progress-bar {
@@ -536,29 +484,6 @@ a:hover {
     color: #EF4444 !important;
 }
 
-/* ===== 🎨 Row/Column 간격 ===== */
-.gr-row {
-    gap: 1.5rem !important;
-}
-
-.gr-column {
-    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 {
@@ -580,7 +505,7 @@ a:hover {
     }
 }
 
-/* ===== 🎨 다크모드 비활성화 (코믹은 밝아야 함) ===== */
+/* ===== 🎨 다크모드 비활성화 ===== */
 @media (prefers-color-scheme: dark) {
     .gradio-container {
         background-color: #FEF9C3 !important;
@@ -631,7 +556,7 @@ with gr.Blocks() as demo:
                 minimum=MIN_DURATION,
                 maximum=MAX_DURATION,
                 step=0.1,
-                value=MAX_DURATION,
+                value=2.0,
                 label="⏱️ Duration (seconds)",
                 info=f"Range: {MIN_DURATION}s - {MAX_DURATION}s at {FIXED_FPS}fps"
             )
@@ -663,25 +588,18 @@ with gr.Blocks() as demo:
                     interactive=True
                 )
                 steps_slider = gr.Slider(
-                    minimum=1,
-                    maximum=30,
+                    minimum=10,
+                    maximum=50,
                     step=1,
-                    value=4,
+                    value=30,
                     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,
+                    minimum=1.0,
+                    maximum=15.0,
                     step=0.5,
-                    value=3,
-                    label="Guidance Scale 2 (Low Noise)"
+                    value=5.0,
+                    label="Guidance Scale"
                 )
             
             with gr.Accordion("📜 Generation Info", open=True):
@@ -721,7 +639,7 @@ with gr.Blocks() as demo:
     # 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."],
+            ["POV selfie video, white cat with sunglasses standing on surfboard, relaxed smile, tropical beach behind."],
             ["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."],
@@ -738,7 +656,6 @@ with gr.Blocks() as demo:
         enhance_prompt_checkbox,
         duration_seconds_input,
         guidance_scale_input,
-        guidance_scale_2_input,
         steps_slider,
         seed_input,
         randomize_seed_checkbox