app.py

import gradio as gr
from gradio_client import Client, handle_file
from google import genai 
from google.genai import types 
import os
from typing import Optional, List
from huggingface_hub import whoami
from PIL import Image
from io import BytesIO
import tempfile

# --- Google Gemini API Configuration ---
GOOGLE_API_KEY = os.getenv("GOOGLE_API_KEY", "")
if not GOOGLE_API_KEY:
    raise ValueError("GOOGLE_API_KEY environment variable not set.")

client = genai.Client(
    api_key=os.environ.get("GOOGLE_API_KEY"),
)

GEMINI_MODEL_NAME = 'gemini-2.5-flash-image-preview'

def verify_pro_status(token: Optional[gr.OAuthToken]) -> bool:
    """Verifies if the user is a Hugging Face PRO user or part of an enterprise org."""
    if not token:
        return False
    try:
        user_info = whoami(token=token.token)
        if user_info.get("isPro", False):
            return True
        orgs = user_info.get("orgs", [])
        if any(org.get("isEnterprise", False) for org in orgs):
            return True
        return False
    except Exception as e:
        print(f"Could not verify user's PRO/Enterprise status: {e}")
        return False

def _extract_image_data_from_response(response) -> Optional[bytes]:
    """Helper to extract image data from the model's response."""
    if hasattr(response, 'candidates') and response.candidates:
        for candidate in response.candidates:
            if hasattr(candidate, 'content') and hasattr(candidate.content, 'parts') and candidate.content.parts:
                for part in candidate.content.parts:
                    if hasattr(part, 'inline_data') and hasattr(part.inline_data, 'data'):
                        return part.inline_data.data
    return None

def unified_image_generator(
    prompt: str, 
    images: Optional[List[str]] = None,
    oauth_token: Optional[gr.OAuthToken] = None
) -> tuple:
    """
    Handles all image generation tasks based on the number of input images.
    Returns: (output_image_path, video_button_visible, video_output_visible)
    """
    if not verify_pro_status(oauth_token):
        raise gr.Error("Access Denied. This service is for PRO users only.")

    try:
        # Dynamically build the 'contents' list for the API
        contents = []
        if images:
            # If there are images, open them and add to contents
            for image_path in images:
                print(image_path)
                contents.append(Image.open(image_path[0]))
        
        # Always add the prompt to the contents
        contents.append(prompt)

        response = client.models.generate_content( 
            model=GEMINI_MODEL_NAME,
            contents=contents,
        )
        
        image_data = _extract_image_data_from_response(response)
        
        if not image_data:
            raise ValueError("No image data found in the model response.")

        # Save the generated image to a temporary file to return its path
        pil_image = Image.open(BytesIO(image_data))
        with tempfile.NamedTemporaryFile(delete=False, suffix=".png") as tmpfile:
            pil_image.save(tmpfile.name)
            output_path = tmpfile.name
        
        # Determine if video button should be shown (only if exactly 1 input image)
        show_video_button = images and len(images) == 1
        
        # Return output image path, video button visibility, and hide video output
        return output_path, gr.update(visible=show_video_button), gr.update(visible=False)

    except Exception as e:
        raise gr.Error(f"Image generation failed: {e}")

def create_video_transition(
    input_image_gallery: List[str],
    prompt_input: str,
    output_image: str,
    oauth_token: Optional[gr.OAuthToken] = None
) -> tuple:
    """
    Creates a video transition between the input and output images.
    Returns: (video_path, video_visible)
    """
    if not verify_pro_status(oauth_token):
        raise gr.Error("Access Denied. This service is for PRO users only.")
    
    if not input_image_gallery or not output_image:
        raise gr.Error("Both input and output images are required for video creation.")
    
    try:
        video_client = Client("multimodalart/wan-2-2-first-last-frame", hf_token=oauth_token.token)
        
        input_image_path = input_image_gallery[0][0]
        
        result = video_client.predict(
            start_image_pil=handle_file(input_image_path),
            end_image_pil=handle_file(output_image),
            prompt=prompt_input,
            api_name="/generate_video"
        )
        print(result)
        return result["video"]
        
    except Exception as e:
        raise gr.Error(f"Video creation failed: {e}")

