Update app.py

app.py

@@ -14,6 +14,7 @@ import numpy as np
 from PIL import Image
 import random
 import gc
+from optimization import optimize_pipeline_
 
 
 # =========================================================
@@ -26,14 +27,13 @@ GROQ_API_KEY = os.getenv("GROQ_API_KEY")
 # =========================================================
 MODEL_ID = "Wan-AI/Wan2.2-T2V-A14B-Diffusers"
 
-# Reduced resolution for memory efficiency
-LANDSCAPE_WIDTH = 640
-LANDSCAPE_HEIGHT = 384
+LANDSCAPE_WIDTH = 832
+LANDSCAPE_HEIGHT = 480
 MAX_SEED = np.iinfo(np.int32).max
 
 FIXED_FPS = 16
-MIN_FRAMES_MODEL = 9  # Must be 4k+1: 5, 9, 13, 17...
-MAX_FRAMES_MODEL = 49  # Reduced from 81 for memory
+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)
@@ -41,14 +41,12 @@ MAX_DURATION = round(MAX_FRAMES_MODEL/FIXED_FPS, 1)
 # =========================================================
 # MODEL LOADING
 # =========================================================
-print("Loading VAE...")
 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,
     transformer=WanTransformer3DModel.from_pretrained(
@@ -67,13 +65,18 @@ pipe = WanPipeline.from_pretrained(
     torch_dtype=torch.bfloat16,
 ).to('cuda')
 
-# Memory cleanup
 for i in range(3): 
     gc.collect()
     torch.cuda.synchronize() 
     torch.cuda.empty_cache()
 
-print("Pipeline loaded successfully!")
+optimize_pipeline_(
+    pipe,
+    prompt='prompt',
+    height=LANDSCAPE_HEIGHT,
+    width=LANDSCAPE_WIDTH,
+    num_frames=MAX_FRAMES_MODEL,
+)
 
 # =========================================================
 # DEFAULT PROMPTS
@@ -161,23 +164,13 @@ def generate_video(
     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}")
     
-    # Calculate num_frames - must satisfy (num_frames - 1) % 4 == 0
-    # Valid values: 5, 9, 13, 17, 21, 25, 29, 33, 37, 41, 45, 49
-    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)
-    
+    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)
 
     output_frames_list = pipe(
@@ -196,10 +189,6 @@ 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