app.py

import gradio as gr
from gradio_image_annotation import image_annotator
import fal_client
from PIL import Image
import io
import base64
import numpy as np
import os

def process_images(annotated_image, second_image, user_api_key=None, progress=gr.Progress()):
    """
    Process the annotated image and second image using fal API
    """
    # Check if annotated_image is provided
    if annotated_image is None:
        return None, "Please provide the first image and draw an annotation box"
    
    # Check if second_image is provided (could be None or numpy array)
    if second_image is None or (isinstance(second_image, np.ndarray) and second_image.size == 0):
        return None, "Please provide the second image"
    
    # Check if annotation box exists
    if not annotated_image.get("boxes") or len(annotated_image["boxes"]) == 0:
        return None, "Please draw an annotation box on the first image"
    
    # Extract bounding box coordinates
    box = annotated_image["boxes"][0]  # Get the first (and only) box
    xmin = box.get("xmin")
    ymin = box.get("ymin")
    xmax = box.get("xmax")
    ymax = box.get("ymax")
    
    # Construct the dynamic prompt with the actual box coordinates
    prompt = f"""add the <central object in the second image> in the first image only inside an imaginary box defined by pixels values "xmin": {xmin}, "ymin": {ymin}, "xmax": {xmax}, "ymax": {ymax}. Take care of shadows, lighting, style, and general concept of objects as per the first image."""
    
    progress(0.2, desc="Gradio is preparing your images...")
    
    try:
        # Set API key - prioritize user input, then environment variable
        original_key = os.environ.get("FAL_KEY", "")
        
        if user_api_key and user_api_key.strip():
            # Use user-provided key
            os.environ["FAL_KEY"] = user_api_key.strip()
            api_key_source = "user-provided"
        elif original_key:
            # Use environment variable (secret)
            api_key_source = "environment"
        else:
            # No API key available
            return None, "⚠️ No FAL API key found. Please either:\n1. Duplicate this app and set your FAL_KEY as a secret, or\n2. Enter your FAL API key in the field provided above."
        
        # Convert first image to file for upload
        first_img = annotated_image["image"]
        if isinstance(first_img, np.ndarray):
            # Convert numpy array to PIL Image
            first_img_pil = Image.fromarray(first_img.astype('uint8'))
            # Save to bytes
            img1_bytes = io.BytesIO()
            first_img_pil.save(img1_bytes, format='PNG')
            img1_bytes.seek(0)
            uploaded_file1 = fal_client.upload(img1_bytes.getvalue(), "image/png")
        elif isinstance(first_img, str):
            # If it's a file path
            uploaded_file1 = fal_client.upload_file(first_img)
        else:
            # If it's already a PIL Image
            img1_bytes = io.BytesIO()
            first_img.save(img1_bytes, format='PNG')
            img1_bytes.seek(0)
            uploaded_file1 = fal_client.upload(img1_bytes.getvalue(), "image/png")
        
        # Convert second image to file for upload
        if isinstance(second_image, np.ndarray):
            second_img_pil = Image.fromarray(second_image.astype('uint8'))
            img2_bytes = io.BytesIO()
            second_img_pil.save(img2_bytes, format='PNG')
            img2_bytes.seek(0)
            uploaded_file2 = fal_client.upload(img2_bytes.getvalue(), "image/png")
        elif isinstance(second_image, str):
            uploaded_file2 = fal_client.upload_file(second_image)
        else:
            img2_bytes = io.BytesIO()
            second_image.save(img2_bytes, format='PNG')
            img2_bytes.seek(0)
            uploaded_file2 = fal_client.upload(img2_bytes.getvalue(), "image/png")
        
        progress(0.4, desc="Processing with nano-banana...")
        
        # Setup progress callback
        def on_queue_update(update):
            if isinstance(update, fal_client.InProgress):
                # InProgress updates don't have a progress attribute, just show we're processing
                progress(0.6, desc="nano-banana is working on your image...")
                # Optionally log any messages if they exist
                if hasattr(update, 'logs') and update.logs:
                    for log in update.logs:
                        print(log.get("message", ""))
        
        # Call fal API with the dynamic prompt including box coordinates
        result = fal_client.subscribe(
            "fal-ai/nano-banana/edit",
            arguments={
                "prompt": prompt,
                "image_urls": [f"{uploaded_file1}", f"{uploaded_file2}"]
            },
            with_logs=True,
            on_queue_update=on_queue_update,
        )
        
        progress(0.95, desc="Finalizing...")
        
        # Extract the result image URL
        if result and "images" in result and len(result["images"]) > 0:
            output_url = result["images"][0]["url"]
            description = result.get("description", "Image processed successfully!")
            progress(1.0, desc="Complete!")
            return output_url, description
        else:
            return None, "Failed to generate image. Please check your API key or try again."
            
    except Exception as e:
        error_message = str(e).lower()
        
        # Check for authentication errors
        if "401" in error_message or "unauthorized" in error_message or "api key" in error_message:
            return None, f"⚠️ API Authentication Error: Invalid or missing FAL API key.\n\nPlease either:\n1. Duplicate this app and set your FAL_KEY as a secret, or\n2. Enter your valid FAL API key in the field provided above.\n\nGet your API key at: https://fal.ai"
        
        # Check for rate limit errors
        elif "429" in error_message or "rate limit" in error_message:
            return None, "⚠️ Rate limit exceeded. Please wait a moment and try again, or use your own API key for higher limits."
        
        # Check for server errors
        elif "500" in error_message or "502" in error_message or "503" in error_message:
            return None, f"⚠️ FAL API server error. The service might be temporarily unavailable.\n\nPlease either:\n1. Try again in a few moments, or\n2. Use your own API key by entering it in the field above.\n\nError details: {str(e)}"
        
        # Generic error with fallback message
        else:
            return None, f"⚠️ Error occurred: {str(e)}\n\nIf the error persists, please either:\n1. Duplicate this app and set your FAL_KEY as a secret, or\n2. Enter your FAL API key in the field provided above.\n\nGet your API key at: https://fal.ai"
    
    finally:
        # Restore original API key if we temporarily changed it
        if user_api_key and user_api_key.strip():
            if original_key:
                os.environ["FAL_KEY"] = original_key
            else:
                os.environ.pop("FAL_KEY", None)