# --- Gradio App UI ---
css = '''
#sub_title{margin-top: -35px !important}
.tab-wrapper{margin-bottom: -33px !important}
.tabitem{padding: 0px !important}
.fillable{max-width: 980px !important}
.dark .progress-text {color: white}
.logo-dark{display: none}
.dark .logo-dark{display: block !important}
.dark .logo-light{display: none}
.grid-container img{object-fit: contain}
.grid-container {display: grid;grid-template-columns: repeat(2, 1fr)}
.grid-container:has(> .gallery-item:only-child) {grid-template-columns: 1fr}
#wan_ad p{text-align: center;padding: .5em}
'''

with gr.Blocks(theme=gr.themes.Citrus(), css=css) as demo:
    gr.HTML('''
    <img class="logo-dark" src='https://huggingface.co/spaces/multimodalart/nano-banana/resolve/main/nano_banana_pros.png' style='margin: 0 auto; max-width: 500px' />
    <img class="logo-light" src='https://huggingface.co/spaces/multimodalart/nano-banana/resolve/main/nano_banana_pros_light.png' style='margin: 0 auto; max-width: 500px' />
    ''')
            
    gr.HTML("<h3 style='text-align:center'>Hugging Face PRO users can use Google's Nano Banana (Gemini 2.5 Flash Image Preview) on this Space. <a href='http://huggingface.co/subscribe/pro?source=nana_banana' target='_blank'>Subscribe to PRO</a></h3>", elem_id="sub_title")

    pro_message = gr.Markdown(visible=False)
    main_interface = gr.Column(visible=False)

    with main_interface:
        with gr.Row():
            with gr.Column(scale=1):
                with gr.Group():
                    image_input_gallery = gr.Gallery(
                        label="Upload one or more images here. Leave empty for text-to-image",
                        file_types=["image"],
                        height="auto"
                    )
                            
                    prompt_input = gr.Textbox(
                        label="Prompt",
                        placeholder="Turns this photo into a masterpiece"
                    )
                    generate_button = gr.Button("Generate", variant="primary")

            with gr.Column(scale=1):
                output_image = gr.Image(label="Output", interactive=False, elem_id="output", type="filepath")
                use_image_button = gr.Button("♻️ Use this Image for Next Edit")
                create_video_button = gr.Button("Create a video between the two images 🎥", variant="primary", visible=False)
                with gr.Group(visible=False) as video_group:
                    video_output = gr.Video(label="Generated Video", show_download_button=True, autoplay=True)
                    gr.Markdown("Generate more with [Wan 2.2 first-last-frame](https://huggingface.co/spaces/multimodalart/wan-2-2-first-last-frame)", elem_id="wan_ad")
        gr.Markdown("## Thank you for being a PRO! 🤗")
    
    login_button = gr.LoginButton()
    
    # --- Event Handlers ---
    gr.on(
        triggers=[generate_button.click, prompt_input.submit],
        fn=lambda: [gr.update(visible=False), gr.update(visible=False)],
        inputs=[],
        outputs=[create_video_button, video_group],
    ).then(
        fn=unified_image_generator,
        inputs=[prompt_input, image_input_gallery],
        outputs=[output_image, create_video_button, video_group],
    )

    use_image_button.click(
        lambda img_path: [img_path] if img_path else None, 
        inputs=[output_image],
        outputs=[image_input_gallery]
    )
    
    # Video creation handler
    create_video_button.click(
        fn=lambda: gr.update(visible=True),
        inputs=[],
        outputs=[video_group],
    ).then(
        fn=create_video_transition,
        inputs=[image_input_gallery, prompt_input, output_image],
        outputs=[video_output],
    )

    # --- Access Control Logic ---
    def control_access(
        profile: Optional[gr.OAuthProfile] = None,
        oauth_token: Optional[gr.OAuthToken] = None
    ):
        if not profile:
            return gr.update(visible=False), gr.update(visible=False)
        if verify_pro_status(oauth_token):
            return gr.update(visible=True), gr.update(visible=False)
        else:
            message = (
                "## ✨ Exclusive Access for PRO Users\n\n"
                "Thank you for your interest! This app is available exclusively for our Hugging Face **PRO** members.\n\n"
                "To unlock this and many other cool stuff, please consider upgrading your account.\n\n"
                "### [**Become a PRO Today!**](http://huggingface.co/subscribe/pro?source=nana_banana)"
            )
            return gr.update(visible=False), gr.update(visible=True, value=message)

    demo.load(control_access, inputs=None, outputs=[main_interface, pro_message])

if __name__ == "__main__":
    demo.queue(max_size=None, default_concurrency_limit=None)
    demo.launch()