Update app.py

app.py

@@ -4,7 +4,6 @@ import os
 from PIL import Image
 import requests
 from io import BytesIO
-import time
 import tempfile
 import base64
 import spaces
@@ -12,6 +11,7 @@ import torch
 import numpy as np
 import random
 import gc
+import time
 
 # ===========================
 # Configuration
@@ -21,26 +21,35 @@ import gc
 os.environ['REPLICATE_API_TOKEN'] = os.getenv('REPLICATE_API_TOKEN')
 
 # Video Model Configuration
-VIDEO_MODEL_ID = "cjwbw/videocrafter2:02e509c789964be7d70de8d8fef3a6dd18f160b37272bcccc742d5adabb9f38f"  # Using public model
-LANDSCAPE_WIDTH = 512  # Reduced for stability
-LANDSCAPE_HEIGHT = 320  # Reduced for stability
+LANDSCAPE_WIDTH = 512
+LANDSCAPE_HEIGHT = 320
 MAX_SEED = np.iinfo(np.int32).max
-FIXED_FPS = 8  # Reduced FPS
-MIN_FRAMES_MODEL = 8
-MAX_FRAMES_MODEL = 32  # Reduced max frames
+FIXED_FPS = 8
 
 default_prompt_i2v = "make this image come alive, smooth animation"
 default_negative_prompt = "static, still, blurry, low quality"
 
 # ===========================
-# Image Processing Functions
+# Helper Functions
 # ===========================
 
+def check_api_token():
+    """Check if Replicate API token is set"""
+    token = os.getenv('REPLICATE_API_TOKEN')
+    return token is not None and token.strip() != ""
+
 def upload_image_to_hosting(image):
     """Upload image to hosting service"""
     try:
+        # Convert to RGB if needed
+        if image.mode == 'RGBA':
+            background = Image.new('RGB', image.size, (255, 255, 255))
+            background.paste(image, mask=image.split()[3])
+            image = background
+        
+        # Try imgbb.com first
         buffered = BytesIO()
-        image.save(buffered, format="PNG")
+        image.save(buffered, format="PNG", optimize=True, quality=95)
         buffered.seek(0)
         img_base64 = base64.b64encode(buffered.getvalue()).decode()
         
@@ -58,7 +67,7 @@ def upload_image_to_hosting(image):
             if data.get('success'):
                 return data['data']['url']
     except Exception as e:
-        print(f"Upload failed: {e}")
+        print(f"Upload error: {e}")
     
     # Fallback to base64
     buffered = BytesIO()
@@ -67,117 +76,152 @@ def upload_image_to_hosting(image):
     img_base64 = base64.b64encode(buffered.getvalue()).decode()
     return f"data:image/png;base64,{img_base64}"
 
+# ===========================
+# Image Generation Functions
+# ===========================
+
 def process_images(prompt, image1, image2=None):
     """Process images using Replicate API"""
-    if not image1:
-        return None, "Please upload at least one image", None
+    if not prompt or prompt.strip() == "":
+        return None, "Please enter a prompt", None
     
-    if not os.getenv('REPLICATE_API_TOKEN'):
-        return None, "Please set REPLICATE_API_TOKEN in Space settings", None
+    if not check_api_token():
+        return None, "⚠️ Please set REPLICATE_API_TOKEN in Space settings (Settings > Repository secrets)", None
     
     try:
-        # Upload image
-        url1 = upload_image_to_hosting(image1)
+        # Simple text-to-image generation using SDXL
+        print(f"Generating image with prompt: {prompt}")
         
-        # Use SDXL for image generation/editing
         output = replicate.run(
             "stability-ai/sdxl:39ed52f2a78e934b3ba6e2a89f5b1c712de7dfea535525255b1aa35c5565e08b",
             input={
-                "prompt": prompt + ", high quality, detailed",
-                "negative_prompt": "low quality, blurry, distorted",
+                "prompt": prompt + ", high quality, detailed, 8k",
+                "negative_prompt": "blurry, low quality, distorted, deformed",
                 "width": 1024,
                 "height": 1024,
-                "num_inference_steps": 25
+                "num_inference_steps": 30,
+                "guidance_scale": 7.5
             }
         )
         
