app.py

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

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

# Set up Replicate API key
os.environ['REPLICATE_API_TOKEN'] = os.getenv('REPLICATE_API_TOKEN')

# Video Model Configuration
MAX_SEED = np.iinfo(np.int32).max
FIXED_FPS = 16

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

# ===========================
# 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 - exact same as example"""
    # Method 1: Try imgbb.com (most reliable)
    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,
            }
        )
        
        if response.status_code == 200:
            data = response.json()
            if data.get('success'):
                return data['data']['url']
    except:
        pass
    
    # Method 2: Try 0x0.st (simple and reliable)
    try:
        buffered = BytesIO()
        image.save(buffered, format="PNG")
        buffered.seek(0)
        
        files = {'file': ('image.png', buffered, 'image/png')}
        response = requests.post("https://0x0.st", files=files)
        
        if response.status_code == 200:
            return response.text.strip()
    except:
        pass
    
    # Method 3: 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}"

# ===========================
# Image Generation with google/nano-banana
# ===========================

def process_images(prompt, image1, image2=None):
    """Process images using google/nano-banana model - exact same logic as example"""
    if not image1:
        return None, "Please upload at least one image", None
    
    if not check_api_token():
        return None, "⚠️ Please set REPLICATE_API_TOKEN in Space settings", None
    
    try:
        image_urls = []
        
        # Upload images
        url1 = upload_image_to_hosting(image1)
        image_urls.append(url1)
        
        if image2:
            url2 = upload_image_to_hosting(image2)
            image_urls.append(url2)
        
        print(f"Running google/nano-banana with prompt: {prompt}")
        print(f"Image URLs: {image_urls}")
        
        # Run the model - exactly as in example
        output = replicate.run(
            "google/nano-banana",
            input={
                "prompt": prompt,
                "image_input": image_urls
            }
        )
        
        if output is None:
            return None, "No output received", None
        
        # Get the generated image - exact same handling as example
        img = None
        
        # Try method 1
        try:
            if hasattr(output, 'read'):
                img_data = output.read()
                img = Image.open(BytesIO(img_data))
        except:
            pass
        
        # Try method 2
        if img is None:
            try:
                if hasattr(output, 'url'):
                    output_url = output.url()
                    response = requests.get(output_url, timeout=30)
                    if response.status_code == 200:
                        img = Image.open(BytesIO(response.content))
            except:
                pass
        
        # Try method 3
        if img is None:
            output_url = None
            if isinstance(output, str):
                output_url = output
            elif isinstance(output, list) and len(output) > 0:
                output_url = output[0]
            
            if output_url:
                response = requests.get(output_url, timeout=30)
                if response.status_code == 200:
                    img = Image.open(BytesIO(response.content))
        
        if img:
            return img, "✨ Image generated successfully! You can now create a video.", img
        else:
            return None, "Could not process output", None
        
    except Exception as e:
        error_msg = str(e)
        print(f"Error in process_images: {error_msg}")
        
        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: {str(e)[:100]}", None

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