# Create the Gradio interface
with gr.Blocks(theme='ocean') as demo:
    # Add navigation bar
    navbar = gr.Navbar(
        value=[
            ("Documentation", "https://docs.fal.ai"),
            ("FAL.AI nano-banana", "https://fal.ai/models/fal-ai/nano-banana/edit/api"),
            ("Learn more about Gradio Navbar", "https://www.gradio.app/guides/multipage-apps#customizing-the-navbar")
        ],
        visible=True,
        main_page_name="🎨 guided nano banana"
    )
    
    # Add the animated banner
    gr.HTML(
        """
        <style>
            .animation-container {
                width: 100%;
                height: 400px;
                background: linear-gradient(135deg, #fff4e6 0%, #ffe8cc 25%, #ffeaa7 50%, #fdcb6e 75%, #ffecb3 100%);
                border-radius: 10px;
                margin-bottom: 20px;
                position: relative;
                overflow: hidden;
                display: flex;
                justify-content: center;
                align-items: center;
            }
            
            .mini-container {
                position: relative;
                width: 600px;
                height: 300px;
                perspective: 1000px;
                transform: scale(1.0);
            }
            
            .mini-image-wrapper {
                position: absolute;
                width: 300px;
                height: 300px;
                overflow: hidden;
                box-shadow: 0 10px 30px rgba(0,0,0,0.2);
                transition: all 1.5s cubic-bezier(0.4, 0, 0.2, 1);
                background: linear-gradient(135deg, rgba(255,255,255,0.5) 0%, rgba(255,255,255,0.3) 100%);
                display: flex;
                align-items: center;
                justify-content: center;
            }
            
            .mini-image-wrapper img {
                width: 100%;
                height: 100%;
                object-fit: contain;
            }
            
            .mini-left-image {
                left: 0;
                border-radius: 15px 0 0 15px;
            }
            
            .mini-right-image {
                right: 0;
                border-radius: 0 15px 15px 0;
            }
            
            .mini-result-image {
                width: 600px;
                height: 300px;
                position: absolute;
                left: 0;
                top: 0;
                border-radius: 15px;
                box-shadow: 0 15px 40px rgba(0,0,0,0.25);
                background: linear-gradient(135deg, rgba(255,255,255,0.5) 0%, rgba(255,255,255,0.3) 100%);
                display: flex;
                align-items: center;
                justify-content: center;
            }
            
            .mini-result-image img {
                width: 100%;
                height: 100%;
                object-fit: contain;
                border-radius: 15px;
            }
            
            /* Animations for Set 1 */
            .mini-left-image.set1 {
                animation: miniLeftSet1 18s infinite;
            }
            
            .mini-right-image.set1 {
                animation: miniRightSet1 18s infinite;
            }
            
            .mini-result-image.set1 {
                animation: miniResultSet1 18s infinite;
            }
            
            /* Animations for Set 2 */
            .mini-left-image.set2 {
                animation: miniLeftSet2 18s infinite;
            }
            
            .mini-right-image.set2 {
                animation: miniRightSet2 18s infinite;
            }
            
            .mini-result-image.set2 {
                animation: miniResultSet2 18s infinite;
            }
            
            /* Animations for Set 3 */
            .mini-left-image.set3 {
                animation: miniLeftSet3 18s infinite;
            }
            
            .mini-right-image.set3 {
                animation: miniRightSet3 18s infinite;
            }
            
            .mini-result-image.set3 {
                animation: miniResultSet3 18s infinite;
            }
            
            /* Set 1 Keyframes (0-6s) */
            @keyframes miniLeftSet1 {
                0% { transform: translateX(-120px) rotateY(15deg); opacity: 0; }
                2% { transform: translateX(0) rotateY(0); opacity: 1; }
                16% { transform: translateX(0) rotateY(0); opacity: 1; }
                22% { transform: translateX(0) rotateY(0); opacity: 0; }
                33%, 100% { transform: translateX(-120px) rotateY(15deg); opacity: 0; }
            }
            
            @keyframes miniRightSet1 {
                0% { transform: translateX(120px) rotateY(-15deg); opacity: 0; }
                2% { transform: translateX(0) rotateY(0); opacity: 1; }
                16% { transform: translateX(0) rotateY(0); opacity: 1; }
                22% { transform: translateX(0) rotateY(0); opacity: 0; }
                33%, 100% { transform: translateX(120px) rotateY(-15deg); opacity: 0; }
            }
            
            @keyframes miniResultSet1 {
                0%, 16% { opacity: 0; transform: scale(0.8); z-index: -1; }
                22% { opacity: 1; transform: scale(1.05); z-index: 10; }
                27% { opacity: 1; transform: scale(1); z-index: 10; }
                30% { opacity: 0; transform: scale(0.8); z-index: -1; }
                100% { opacity: 0; transform: scale(0.8); z-index: -1; }
            }
            
            /* Set 2 Keyframes (6-12s) */
            @keyframes miniLeftSet2 {
                0%, 33% { transform: translateX(-120px) rotateY(15deg); opacity: 0; }
                35% { transform: translateX(0) rotateY(0); opacity: 1; }
                50% { transform: translateX(0) rotateY(0); opacity: 1; }
                55% { transform: translateX(0) rotateY(0); opacity: 0; }
                66%, 100% { transform: translateX(-120px) rotateY(15deg); opacity: 0; }
            }
            
            @keyframes miniRightSet2 {
                0%, 33% { transform: translateX(120px) rotateY(-15deg); opacity: 0; }
                35% { transform: translateX(0) rotateY(0); opacity: 1; }
                50% { transform: translateX(0) rotateY(0); opacity: 1; }
                55% { transform: translateX(0) rotateY(0); opacity: 0; }
                66%, 100% { transform: translateX(120px) rotateY(-15deg); opacity: 0; }
            }
            
            @keyframes miniResultSet2 {
                0%, 50% { opacity: 0; transform: scale(0.8); z-index: -1; }
                55% { opacity: 1; transform: scale(1.05); z-index: 10; }
                61% { opacity: 1; transform: scale(1); z-index: 10; }
                63% { opacity: 0; transform: scale(0.8); z-index: -1; }
                100% { opacity: 0; transform: scale(0.8); z-index: -1; }
            }
            
            /* Set 3 Keyframes (12-18s) */
            @keyframes miniLeftSet3 {
                0%, 66% { transform: translateX(-120px) rotateY(15deg); opacity: 0; }
                69% { transform: translateX(0) rotateY(0); opacity: 1; }
                83% { transform: translateX(0) rotateY(0); opacity: 1; }
                88% { transform: translateX(0) rotateY(0); opacity: 0; }
                96%, 100% { transform: translateX(-120px) rotateY(15deg); opacity: 0; }
            }
            
            @keyframes miniRightSet3 {
                0%, 66% { transform: translateX(120px) rotateY(-15deg); opacity: 0; }
                69% { transform: translateX(0) rotateY(0); opacity: 1; }
                83% { transform: translateX(0) rotateY(0); opacity: 1; }
                88% { transform: translateX(0) rotateY(0); opacity: 0; }
                96%, 100% { transform: translateX(120px) rotateY(-15deg); opacity: 0; }
            }
            
            @keyframes miniResultSet3 {
                0%, 83% { opacity: 0; transform: scale(0.8); z-index: -1; }
                88% { opacity: 1; transform: scale(1.05); z-index: 10; }
                94% { opacity: 1; transform: scale(1); z-index: 10; }
                96% { opacity: 0; transform: scale(0.8); z-index: -1; }
                100% { opacity: 0; transform: scale(0.8); z-index: -1; }
            }
            
            .mini-progress-bar {
                position: absolute;
                bottom: 20px;
                left: 50%;
                transform: translateX(-50%);
                width: 300px;
                height: 4px;
                background: rgba(255, 165, 0, 0.2);
                border-radius: 2px;
                overflow: hidden;
            }
            
            .mini-progress-fill {
                height: 100%;
                background: linear-gradient(90deg, #ffa502 0%, #ff6348 100%);
                border-radius: 2px;
                animation: miniProgressCycle 18s linear infinite;
            }
            
            @keyframes miniProgressCycle {
                0% { width: 0%; }
                100% { width: 100%; }
            }
            
            .animation-title {
                position: absolute;
                top: 20px;
                left: 50%;
                transform: translateX(-50%);
                color: #ff6348;
                font-size: 18px;
                font-weight: 700;
                text-shadow: 0 2px 4px rgba(255,99,72,0.2);
                z-index: 100;
            }
        </style>
        
        <div class="animation-container">
            <div class="animation-title">Image Fusion Magic ✨</div>
            <div class="mini-container">
                <!-- Set 1 Images -->
                <div class="image-wrapper left-image set1">
                    <img src="https://cdn-uploads.huggingface.co/production/uploads/60d2dc1007da9c17c72708f8/aPCNIVXzUkwGrjDvm4a2q.png" alt="Set1 Left" id="set1LeftImg">
                </div>
                <div class="image-wrapper right-image set1">
                    <img src="https://cdn-uploads.huggingface.co/production/uploads/60d2dc1007da9c17c72708f8/yhuEbtc1s2kC9OUdcIPL2.png" alt="Set1 Right" id="set1RightImg">
                </div>
                <div class="result-image set1">
                    <div class="glow set1"></div>
                    <img src="https://cdn-uploads.huggingface.co/production/uploads/60d2dc1007da9c17c72708f8/Kk5WkJOsbCexc_14HYsVf.jpeg" alt="Set1 Result" id="set1ResultImg">
                </div>

                <!-- Set 2 Images -->
                <div class="image-wrapper left-image set2">
                    <img src="https://cdn-uploads.huggingface.co/production/uploads/60d2dc1007da9c17c72708f8/c4RM2XbJKLOiT5ePGkIoF.png" alt="Set2 Left" id="set2LeftImg">
                </div>
                <div class="image-wrapper right-image set2">
                    <img src="https://cdn-uploads.huggingface.co/production/uploads/60d2dc1007da9c17c72708f8/Eiv6EYFjpHEVHNRYPxpLA.png" alt="Set2 Right" id="set2RightImg">
                </div>
                <div class="result-image set2">
                    <div class="glow set2"></div>
                    <img src="https://cdn-uploads.huggingface.co/production/uploads/60d2dc1007da9c17c72708f8/zWCbcioy4Ary_1_x_0-5a.jpeg" alt="Set2 Result" id="set2ResultImg">
                </div>

                <!-- Set 3 Images -->
                <div class="image-wrapper left-image set3">
                    <img src="https://cdn-uploads.huggingface.co/production/uploads/60d2dc1007da9c17c72708f8/Yk5FqD9brcUNiub2D94K9.png" alt="Set3 Left" id="set3LeftImg">
                </div>
                <div class="image-wrapper right-image set3">
                    <img src="https://cdn-uploads.huggingface.co/production/uploads/60d2dc1007da9c17c72708f8/AYdhb7p7vD5EecbNUqhHb.png" alt="Set3 Right" id="set3RightImg">
                </div>
                <div class="result-image set3">
                    <div class="glow set3"></div>
                    <img src="https://cdn-uploads.huggingface.co/production/uploads/60d2dc1007da9c17c72708f8/99wlXmLCptDGnYkE6IOfE.jpeg" alt="Set3 Result" id="set3ResultImg">
                </div>
                
                <!-- Progress bar -->
                <div class="mini-progress-bar">
                    <div class="mini-progress-fill"></div>
                </div>
            </div>
        </div>
        """
    )
    
    gr.HTML(
        """
        <h1><center>Guide Your Nano Banana👉🍌</center></h1>
        
        <b>How to use:</b><br>
        1. Upload or capture the first image and draw a box where you want to place an object<br>
        2. Upload the second image containing the object you want to insert<br>
        3. Click "Generate Composite Image" and wait for the Gradio and Nano-Banana to blend the images<br>
        
        The Gradio app will intelligently place the object from the second image into the boxed area of the first image,
        taking care of lighting, shadows, and proper integration.
        """
    )
    
    # API Key input section
    with gr.Row():
        with gr.Column():
            with gr.Accordion("🔑 API Configuration (Optional)", open=False):
                gr.Markdown(
                    """
                    **Note:** If you're experiencing API errors or want to use your own FAL account:
                    - Enter your FAL API key below, or
                    - [Duplicate this Space](https://huggingface.co/spaces) and set FAL_KEY as a secret
                    - Get your API key at [fal.ai](https://fal.ai)
                    """
                )
                api_key_input = gr.Textbox(
                    label="FAL API Key",
                    placeholder="Enter your FAL key (optional)",
                    type="password",
                    interactive=True,
                    info="Your key will be used only for this session and won't be stored"
                )
    
    with gr.Row():
        with gr.Column(scale=1):
            with gr.Row():
                with gr.Column(scale=1):
                    gr.Markdown("### Step 1: Annotate First Image")
                    # Image annotator for first image
                    from gradio_image_annotation import image_annotator
                    #first_image = ImageAnnotator(
                    first_image = image_annotator(
                        value=None,
                        label="Draw a box where you want to place the object",
                        image_type="pil",
                        single_box=True,  # Only allow one box
                        disable_edit_boxes=True,  
                        show_download_button=False,
                        show_share_button=False,
                        box_thickness=3,
                        box_selected_thickness=4,
                        show_label=True,
                        #image_mode="RGB",
                        #box_min_size=20,
                    )
                    
                with gr.Column(scale=1):
                    gr.Markdown("### Step 2: Upload Second Image")
                    # Regular image input for second image
                    second_image = gr.Image(
                        label="Image containing the object to insert",
                        type="numpy",
                        height=400,
                    )
                    # Generate button
                    generate_btn = gr.Button("Step 3: 🚀 Generate Composite Image", variant="primary", size="lg")
    
    # Output section
    with gr.Column():
        output_image = gr.Image(
            label="Generated Composite Image",
            type="filepath",
            height=500,
        )
        status_text = gr.Textbox(
            label="Status",
            placeholder="Results will appear here...",
            lines=3,
        )
        
    # Connect the button to the processing function
    generate_btn.click(
        fn=process_images,
        inputs=[first_image, second_image, api_key_input],
        outputs=[output_image, status_text],
        show_progress=True,
    )

with demo.route("Tips", "/tips"):
    gr.Markdown(
        """
        # ℹ️ Tips for Best Results
        - **Box Placement**: Draw the box exactly where you want the object to appear
        - **Image Quality**: Use high-resolution images for better results
        - **Object Selection**: The second image should clearly show the object you want to insert
        - **Lighting**: Images with similar lighting conditions work best
        - **Processing Time**: Generation typically takes 10-30 seconds
        - **API Key**: If you encounter errors, try using your own FAL API key
        """
    )

    # Different navbar for the Settings page
    navbar = gr.Navbar(
        visible=True,
        main_page_name="Home",
    )  
    
if __name__ == "__main__":
    demo.launch(ssr_mode=False)