-        if output and isinstance(output, list) and len(output) > 0:
-            img_url = output[0]
+        # Handle output
+        if output:
+            if isinstance(output, list) and len(output) > 0:
+                img_url = output[0]
+            elif isinstance(output, str):
+                img_url = output
+            else:
+                img_url = str(output)
+            
+            # Download image
             response = requests.get(img_url, timeout=30)
             if response.status_code == 200:
                 img = Image.open(BytesIO(response.content))
-                return img, "✨ Image generated successfully!", img
+                return img, "✨ Image generated successfully! You can now create a video from it.", img
         
-        return None, "Could not process output", None
+        return None, "Failed to generate image", None
         
     except Exception as e:
         error_msg = str(e)
-        if "trial" in error_msg.lower():
-            return None, "Replicate API limit reached. Please try again later.", None
-        return None, f"Error: {error_msg[:200]}", None
+        if "authentication" in error_msg.lower():
+            return None, "❌ Invalid API token. Please check your REPLICATE_API_TOKEN.", None
+        elif "rate limit" in error_msg.lower():
+            return None, "⏳ Rate limit reached. Please try again later.", None
+        else:
+            return None, f"Error: {error_msg[:200]}", None
 
 # ===========================
 # Video Generation Functions
 # ===========================
 
 def resize_image_for_video(image: Image.Image) -> Image.Image:
-    """Resize image for video generation"""
-    # Convert RGBA to RGB if necessary
+    """Resize and prepare image for video generation"""
+    # Convert RGBA to RGB
     if image.mode == 'RGBA':
         background = Image.new('RGB', image.size, (255, 255, 255))
         background.paste(image, mask=image.split()[3])
         image = background
+    elif image.mode != 'RGB':
+        image = image.convert('RGB')
     
     # Resize to target dimensions
     image = image.resize((LANDSCAPE_WIDTH, LANDSCAPE_HEIGHT), Image.LANCZOS)
     return image
 
-# GPU function with proper decorator
 @spaces.GPU(duration=60)
 def generate_video_gpu(
     input_image,
     prompt,
-    steps=25,
-    negative_prompt=default_negative_prompt,
+    steps,
+    negative_prompt,
+    duration_seconds,
+    seed,
+    randomize_seed
+):
+    """GPU-accelerated video generation"""
+    try:
+        # This function runs on GPU
+        # For demo purposes, we'll just process the image
+        # In production, you would run actual video generation here
+        
+        # Clear GPU memory
+        if torch.cuda.is_available():
+            torch.cuda.empty_cache()
+            gc.collect()
+            
+        # Simulate video generation
+        time.sleep(2)  # Simulate processing time
+        
+        # For now, return a placeholder since actual video generation requires specific models
+        return None, seed, "🎬 Video generation simulated (GPU function executed successfully)"
+        
+    except Exception as e:
+        return None, seed, f"GPU Error: {str(e)[:200]}"
+
+def generate_video_replicate(
+    input_image,
+    prompt,
+    steps=30,
+    negative_prompt="",
     duration_seconds=2.0,
     seed=42,
-    randomize_seed=False,
+    randomize_seed=False
 ):
-    """Generate video using Replicate API with GPU"""
+    """Generate video using Replicate API (no GPU needed)"""
+    
+    if not check_api_token():
+        return None, seed, "⚠️ Please set REPLICATE_API_TOKEN"
     
     if input_image is None:
         return None, seed, "Please provide an input image"
     
     try:
-        # Clear GPU memory
-        if torch.cuda.is_available():
-            torch.cuda.empty_cache()
-            gc.collect()
-        
-        # Resize image
+        # Resize and prepare image
         resized_image = resize_image_for_video(input_image)
         
-        # Save resized image temporarily
-        with tempfile.NamedTemporaryFile(suffix=".png", delete=False) as tmp_img:
-            resized_image.save(tmp_img.name)
-            
-            # Upload to hosting
-            img_url = upload_image_to_hosting(resized_image)
+        # Upload image
+        img_url = upload_image_to_hosting(resized_image)
         
         current_seed = random.randint(0, MAX_SEED) if randomize_seed else int(seed)
         
-        # Use Replicate for video generation
         print("Generating video with Replicate...")
