app.py

import gradio as gr
import replicate
import os
from PIL import Image
import requests
from io import BytesIO
import time
import tempfile
import base64
import spaces
import torch
import numpy as np
import random
import gc

# ===========================
# Configuration
# ===========================

# Set up Replicate API key
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
MAX_SEED = np.iinfo(np.int32).max
FIXED_FPS = 8  # Reduced FPS
MIN_FRAMES_MODEL = 8
MAX_FRAMES_MODEL = 32  # Reduced max frames

default_prompt_i2v = "make this image come alive, smooth animation"
default_negative_prompt = "static, still, blurry, low quality"

# ===========================
# Image Processing Functions
# ===========================

def upload_image_to_hosting(image):
    """Upload image to hosting service"""
    try:
        buffered = BytesIO()
        image.save(buffered, format="PNG")
        buffered.seek(0)
        img_base64 = base64.b64encode(buffered.getvalue()).decode()
        
        response = requests.post(
            "https://api.imgbb.com/1/upload",
            data={
                'key': '6d207e02198a847aa98d0a2a901485a5',
                'image': img_base64,
            },
            timeout=30
        )
        
        if response.status_code == 200:
            data = response.json()
            if data.get('success'):
                return data['data']['url']
    except Exception as e:
        print(f"Upload failed: {e}")
    
    # Fallback to base64
    buffered = BytesIO()
    image.save(buffered, format="PNG")
    buffered.seek(0)
    img_base64 = base64.b64encode(buffered.getvalue()).decode()
    return f"data:image/png;base64,{img_base64}"

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 os.getenv('REPLICATE_API_TOKEN'):
        return None, "Please set REPLICATE_API_TOKEN in Space settings", None
    
    try:
        # Upload image
        url1 = upload_image_to_hosting(image1)
        
        # 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",
                "width": 1024,
                "height": 1024,
                "num_inference_steps": 25
            }
        )
        
        if output and isinstance(output, list) and len(output) > 0:
            img_url = output[0]
            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 None, "Could not process output", 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

# ===========================
# Video Generation Functions
# ===========================

def resize_image_for_video(image: Image.Image) -> Image.Image:
    """Resize image for video generation"""
    # Convert RGBA to RGB if necessary
    if image.mode == 'RGBA':
        background = Image.new('RGB', image.size, (255, 255, 255))
        background.paste(image, mask=image.split()[3])
        image = background
    
    # 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,
    duration_seconds=2.0,
    seed=42,
    randomize_seed=False,
):
    """Generate video using Replicate API with GPU"""
    
    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
        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)
        
        current_seed = random.randint(0, MAX_SEED) if randomize_seed else int(seed)
        
        # Use Replicate for video generation
        print("Generating video with Replicate...")
        output = replicate.run(
            VIDEO_MODEL_ID,
            input={
                "prompt": prompt,
                "image": img_url,
                "steps": int(steps),
                "fps": FIXED_FPS,
                "seconds": min(duration_seconds, 3),  # Limit to 3 seconds
                "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)
            
            response = requests.get(video_url, timeout=60)
            if response.status_code == 200:
                with tempfile.NamedTemporaryFile(suffix=".mp4", delete=False) as tmp_video:
                    tmp_video.write(response.content)
                    return tmp_video.name, current_seed, "🎬 Video generated successfully!"
        
        return None, seed, "Failed to generate video"
        
    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"

# ===========================
# CSS Styling
# ===========================

css = """
.gradio-container {
    max-width: 1200px !important;
    margin: 0 auto !important;
}
.header-container {
    background: linear-gradient(135deg, #ffd93d, #ffb347);
    padding: 2rem;
    border-radius: 15px;
    margin-bottom: 2rem;
    text-align: center;
}
.logo-text {
    font-size: 2.5rem;
    font-weight: bold;
    color: #2d3436;
}
.subtitle {
    color: #2d3436;
    font-size: 1.1rem;
    margin-top: 0.5rem;
}
.gr-button {
    font-size: 1rem !important;
    padding: 12px 24px !important;
}
.gr-button-primary {
    background: linear-gradient(135deg, #ffd93d, #ffb347) !important;
    border: none !important;
}
.gr-button-secondary {
    background: linear-gradient(135deg, #667eea, #764ba2) !important;
    color: white !important;
    border: none !important;
}
"""

# ===========================
# 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"
                    )
                    
                    with gr.Row():
                        image1 = gr.Image(
                            label="Reference Image (Optional)",
                            type="pil",
                            height=200
                        )
                        image2 = gr.Image(
                            label="Style Reference (Optional)",
                            type="pil",
                            height=200
                        )
                    
                    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():
                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)"
                        )
                        
                        steps_slider = gr.Slider(
                            minimum=10,
                            maximum=50,
                            step=5,
                            value=25,
                            label="Quality Steps"
                        )
                    
                    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
                        )
                    
                    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]
    )
    
    # 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"
    )

# Launch the app
if __name__ == "__main__":
    print("Starting Nano Banana + Video app...")
    print("Make sure to set REPLICATE_API_TOKEN in your Space settings!")
    
    demo.launch(
        share=False,
        show_error=True
    )