Update app.py

app.py

@@ -1,21 +1,11 @@
 import gradio as gr
-import google.generativeai as genai
+import fal_client
 import os
 from typing import Optional, List
 from huggingface_hub import whoami
-from PIL import Image
-import tempfile
-import io # Import io for handling in-memory binary streams
 
-# --- Google Gemini API Configuration ---
-# Set your Google API key as an environment variable
-GOOGLE_API_KEY = os.getenv("GOOGLE_API_KEY", "")
-if not GOOGLE_API_KEY:
-    raise ValueError("GOOGLE_API_KEY environment variable not set.")
-genai.configure(api_key=GOOGLE_API_KEY)
-
-# --- Define the correct model name ---
-GEMINI_MODEL_NAME = 'gemini-2.5-flash-image-preview'
+FAL_KEY = os.getenv("FAL_KEY", "")
+fal_client.api_key = FAL_KEY
 
 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."""
@@ -23,105 +13,56 @@ def verify_pro_status(token: Optional[gr.OAuthToken]) -> bool:
         return False
     try:
         user_info = whoami(token=token.token)
+
+        # Case 1: User is PRO
         if user_info.get("isPro", False):
             return True
+
+        # Case 2: User is in any enterprise org
         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
 
 # --- Backend Generation Functions ---
 
-def run_single_image_logic(prompt: str, image_path: Optional[str] = None) -> str:
-    """Handles text-to-image or single image-to-image using Google Gemini."""
-    try:
-        model = genai.GenerativeModel(GEMINI_MODEL_NAME) # Use the defined model name
-        contents = [prompt]
-        if image_path:
-            input_image = Image.open(image_path)
-            contents.append(input_image)
-
-        response = model.generate_content(contents)
-        
-        # Access the image data correctly based on the response structure
-        # Assuming the generated content might be in response.candidates[0].content.parts[0].inline_data.data
-        # Or direct from response.parts if it's a single part with inline_data
-        
-        image_data = None
-        if hasattr(response, 'parts') and response.parts:
-            for part in response.parts:
-                if hasattr(part, 'inline_data') and hasattr(part.inline_data, 'data'):
-                    image_data = part.inline_data.data
-                    break
-        elif 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'):
-                            image_data = part.inline_data.data
-                            break
-                if image_data:
-                    break
-        
-        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(io.BytesIO(image_data))
-        with tempfile.NamedTemporaryFile(delete=False, suffix=".png") as tmpfile:
-            pil_image.save(tmpfile.name)
-            return tmpfile.name
-
-    except Exception as e:
-        raise gr.Error(f"Image generation failed: {e}")
-
+def run_single_image_logic(prompt: str, image: Optional[str] = None) -> str:
+    """Handles text-to-image or single image-to-image and returns a single URL string."""
+    if image:
+        image_url = fal_client.upload_file(image)
+        result = fal_client.run(
+            "fal-ai/nano-banana/edit",
+            # CORRECTED: The 'edit' endpoint always expects 'image_urls' as a list.
+            arguments={"prompt": prompt, "image_urls": [image_url]},
+        )
+    else:
+        result = fal_client.run(
+            "fal-ai/nano-banana", arguments={"prompt": prompt}
+        )
+    return result["images"][0]["url"]
 
 def run_multi_image_logic(prompt: str, images: List[str]) -> str:
     """
-    Handles multi-image editing by sending a list of images and a prompt.
+    Handles multi-image editing by sending a list of URLs in a single API call.
     """
     if not images:
         raise gr.Error("Please upload at least one image in the 'Multiple Images' tab.")
-    
-    try:
-        model = genai.GenerativeModel(GEMINI_MODEL_NAME) # Use the defined model name
-        # The prompt should be the last part of the contents list
-        contents = [Image.open(image_path[0]) for image_path in images]
-        contents.append(prompt)
-
-        response = model.generate_content(contents)
-        
-        image_data = None
-        if hasattr(response, 'parts') and response.parts:
-            for part in response.parts:
-                if hasattr(part, 'inline_data') and hasattr(part.inline_data, 'data'):
-                    image_data = part.inline_data.data
-                    break
-        elif 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'):
-                            image_data = part.inline_data.data
-                            break
-                if image_data:
-                    break
-
-        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(io.BytesIO(image_data))
-        with tempfile.NamedTemporaryFile(delete=False, suffix=".png") as tmpfile:
-            pil_image.save(tmpfile.name)
-            return tmpfile.name
-            
-    except Exception as e:
-        raise gr.Error(f"Image generation failed: {e}")
-
+    print(images)
+    image_urls = [fal_client.upload_file(image_path[0]) for image_path in images]
+    result = fal_client.run(
+        "fal-ai/nano-banana/edit",
+        arguments={
+            "prompt": prompt,
+            "image_urls": image_urls,
+            "num_images": 1
+        },
+    )
+    return result["images"][0]["url"]
 
 # --- Gradio App UI ---
 css = '''
@@ -131,10 +72,18 @@ css = '''
 #output{margin-top: 25px}
 .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}
 '''
 with gr.Blocks(theme=gr.themes.Citrus(), css=css) as demo:
-    gr.HTML("<h1 style='text-align:center'>Image Generation with Google Gemini</h1>")
-    gr.HTML("<h3 style='text-align:center'>Hugging Face PRO users can use Google's Gemini 2.5 Flash Image Preview on this Space. <a href='https://huggingface.co/pro' target='_blank'>Subscribe to PRO</a></h3>", elem_id="sub_title")
+    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='https://huggingface.co/pro' target='_blank'>Subscribe to PRO</a></h3>", elem_id="sub_title")
 
     pro_message = gr.Markdown(visible=False)
     main_interface = gr.Column(visible=False)