+        
+        # Use AnimateDiff or similar model
         output = replicate.run(
-            VIDEO_MODEL_ID,
+            "stability-ai/stable-video-diffusion:3f0457e4619daac51203dedb472816fd4af51f3149fa7a9e0b5ffcf1b8172438",
             input={
-                "prompt": prompt,
-                "image": img_url,
-                "steps": int(steps),
-                "fps": FIXED_FPS,
-                "seconds": min(duration_seconds, 3),  # Limit to 3 seconds
+                "input_image": img_url,
+                "frames_per_second": FIXED_FPS,
+                "motion_bucket_id": 127,  # Controls motion amount
+                "cond_aug": 0.02,
+                "decoding_t": 7,  # Number of frames
                 "seed": current_seed
             }
         )
         
         if output:
             # Download video
-            if isinstance(output, str):
-                video_url = output
-            elif hasattr(output, 'url'):
-                video_url = output.url()
-            else:
-                video_url = str(output)
+            video_url = output if isinstance(output, str) else str(output)
             
             response = requests.get(video_url, timeout=60)
             if response.status_code == 200:
@@ -189,44 +233,10 @@ def generate_video_gpu(
         
     except Exception as e:
         error_msg = str(e)
-        if "out of memory" in error_msg.lower():
-            torch.cuda.empty_cache()
-            gc.collect()
-            return None, seed, "GPU memory exceeded. Try reducing duration."
-        return None, seed, f"Error: {error_msg[:200]}"
-
-# Wrapper function for video generation
-def generate_video(
-    input_image,
-    prompt,
-    steps=25,
-    negative_prompt=default_negative_prompt,
-    duration_seconds=2.0,
-    seed=42,
-    randomize_seed=False,
-):
-    """Wrapper function that calls the GPU function"""
-    if not os.getenv('REPLICATE_API_TOKEN'):
-        return None, seed, "Please set REPLICATE_API_TOKEN in Space settings"
-    
-    return generate_video_gpu(
-        input_image,
-        prompt,
-        steps,
-        negative_prompt,
-        duration_seconds,
-        seed,
-        randomize_seed
-    )
-
-# ===========================
-# Simple dummy GPU function for startup
-# ===========================
-
-@spaces.GPU(duration=1)
-def dummy_gpu_function():
-    """Dummy function to satisfy Spaces GPU requirement"""
-    return "GPU initialized"
+        if "authentication" in error_msg.lower():
+            return None, seed, "❌ Invalid API token"
+        else:
+            return None, seed, f"Error: {error_msg[:200]}"
 
 # ===========================
 # CSS Styling
@@ -236,36 +246,40 @@ css = """
 .gradio-container {
     max-width: 1200px !important;
     margin: 0 auto !important;
+    padding: 20px !important;
 }
 .header-container {
     background: linear-gradient(135deg, #ffd93d, #ffb347);
-    padding: 2rem;
-    border-radius: 15px;
+    padding: 2.5rem;
+    border-radius: 20px;
     margin-bottom: 2rem;
     text-align: center;
+    box-shadow: 0 10px 30px rgba(0,0,0,0.1);
 }
 .logo-text {
-    font-size: 2.5rem;
-    font-weight: bold;
+    font-size: 3rem;
+    font-weight: 900;
     color: #2d3436;
+    margin: 0;
+    text-shadow: 2px 2px 4px rgba(0,0,0,0.1);
 }
 .subtitle {
     color: #2d3436;
-    font-size: 1.1rem;
+    font-size: 1.2rem;
     margin-top: 0.5rem;
+    font-weight: 500;
 }
-.gr-button {
-    font-size: 1rem !important;
-    padding: 12px 24px !important;
+.status-box {
+    padding: 10px;
+    border-radius: 8px;
+    margin: 10px 0;
 }
-.gr-button-primary {
-    background: linear-gradient(135deg, #ffd93d, #ffb347) !important;
-    border: none !important;
+.gr-button {
+    transition: all 0.3s ease;
 }
-.gr-button-secondary {
-    background: linear-gradient(135deg, #667eea, #764ba2) !important;
-    color: white !important;
-    border: none !important;
+.gr-button:hover {
+    transform: translateY(-2px);
+    box-shadow: 0 5px 15px rgba(0,0,0,0.2);
 }
 """
 
@@ -273,203 +287,226 @@ css = """
 # Gradio Interface
 # ===========================
 
-with gr.Blocks(css=css, theme=gr.themes.Soft()) as demo:
-    # Initialize GPU on startup
-    startup_status = gr.State(dummy_gpu_function())
-    
-    # Shared state
-    generated_image_state = gr.State(None)
-    
-    gr.HTML("""
-        <div class="header-container">
-            <h1 class="logo-text">🍌 Nano Banana + Video</h1>
-            <p class="subtitle">AI Image Generation with Video Creation</p>
-            <p style="color: #636e72; font-size: 0.9rem; margin-top: 10px;">
-                ⚠️ Note: Add REPLICATE_API_TOKEN in Space Settings > Repository secrets
-            </p>
-        </div>
-    """)
-    
-    with gr.Tabs():
-        # Tab 1: Image Generation
-        with gr.TabItem("🎨 Step 1: Generate Image"):
-            with gr.Row():
-                with gr.Column(scale=1):
-                    style_prompt = gr.Textbox(
-                        label="Image Description",
-                        placeholder="Describe what you want to create...",
-                        lines=3,
-                        value="A beautiful fantasy landscape with mountains and a river, studio ghibli style"
-                    )
+def create_interface():
+    with gr.Blocks(css=css, theme=gr.themes.Soft()) as demo:
+        # Shared state
+        generated_image_state = gr.State(None)
+        
+        # Header
+        gr.HTML("""
+            <div class="header-container">
+                <h1 class="logo-text">🍌 Nano Banana + Video</h1>
+                <p class="subtitle">Transform Text to Image, Then Bring It to Life!</p>
+                <div style="margin-top: 1rem;">
+                    <p style="color: #636e72; font-size: 0.9rem;">
+                        Step 1: Generate an image from text | Step 2: Convert it to video
+                    </p>
+                </div>
+            </div>
+        """)
+        
+        # API Token Status
+        with gr.Row():
+            gr.HTML(f"""
+                <div class="status-box" style="background: {'#d4edda' if check_api_token() else '#f8d7da'}; 
+                     color: {'#155724' if check_api_token() else '#721c24'};">
+                    <b>API Status:</b> {'✅ Token configured' if check_api_token() else '❌ Token missing - Add REPLICATE_API_TOKEN in Settings > Repository secrets'}
+                </div>
+            """)
+        
+        # Tabs
+        with gr.Tabs() as tabs:
+            # Image Generation Tab
+            with gr.TabItem("🎨 Step 1: Generate Image", id=1):
+                with gr.Row():
+                    with gr.Column(scale=3):
+                        style_prompt = gr.Textbox(
+                            label="Describe your image",
+                            placeholder="E.g., A magical forest with glowing mushrooms, fantasy art style",
+                            lines=3,
+                            value=""
+                        )
+                        
+                        # Example prompts
+                        gr.Examples(
+                            examples=[
+                                ["A majestic dragon flying over a medieval castle at sunset, fantasy art"],
+                                ["Cyberpunk city street with neon lights, rainy night, blade runner style"],
+                                ["Cute robot watering flowers in a garden, pixar animation style"],
+                                ["Northern lights dancing over snowy mountains, photorealistic"],
+                                ["Steampunk airship floating among clouds, detailed mechanical parts"],
+                            ],
+                            inputs=[style_prompt],
+                            label="💡 Example Prompts"
+                        )
+                        
+                        generate_img_btn = gr.Button(
+                            "🎨 Generate Image",
+                            variant="primary",
+                            size="lg"
+                        )
                     
-                    with gr.Row():
-                        image1 = gr.Image(
-                            label="Reference Image (Optional)",
+                    with gr.Column(scale=3):
+                        output_image = gr.Image(
+                            label="Generated Image",
                             type="pil",
-                            height=200
+                            height=400
                         )
-                        image2 = gr.Image(
-                            label="Style Reference (Optional)",
-                            type="pil",
-                            height=200
+                        
+                        img_status = gr.Textbox(
+                            label="Status",
+                            interactive=False,
+                            value="Ready to generate..."
                         )
-                    
-                    generate_img_btn = gr.Button(
-                        "🎨 Generate Image",
-                        variant="primary",
-                        size="lg"
-                    )
-                
-                with gr.Column(scale=1):
-                    output_image = gr.Image(
-                        label="Generated Result",
-                        type="pil",
-                        height=400
-                    )
-                    
-                    img_status = gr.Textbox(
-                        label="Status",
-                        interactive=False,
-                        value="Ready to generate..."
-                    )
-                    
-                    send_to_video_btn = gr.Button(
-                        "➡️ Send to Video Generation",
-                        variant="secondary",
-                        visible=False
-                    )
-        
-        # Tab 2: Video Generation  
-        with gr.TabItem("🎬 Step 2: Generate Video"):
-            gr.Markdown("### Transform your image into a video")
+                        
+                        with gr.Row():
+                            send_to_video_btn = gr.Button(
+                                "➡️ Use for Video",
+                                variant="secondary",
+                                visible=False
+                            )
             
-            with gr.Row():
-                with gr.Column(scale=1):
-                    video_input_image = gr.Image(
-                        type="pil",
-                        label="Input Image",
-                        height=300
-                    )
-                    
-                    video_prompt = gr.Textbox(
-                        label="Animation Description",
-                        value=default_prompt_i2v,
-                        lines=2
-                    )
-                    
-                    with gr.Row():
-                        duration_input = gr.Slider(
-                            minimum=1.0,
-                            maximum=3.0,
-                            step=0.5,
-                            value=2.0,
-                            label="Duration (seconds)"
+            # Video Generation Tab
+            with gr.TabItem("🎬 Step 2: Generate Video", id=2):
+                gr.Markdown("### Transform your static image into a dynamic video!")
+                
+                with gr.Row():
+                    with gr.Column(scale=3):
+                        video_input_image = gr.Image(
+                            type="pil",
+                            label="Input Image (from Step 1 or upload)",
+                            height=300
                         )
                         
-                        steps_slider = gr.Slider(
-                            minimum=10,
-                            maximum=50,
-                            step=5,
-                            value=25,
-                            label="Quality Steps"
+                        video_prompt = gr.Textbox(
+                            label="Video Motion Description (optional)",
+                            value="smooth camera pan, gentle movement",
+                            lines=2
                         )
+                        
+                        with gr.Row():
+                            video_seed = gr.Slider(
+                                label="Seed",
+                                minimum=0,
+                                maximum=MAX_SEED,
+                                step=1,
+                                value=42
+                            )
+                            randomize_seed = gr.Checkbox(
+                                label="Random seed",
+                                value=True
+                            )
+                        
+                        generate_video_btn = gr.Button(
+                            "🎬 Generate Video (Replicate API)",
+                            variant="primary",
+                            size="lg"
+                        )
+                        
+                        # GPU Test Button (optional)
+                        with gr.Accordion("Advanced Options", open=False):
+                            test_gpu_btn = gr.Button(
+                                "🖥️ Test GPU Function",
+                                variant="secondary"
+                            )
                     
-                    with gr.Row():
-                        video_seed = gr.Slider(
-                            label="Seed",
-                            minimum=0,
-                            maximum=MAX_SEED,
-                            step=1,
-                            value=42
+                    with gr.Column(scale=3):
+                        video_output = gr.Video(
+                            label="Generated Video",
+                            autoplay=True
                         )
                         
-                        randomize_seed = gr.Checkbox(
-                            label="Random seed",
-                            value=True
+                        video_status = gr.Textbox(
+                            label="Status",
+                            interactive=False,
+                            value="Ready to generate video..."
                         )
-                    
-                    video_negative_prompt = gr.Textbox(
-                        label="Negative Prompt",
-                        value=default_negative_prompt,
-                        lines=2
-                    )
-                    
-                    generate_video_btn = gr.Button(
-                        "🎬 Generate Video",
-                        variant="primary",
-                        size="lg"
-                    )
-                
-                with gr.Column(scale=1):
-                    video_output = gr.Video(
-                        label="Generated Video",
-                        autoplay=True,
-                        height=400
-                    )
-                    
-                    video_status = gr.Textbox(
-                        label="Status",
-                        interactive=False,
-                        value="Ready to generate video..."
-                    )
-    
-    # Event Handlers
-    def on_image_generated(prompt, img1, img2):
-        img, status, state_img = process_images(prompt, img1, img2)
-        if img:
-            return img, status, state_img, gr.update(visible=True)
-        return None, status, None, gr.update(visible=False)
-    
-    def send_image_to_video(img):
-        if img:
-            return img, "Image loaded! Ready to generate video."
-        return None, "No image to send."
-    
-    # Connect events
-    generate_img_btn.click(
-        fn=on_image_generated,
-        inputs=[style_prompt, image1, image2],
-        outputs=[output_image, img_status, generated_image_state, send_to_video_btn]
-    )
-    
-    send_to_video_btn.click(
-        fn=send_image_to_video,
-        inputs=[generated_image_state],
-        outputs=[video_input_image, video_status]
-    )
-    
-    generate_video_btn.click(
-        fn=generate_video,
-        inputs=[
-            video_input_image,
-            video_prompt,
-            steps_slider,
-            video_negative_prompt,
-            duration_input,
-            video_seed,
-            randomize_seed
-        ],
-        outputs=[video_output, video_seed, video_status]
-    )
+        
+        # Event Handlers
+        def on_image_generated(prompt, dummy1=None, dummy2=None):
+            img, status, state_img = process_images(prompt, dummy1, dummy2)
+            if img:
+                return img, status, state_img, gr.update(visible=True)
+            return None, status, None, gr.update(visible=False)
+        
+        def send_image_to_video(img):
+            if img:
+                return img, "✅ Image loaded! Ready to generate video."
+            return None, "❌ No image to send."
+        
+        # Connect events
+        generate_img_btn.click(
+            fn=on_image_generated,
+            inputs=[style_prompt],
+            outputs=[output_image, img_status, generated_image_state, send_to_video_btn]
+        )
+        
+        send_to_video_btn.click(
+            fn=send_image_to_video,
+            inputs=[generated_image_state],
+            outputs=[video_input_image, video_status]
+        )
+        
+        # Video generation with Replicate
+        generate_video_btn.click(
+            fn=generate_video_replicate,
+            inputs=[
+                video_input_image,
+                video_prompt,
+                gr.State(30),  # steps
+                gr.State("blurry, distorted"),  # negative prompt
+                gr.State(2.0),  # duration
+                video_seed,
+                randomize_seed
+            ],
+            outputs=[video_output, video_seed, video_status]
+        )
+        
+        # GPU test (optional)
+        test_gpu_btn.click(
+            fn=generate_video_gpu,
+            inputs=[
+                video_input_image,
+                video_prompt,
+                gr.State(30),
+                gr.State("blurry"),
+                gr.State(2.0),
+                video_seed,
+                randomize_seed
+            ],
+            outputs=[video_output, video_seed, video_status]
+        )
+        
+        # Footer
+        gr.HTML("""
+            <div style="margin-top: 2rem; padding: 1rem; background: #f8f9fa; border-radius: 10px;">
+                <p style="text-align: center; color: #6c757d;">
+                    Made with ❤️ using Gradio and Replicate API<br>
+                    <a href="https://replicate.com/" target="_blank">Get your API token</a> | 
+                    <a href="https://huggingface.co/spaces" target="_blank">Hugging Face Spaces</a>
+                </p>
+            </div>
+        """)
     
-    # Examples
-    gr.Examples(
-        examples=[
-            ["A majestic castle on a hilltop at sunset, fantasy art style"],
-            ["Cute robot in a flower garden, pixar animation style"],
-            ["Northern lights over a frozen lake, photorealistic"],
-            ["Ancient temple in a jungle, mysterious atmosphere"],
-        ],
-        inputs=[style_prompt],
-        label="Example Prompts"
-    )
+    return demo
 
-# Launch the app
+# Launch
 if __name__ == "__main__":
-    print("Starting Nano Banana + Video app...")
-    print("Make sure to set REPLICATE_API_TOKEN in your Space settings!")
+    print("=" * 50)
+    print("Starting Nano Banana + Video Application")
+    print("=" * 50)
+    
+    if check_api_token():
+        print("✅ Replicate API token found")
+    else:
+        print("⚠️ REPLICATE_API_TOKEN not found")
+        print("Please add it in Settings > Repository secrets")
+    
+    print("=" * 50)
     
+    # Create and launch the interface
+    demo = create_interface()
     demo.launch(
-        share=False,
-        show_error=True
+        show_error=True,
+        share=False
     )