Swap to Google

app.py

@@ -1,11 +1,22 @@
 import gradio as gr
-import fal_client
+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
 
-FAL_KEY = os.getenv("FAL_KEY", "")
-fal_client.api_key = FAL_KEY
+# --- 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.")
+
+genai.configure(api_key=GOOGLE_API_KEY) 
+client = genai.Client() 
+
+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."""
@@ -13,56 +24,83 @@ 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: 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}
+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 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:
+        contents = [prompt]
+        if image_path:
+            input_image = Image.open(image_path)
+            contents.append(input_image)
+
+        response = client.models.generate_content( 
+            model=GEMINI_MODEL_NAME,
+            contents=contents,
         )
-    return result["images"][0]["url"]
+        
+        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)
+            return tmpfile.name
+
+    except Exception as e:
+        raise gr.Error(f"Image generation failed: {e}")
+
 
 def run_multi_image_logic(prompt: str, images: List[str]) -> str:
     """
-    Handles multi-image editing by sending a list of URLs in a single API call.
+    Handles multi-image editing by sending a list of images and a prompt.
     """
     if not images:
         raise gr.Error("Please upload at least one image in the 'Multiple Images' tab.")
-    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"]
+    
+    try:
+        contents = [Image.open(image_path[0]) for image_path in images]
+        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.")
+            
+        pil_image = Image.open(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}")
+
 
 # --- Gradio App UI ---
 css = '''
@@ -72,18 +110,10 @@ 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('''
-    <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")
+    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")
 
     pro_message = gr.Markdown(visible=False)
     main_interface = gr.Column(visible=False)