Update app.py

app.py

@@ -138,6 +138,7 @@ 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'
+GEMINI_PRO_MODEL_NAME = 'gemini-3-pro-image-preview'
 
 def verify_pro_status(token: Optional[Union[gr.OAuthToken, str]]) -> bool:
     """Verifies if the user is a Hugging Face PRO user or part of an enterprise org."""
@@ -261,7 +262,7 @@ def _generate_video_segment(input_image_path: str, output_image_path: str, promp
     )
     return result[0]["video"]
 
-def unified_image_generator(prompt: str, images: Optional[List[str]], previous_video_path: Optional[str], last_frame_path: Optional[str], aspect_ratio: str, manual_token: str, oauth_token: Optional[gr.OAuthToken]) -> tuple:
+def unified_image_generator(prompt: str, images: Optional[List[str]], previous_video_path: Optional[str], last_frame_path: Optional[str], aspect_ratio: str, use_pro_model: bool, manual_token: str, oauth_token: Optional[gr.OAuthToken]) -> tuple:
     if not (verify_pro_status(oauth_token) or verify_pro_status(manual_token)): 
         raise gr.Error("Access Denied.")
     
@@ -277,21 +278,35 @@ def unified_image_generator(prompt: str, images: Optional[List[str]], previous_v
         contents = [Image.open(image_path[0]) for image_path in images] if images else []
         contents.append(prompt)
         
+        # Select model based on checkbox
+        model_name = GEMINI_PRO_MODEL_NAME if use_pro_model else GEMINI_MODEL_NAME
+        
         # Create config with aspect ratio (omit image_config if Auto is selected)
-        if aspect_ratio == "Auto":
-            generate_content_config = types.GenerateContentConfig(
-                response_modalities=["IMAGE", "TEXT"],
-            )
-        else:
+        # Note: PRO model uses image_size instead of aspect_ratio
+        if use_pro_model:
+            # Gemini 3 Pro uses image_size parameter
             generate_content_config = types.GenerateContentConfig(
                 response_modalities=["IMAGE", "TEXT"],
                 image_config=types.ImageConfig(
-                    aspect_ratio=aspect_ratio,
+                    image_size="1K",  # Options: "1K", "2K", "4K"
                 ),
             )
+        else:
+            # Gemini 2.5 Flash uses aspect_ratio parameter
+            if aspect_ratio == "Auto":
+                generate_content_config = types.GenerateContentConfig(
+                    response_modalities=["IMAGE", "TEXT"],
+                )
+            else:
+                generate_content_config = types.GenerateContentConfig(
+                    response_modalities=["IMAGE", "TEXT"],
+                    image_config=types.ImageConfig(
+                        aspect_ratio=aspect_ratio,
+                    ),
+                )
         
         response = client.models.generate_content(
-            model=GEMINI_MODEL_NAME, 
+            model=model_name, 
             contents=contents,
             config=generate_content_config
         )
@@ -356,7 +371,7 @@ with gr.Blocks(theme=gr.themes.Citrus(), css=css) as demo:
     <img class="logo-dark" src='https://huggingface.co/spaces/multimodalart/nano-banana/resolve/main/nano_banana_pros.png' style='margin: 0 auto; max-width: 650px' />
     <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: 650px' />
     ''')
-    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")
+    gr.HTML("<h3 style='text-align:center'>Hugging Face PRO users can use Google's Nano Banana and Nano Banana PRO 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)
     
@@ -368,12 +383,23 @@ with gr.Blocks(theme=gr.themes.Citrus(), css=css) as demo:
             with gr.Column(scale=1):
                 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")
+                
+                # Model selection checkbox
+                use_pro_checkbox = gr.Checkbox(
+                    label="🌟 Use Nano Banana PRO (Gemini 3 Pro Image) - Higher quality, slower generation",
+                    value=False,
+                    interactive=True
+                )
+                
                 aspect_ratio_dropdown = gr.Dropdown(
-                    label="Aspect Ratio",
+                    label="Aspect Ratio (only for standard model)",
                     choices=["Auto", "1:1", "9:16", "16:9", "3:4", "4:3", "3:2", "2:3", "5:4", "4:5", "21:9"],
                     value="Auto",
                     interactive=True
                 )
+                
+                gr.Markdown("💡 **Tip:** PRO model uses fixed 1K resolution. Aspect ratio setting is ignored when PRO is enabled.")
+                
                 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")
@@ -392,7 +418,7 @@ with gr.Blocks(theme=gr.themes.Citrus(), css=css) as demo:
     gr.on(
         triggers=[generate_button.click, prompt_input.submit],
         fn=unified_image_generator,
-        inputs=[prompt_input, image_input_gallery, previous_video_state, last_frame_of_video_state, aspect_ratio_dropdown, manual_token],
+        inputs=[prompt_input, image_input_gallery, previous_video_state, last_frame_of_video_state, aspect_ratio_dropdown, use_pro_checkbox, manual_token],
         outputs=[output_image, create_video_button, extend_video_button, video_group],
         api_name=False
     )