Update app.py

app.py

@@ -14,7 +14,7 @@ if not GOOGLE_API_KEY:
     raise ValueError("GOOGLE_API_KEY environment variable not set.")
 
 client = genai.Client(
-    api_key=os.environ.get("GEMINI_API_KEY"),
+    api_key=os.environ.get("GOOGLE_API_KEY"),
 )
 
 GEMINI_MODEL_NAME = 'gemini-2.5-flash-image-preview'
@@ -45,13 +45,29 @@ def _extract_image_data_from_response(response) -> Optional[bytes]:
                         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."""
+def unified_image_generator(
+    prompt: str, 
+    images: Optional[List[str]] = None,
+    oauth_token: Optional[gr.OAuthToken] = None
+) -> str:
+    """
+    Handles all image generation tasks based on the number of input images.
+    - 0 images: Text-to-image
+    - 1+ images: Image-to-image (single or multi-modal)
+    """
+    if not verify_pro_status(oauth_token):
+        raise gr.Error("Access Denied. This service is for PRO users only.")
+
     try:
-        contents = [prompt]
-        if image_path:
-            input_image = Image.open(image_path)
-            contents.append(input_image)
+        # 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:
+                contents.append(Image.open(image_path))
+        
+        # Always add the prompt to the contents
+        contents.append(prompt)
 
         response = client.models.generate_content( 
             model=GEMINI_MODEL_NAME,
@@ -73,36 +89,6 @@ def run_single_image_logic(prompt: str, image_path: Optional[str] = None) -> str
         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 images and a prompt.
-    """
-    if not images:
-        raise gr.Error("Please upload at least one image in the 'Multiple Images' tab.")
-    
-    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 = '''
 #sub_title{margin-top: -35px !important}
@@ -129,17 +115,12 @@ with gr.Blocks(theme=gr.themes.Citrus(), css=css) as demo:
     with main_interface:
         with gr.Row():
             with gr.Column(scale=1):
-                active_tab_state = gr.State(value="single")
-                with gr.Tabs() as tabs:
-                    with gr.TabItem("Single Image", id="single") as single_tab:
-                        image_input = gr.Image(
-                            type="filepath",
-                            label="Input Image (Leave blank for text-to-image)"
-                        )
-                    with gr.TabItem("Multiple Images", id="multiple") as multi_tab:
-                        gallery_input = gr.Gallery(
-                            label="Input Images (drop all images here)", file_types=["image"]
-                        )
+                image_input_gallery = gr.Gallery(
+                    label="Input Image(s)",
+                    info="Upload one or more images here. Leave empty for text-to-image",
+                    file_types=["image"],
+                    height="auto"
+                )
                         
                 prompt_input = gr.Textbox(
                     label="Prompt",
@@ -154,37 +135,20 @@ with gr.Blocks(theme=gr.themes.Citrus(), css=css) as demo:
     
     login_button = gr.LoginButton()
     
-    # --- Event Handlers ---
-    def unified_generator(
-        prompt: str,
-        single_image: Optional[str],
-        multi_images: Optional[List[str]],
-        active_tab: str,
-        oauth_token: Optional[gr.OAuthToken] = None,
-    ) -> str:
-        if not verify_pro_status(oauth_token):
-            raise gr.Error("Access Denied. This service is for PRO users only.")
-        if active_tab == "multiple" and multi_images:
-            return run_multi_image_logic(prompt, multi_images)
-        else:
-            return run_single_image_logic(prompt, single_image)
-
-    single_tab.select(lambda: "single", None, active_tab_state)
-    multi_tab.select(lambda: "multiple", None, active_tab_state)
-
+    # --- Event Handlers (SIMPLIFIED) ---
     generate_button.click(
-        unified_generator,
-        inputs=[prompt_input, image_input, gallery_input, active_tab_state],
+        unified_image_generator,
+        inputs=[prompt_input, image_input_gallery], # Inputs are now just the prompt and the single gallery
         outputs=[output_image],
     )
 
     use_image_button.click(
-        lambda img: img, 
+        lambda img_path: [img_path] if img_path else None, 
         inputs=[output_image],
-        outputs=[image_input]
+        outputs=[image_input_gallery]
     )
 
-    # --- Access Control Logic ---
+    # --- Access Control Logic (UNCHANGED) ---
     def control_access(
         profile: Optional[gr.OAuthProfile] = None,
         oauth_token: Optional[gr.OAuthToken] = None