Update app.py

app.py

@@ -51,6 +51,13 @@ def initialize_video_pipeline():
             # Install PyTorch 2.8 (if needed)
             os.system('pip install --upgrade --pre --extra-index-url https://download.pytorch.org/whl/nightly/cu126 "torch<2.9" spaces')
             
+            # Import optimization module
+            try:
+                from optimization import optimize_pipeline_
+            except ImportError:
+                print("Warning: optimization module not found, skipping optimization")
+                optimize_pipeline_ = None
+            
             video_pipe = WanImageToVideoPipeline.from_pretrained(VIDEO_MODEL_ID,
                 transformer=WanTransformer3DModel.from_pretrained('cbensimon/Wan2.2-I2V-A14B-bf16-Diffusers',
                     subfolder='transformer',
@@ -70,6 +77,16 @@ def initialize_video_pipeline():
                 gc.collect()
                 torch.cuda.synchronize()
                 torch.cuda.empty_cache()
+            
+            # Optimize pipeline if module available
+            if optimize_pipeline_ is not None:
+                optimize_pipeline_(video_pipe,
+                    image=Image.new('RGB', (LANDSCAPE_WIDTH, LANDSCAPE_HEIGHT)),
+                    prompt='prompt',
+                    height=LANDSCAPE_HEIGHT,
+                    width=LANDSCAPE_WIDTH,
+                    num_frames=MAX_FRAMES_MODEL,
+                )
                 
             print("Video pipeline initialized successfully!")
         except Exception as e:
@@ -225,7 +242,10 @@ def resize_image_landscape(image: Image.Image) -> Image.Image:
     
     return image.resize((LANDSCAPE_WIDTH, LANDSCAPE_HEIGHT), Image.LANCZOS)
 
-@spaces.GPU(duration=120)
+def get_duration(steps):
+    return int(steps) * 15
+
+@spaces.GPU(duration=get_duration)
 def generate_video(
     input_image,
     prompt,
@@ -248,29 +268,55 @@ def generate_video(
     if video_pipe is None:
         raise gr.Error("Video pipeline not initialized. Please check GPU availability.")
     
-    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)
-    resized_image = resize_image_for_video(input_image)
-    
-    output_frames_list = video_pipe(
-        image=resized_image,
-        prompt=prompt,
-        negative_prompt=negative_prompt,
-        height=resized_image.height,
-        width=resized_image.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]
-    
-    with tempfile.NamedTemporaryFile(suffix=".mp4", delete=False) as tmpfile:
-        video_path = tmpfile.name
-    
-    export_to_video(output_frames_list, video_path, fps=FIXED_FPS)
-    
-    return video_path, current_seed, "🎬 Video generated successfully!"
+    try:
+        # Ensure frames are divisible by 4
+        num_frames = int(round(duration_seconds * FIXED_FPS))
+        num_frames = np.clip(num_frames, MIN_FRAMES_MODEL, MAX_FRAMES_MODEL)
+        # Round to nearest number divisible by 4
+        num_frames = ((num_frames - 1) // 4) * 4 + 1
+        
+        current_seed = random.randint(0, MAX_SEED) if randomize_seed else int(seed)
+        resized_image = resize_image_for_video(input_image)
+        
+        # Clear cache before generation
+        torch.cuda.empty_cache()
+        gc.collect()
+        
+        # Generate video with memory management
+        with torch.inference_mode():
+            output_frames_list = video_pipe(
+                image=resized_image,
+                prompt=prompt,
+                negative_prompt=negative_prompt,
+                height=resized_image.height,
+                width=resized_image.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]
+        
+        # Clear cache after generation
+        torch.cuda.empty_cache()
+        gc.collect()
+        
+        with tempfile.NamedTemporaryFile(suffix=".mp4", delete=False) as tmpfile:
+            video_path = tmpfile.name
+        
+        export_to_video(output_frames_list, video_path, fps=FIXED_FPS)
+        
+        return video_path, current_seed, "🎬 Video generated successfully!"
+        
+    except RuntimeError as e:
+        if "out of memory" in str(e).lower() or "CUDA" in str(e):
+            torch.cuda.empty_cache()
+            gc.collect()
+            raise gr.Error("GPU memory error. Try reducing the duration or number of steps.")
+        else:
+            raise gr.Error(f"Video generation error: {str(e)}")
+    except Exception as e:
+        raise gr.Error(f"Unexpected error: {str(e)}")
 
 # ===========================
 # Enhanced CSS
@@ -383,7 +429,7 @@ with gr.Blocks(css=css, theme=gr.themes.Base()) as demo:
     
     with gr.Column(elem_classes="header-container"):
         gr.HTML("""
-            <h1 class="logo-text">🍌 Nano Banana + Video</h1>
+            <h1 class="logo-text">🍌 Open Nano Banana + Video</h1>
             <p class="subtitle">AI-Powered Image Style Transfer with Video Generation</p>
             <div style="display: flex; justify-content: center; align-items: center; gap: 10px; margin-top: 20px;">
                 <a href="https://huggingface.co/spaces/openfree/Nano-Banana-Upscale" target="_blank">
@@ -570,12 +616,15 @@ with gr.Blocks(css=css, theme=gr.themes.Base()) as demo:
         guidance_1, guidance_2, video_seed, randomize_seed
     ]
     
-    def generate_video_wrapper(*args):
+    def generate_video_wrapper(img, prompt, steps, neg_prompt, duration, g1, g2, seed, rand_seed):
         try:
-            video_path, seed, status = generate_video(*args)
-            return video_path, seed, status
+            # Pass steps as first argument for GPU duration
+            video_path, new_seed, status = generate_video(
+                img, prompt, steps, neg_prompt, duration, g1, g2, seed, rand_seed
+            )
+            return video_path, new_seed, status
         except Exception as e:
-            return None, args[7], f"Error: {str(e)}"
+            return None, seed, f"Error: {str(e)}"
     
     generate_video_btn.click(
         fn=generate_video_wrapper,
@@ -583,7 +632,18 @@ with gr.Blocks(css=css, theme=gr.themes.Base()) as demo:
         outputs=[video_output, video_seed, video_status]
     )
     
-
+    # Examples for image generation
+    gr.Examples(
+        examples=[
+            ["Create a dreamy watercolor style with soft pastels", "examples/photo1.jpg", None],
+            ["Transform into cyberpunk neon aesthetic", "examples/photo2.jpg", "examples/style.jpg"],
+            ["Make it look like Studio Ghibli animation", "examples/landscape.jpg", None],
+        ],
+        inputs=[style_prompt, image1, image2],
+        outputs=[output_image, img_status],
+        fn=process_images,
+        cache_examples=False
+    )
 
 # Launch
 if __name__ == "__main__":