def resize_image_for_video(image: Image.Image, target_width=None, target_height=None):
    """Resize image for video generation while maintaining aspect ratio"""
    
    # 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')
    
    # Get original dimensions
    orig_width, orig_height = image.size
    aspect_ratio = orig_width / orig_height
    
    # If no target dimensions specified, use original aspect ratio with constraints
    if target_width is None or target_height is None:
        # Determine if landscape or portrait
        if aspect_ratio > 1:  # Landscape
            target_width = min(1024, orig_width)
            target_height = int(target_width / aspect_ratio)
        else:  # Portrait or square
            target_height = min(1024, orig_height)
            target_width = int(target_height * aspect_ratio)
        
        # Ensure dimensions are divisible by 8 (required by many models)
        target_width = (target_width // 8) * 8
        target_height = (target_height // 8) * 8
        
        # Minimum size constraints
        target_width = max(256, target_width)
        target_height = max(256, target_height)
    
    # Resize image
    resized = image.resize((target_width, target_height), Image.LANCZOS)
    
    return resized, target_width, target_height

@spaces.GPU(duration=60)
def generate_video_gpu(
    input_image,
    prompt,
    steps,
    negative_prompt,
    duration_seconds,
    seed,
    randomize_seed,
    maintain_aspect_ratio
):
    """GPU-accelerated video generation"""
    try:
        # Clear GPU memory
        if torch.cuda.is_available():
            torch.cuda.empty_cache()
            gc.collect()
            
        # Simulate processing
        time.sleep(2)
        
        return None, seed, "🎬 GPU test completed 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,
    maintain_aspect_ratio=True
):
    """Generate video using Replicate API with aspect ratio preservation"""
    
    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:
        # Get image dimensions while maintaining aspect ratio
        if maintain_aspect_ratio:
            resized_image, video_width, video_height = resize_image_for_video(input_image)
            print(f"Video dimensions: {video_width}x{video_height} (maintaining aspect ratio)")
        else:
            # Default landscape dimensions
            resized_image, video_width, video_height = resize_image_for_video(input_image, 768, 512)
            print(f"Video dimensions: {video_width}x{video_height} (fixed landscape)")
        
        # Upload image
        img_url = upload_image_to_hosting(resized_image)
        
        current_seed = random.randint(0, MAX_SEED) if randomize_seed else int(seed)
        
        print("Generating video with Stable Video Diffusion...")
        
        # Use Stable Video Diffusion
        output = replicate.run(
            "stability-ai/stable-video-diffusion:3f0457e4619daac51203dedb472816fd4af51f3149fa7a9e0b5ffcf1b8172438",
            input={
                "input_image": img_url,
                "frames_per_second": FIXED_FPS,
                "motion_bucket_id": 127,  # Controls motion amount (0-255)
                "cond_aug": 0.02,  # Conditioning augmentation
                "decoding_t": min(14, int(duration_seconds * 7)),  # Number of frames
                "seed": current_seed,
                "sizing_strategy": "maintain_aspect_ratio"  # Preserve aspect ratio
            }
        )
        
        if output:
            # Download video
            video_url = output if isinstance(output, str) else 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, f"🎬 Video generated! ({video_width}x{video_height})"
        
        return None, seed, "Failed to generate video"
        
    except Exception as e:
        error_msg = str(e)
        if "authentication" in error_msg.lower():
            return None, seed, "❌ Invalid API token"
        else:
            return None, seed, f"Error: {error_msg[:200]}"

# ===========================
# Enhanced CSS (same as example)
# ===========================

css = """
.gradio-container {
    background: linear-gradient(135deg, #f5f7fa 0%, #c3cfe2 100%);
    font-family: -apple-system, BlinkMacSystemFont, 'Segoe UI', sans-serif;
    min-height: 100vh;
}
.header-container {
    background: linear-gradient(135deg, #ffd93d 0%, #ffb347 100%);
    padding: 2.5rem;
    border-radius: 24px;
    margin-bottom: 2.5rem;
    box-shadow: 0 20px 60px rgba(255, 179, 71, 0.25);
}
.logo-text {
    font-size: 3.5rem;
    font-weight: 900;
    color: #2d3436;
    text-align: center;
    margin: 0;
    letter-spacing: -2px;
}
.subtitle {
    color: #2d3436;
    text-align: center;
    font-size: 1.2rem;
    margin-top: 0.5rem;
    opacity: 0.9;
    font-weight: 600;
}
.main-content {
    background: rgba(255, 255, 255, 0.95);
    backdrop-filter: blur(20px);
    border-radius: 24px;
    padding: 2.5rem;
    box-shadow: 0 10px 40px rgba(0, 0, 0, 0.08);
    margin-bottom: 2rem;
}
.gr-button-primary {
    background: linear-gradient(135deg, #ffd93d 0%, #ffb347 100%) !important;
    border: none !important;
    color: #2d3436 !important;
    font-weight: 700 !important;
    font-size: 1.1rem !important;
    padding: 1.2rem 2rem !important;
    border-radius: 14px !important;
    transition: all 0.3s cubic-bezier(0.4, 0, 0.2, 1) !important;
    text-transform: uppercase;
    letter-spacing: 1px;
    width: 100%;
    margin-top: 1rem !important;
}
.gr-button-primary:hover {
    transform: translateY(-3px) !important;
    box-shadow: 0 15px 40px rgba(255, 179, 71, 0.35) !important;
}
.gr-button-secondary {
    background: linear-gradient(135deg, #667eea 0%, #764ba2 100%) !important;
    border: none !important;
    color: white !important;
    font-weight: 700 !important;
    font-size: 1.1rem !important;
    padding: 1.2rem 2rem !important;
    border-radius: 14px !important;
}
.section-title {
    font-size: 1.8rem;
    font-weight: 800;
    color: #2d3436;
    margin-bottom: 1rem;
    padding-bottom: 0.5rem;
    border-bottom: 3px solid #ffd93d;
}
.status-text {
    font-family: 'SF Mono', 'Monaco', monospace;
    color: #00b894;
    font-size: 0.9rem;
}
.image-container {
    border-radius: 14px !important;
    overflow: hidden;
    border: 2px solid #e1e8ed !important;
    background: #fafbfc !important;
}
footer {
    display: none !important;
}
"""

# ===========================
# Gradio Interface
# ===========================

def create_interface():
    with gr.Blocks(css=css, theme=gr.themes.Base()) as demo:
        # Shared state
        generated_image_state = gr.State(None)
        
        # Header
        with gr.Column(elem_classes="header-container"):
            gr.HTML("""
                <h1 class="logo-text">🍌 Nano Banana VIDEO</h1>
                <p class="subtitle">AI-Powered Image Style Transfer</p>
                <div style="display: flex; justify-content: center; align-items: center; gap: 10px; margin-top: 20px;">
                    <a href="https://huggingface.co/spaces/ginigen/Nano-Banana-PRO" target="_blank">
                        <img src="https://img.shields.io/static/v1?label=NANO%20BANANA&message=PRO&color=%230000ff&labelColor=%23800080&logo=HUGGINGFACE&logoColor=white&style=for-the-badge" alt="badge">
                    </a>                     
                    <a href="https://huggingface.co/spaces/openfree/Nano-Banana-Upscale" target="_blank">
                        <img src="https://img.shields.io/static/v1?label=NANO%20BANANA&message=UPSCALE&color=%230000ff&labelColor=%23800080&logo=GOOGLE&logoColor=white&style=for-the-badge" alt="Nano Banana Upscale">
                    </a>
                    <a href="https://huggingface.co/spaces/openfree/Free-Nano-Banana" target="_blank">
                        <img src="https://img.shields.io/static/v1?label=NANO%20BANANA&message=FREE&color=%230000ff&labelColor=%23800080&logo=GOOGLE&logoColor=white&style=for-the-badge" alt="Free Nano Banana">
                    </a>
                    <a href="https://huggingface.co/spaces/aiqtech/Nano-Banana-API" target="_blank">
                        <img src="https://img.shields.io/static/v1?label=NANO%20BANANA&message=API&color=%230000ff&labelColor=%23800080&logo=GOOGLE&logoColor=white&style=for-the-badge" alt="Nano Banana API">
                    </a>
                    <a href="https://discord.gg/openfreeai" target="_blank">
                        <img src="https://img.shields.io/static/v1?label=Discord&message=Openfree%20AI&color=%230000ff&labelColor=%23800080&logo=discord&logoColor=white&style=for-the-badge" alt="Discord Openfree AI">
                    </a>
                </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'}; 
                     padding: 12px; border-radius: 10px; margin: 15px 0;">
                    <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():
            # Tab 1: Image Generation
            with gr.TabItem("🎨 Step 1: Generate Image"):
                with gr.Column(elem_classes="main-content"):
                    gr.HTML('<h2 class="section-title">🎨 Image Style Transfer</h2>')
                    
                    with gr.Row(equal_height=True):
                        with gr.Column(scale=1):
                            style_prompt = gr.Textbox(
                                label="Style Description",
                                placeholder="Describe your style...",
                                lines=3,
                                value="Make the sheets in the style of the logo. Make the scene natural.",
                            )
                            
                            with gr.Row(equal_height=True):
                                image1 = gr.Image(
                                    label="Primary Image",
                                    type="pil",
                                    height=200,
                                    elem_classes="image-container"
                                )
                                image2 = gr.Image(
                                    label="Secondary Image (Optional)",
                                    type="pil",
                                    height=200,
                                    elem_classes="image-container"
                                )
                            
                            generate_img_btn = gr.Button(
                                "Generate Magic ✨",
                                variant="primary",
                                size="lg"
                            )
                        
                        with gr.Column(scale=1):
                            output_image = gr.Image(
                                label="Generated Result",
                                type="pil",
                                height=420,
                                elem_classes="image-container"
                            )
                            
                            img_status = gr.Textbox(
                                label="Status",
                                interactive=False,
                                lines=1,
                                elem_classes="status-text",
                                value="Ready to generate..."
                            )
                            
                            send_to_video_btn = gr.Button(
                                "Send to Video Generation →",
                                variant="secondary",
                                size="lg",
                                visible=False
                            )
            
            # Tab 2: Video Generation
            with gr.TabItem("🎬 Step 2: Generate Video"):
                with gr.Column(elem_classes="main-content"):
                    gr.HTML('<h2 class="section-title">🎬 Video Generation from Image</h2>')
                    
                    with gr.Row():
                        with gr.Column():
                            video_input_image = gr.Image(
                                type="pil", 
                                label="Input Image (from Step 1 or upload new)",
                                elem_classes="image-container"
                            )
                            video_prompt = gr.Textbox(
                                label="Animation Prompt", 
                                value=default_prompt_i2v,
                                lines=3
                            )
                            
                            with gr.Row():
                                duration_input = gr.Slider(
                                    minimum=1.0, 
                                    maximum=4.0, 
                                    step=0.5, 
                                    value=2.0, 
                                    label="Duration (seconds)"
                                )
                                
                                maintain_aspect = gr.Checkbox(
                                    label="Maintain Original Aspect Ratio",
                                    value=True
                                )
                            
                            with gr.Accordion("Advanced Settings", open=False):
                                video_negative_prompt = gr.Textbox(
                                    label="Negative Prompt", 
                                    value=default_negative_prompt, 
                                    lines=3
                                )
                                video_seed = gr.Slider(
                                    label="Seed", 
                                    minimum=0, 
                                    maximum=MAX_SEED, 
                                    step=1, 
                                    value=42
                                )
                                randomize_seed = gr.Checkbox(
                                    label="Randomize seed", 
                                    value=True
                                )
                                steps_slider = gr.Slider(
                                    minimum=10, 
                                    maximum=50, 
                                    step=5, 
                                    value=30, 
                                    label="Quality Steps"
                                )
                            
                            generate_video_btn = gr.Button(
                                "Generate Video 🎬",
                                variant="primary",
                                size="lg"
                            )
                        
                        with gr.Column():
                            video_output = gr.Video(
                                label="Generated Video", 
                                autoplay=True
                            )
                            video_status = gr.Textbox(
                                label="Status",
                                interactive=False,
                                lines=1,
                                elem_classes="status-text",
                                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 img, status, state_img, 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."
        
        # Image generation 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 tab
        send_to_video_btn.click(
            fn=send_image_to_video,
            inputs=[generated_image_state],
            outputs=[video_input_image, video_status]
        )
        
        # Video generation events
        generate_video_btn.click(
            fn=generate_video_replicate,
            inputs=[
                video_input_image,
                video_prompt,
                steps_slider,
                video_negative_prompt,
                duration_input,
                video_seed,
                randomize_seed,
                maintain_aspect
            ],
            outputs=[video_output, video_seed, video_status]
        )
    
    return demo

# Launch
if __name__ == "__main__":
    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(
        show_error=True,
        share=False
    )