Update app.py

app.py

@@ -299,19 +299,22 @@ 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, model_selection: 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, model_selection: str, resolution: str, 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.")
     
     # Determine if using PRO model based on radio selection
     use_pro_model = (model_selection == "Nano Banana PRO")
     
+    # Calculate credit cost based on resolution (only for PRO model)
+    credits_to_use = get_credit_cost(resolution) if use_pro_model else 1
+    
     # Check rate limit
     username = get_username(oauth_token) or get_username(manual_token)
     if not username:
         raise gr.Error("Could not identify user.")
     
-    can_generate = check_and_update_usage(username, use_pro_model)
+    can_generate = check_and_update_usage(username, use_pro_model, credits_to_use)
     
     if not can_generate:
         # Check if user has quota on the other model
@@ -320,10 +323,16 @@ def unified_image_generator(prompt: str, images: Optional[List[str]], previous_v
         model_name = "Nano Banana PRO" if use_pro_model else "Nano Banana"
         other_model_name = "Nano Banana" if use_pro_model else "Nano Banana PRO"
         
+        # Get remaining credits for current model
+        remaining_current = get_remaining_generations(username, use_pro_model)
+        
+        if use_pro_model and remaining_current > 0 and remaining_current < credits_to_use:
+            gr.Info(f"You need {credits_to_use} credits for {get_resolution_value(resolution)} but only have {remaining_current} credits remaining. Try a lower resolution or use Nano Banana.")
+        
         if remaining_other > 0:
-            gr.Info(f"You've reached your daily limit for {model_name} ({limit_current} images). You still have {remaining_other} generations left with {other_model_name}!")
+            gr.Info(f"You've reached your daily limit for {model_name}. You still have {remaining_other} generations left with {other_model_name}!")
         
-        raise gr.Error(f"This demo is made for interactive generations, not automated workflows. You have generated {limit_current} images with {model_name} today, come back tomorrow for more.")
+        raise gr.Error(f"Insufficient credits. You need {credits_to_use} credits for this generation.")
     
     try:
         contents = [Image.open(image_path[0]) for image_path in images] if images else []
@@ -332,18 +341,38 @@ def unified_image_generator(prompt: str, images: Optional[List[str]], previous_v
         # Select model based on radio selection
         model_name = GEMINI_PRO_MODEL_NAME if use_pro_model else GEMINI_MODEL_NAME
         
-        # Create config with aspect ratio
-        if aspect_ratio == "Auto":
-            generate_content_config = types.GenerateContentConfig(
-                response_modalities=["IMAGE", "TEXT"],
-            )
+        # Create config with aspect ratio and resolution (for PRO model)
+        if use_pro_model:
+            # PRO model: use both aspect_ratio and image_size
+            resolution_value = get_resolution_value(resolution)
+            if aspect_ratio == "Auto":
+                generate_content_config = types.GenerateContentConfig(
+                    response_modalities=["IMAGE", "TEXT"],
+                    image_config=types.ImageConfig(
+                        image_size=resolution_value,
+                    ),
+                )
+            else:
+                generate_content_config = types.GenerateContentConfig(
+                    response_modalities=["IMAGE", "TEXT"],
+                    image_config=types.ImageConfig(
+                        aspect_ratio=aspect_ratio,
+                        image_size=resolution_value,
+                    ),
+                )
         else:
-            generate_content_config = types.GenerateContentConfig(
-                response_modalities=["IMAGE", "TEXT"],
-                image_config=types.ImageConfig(
-                    aspect_ratio=aspect_ratio,
-                ),
-            )
+            # Standard model: only aspect_ratio
+            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=model_name, 
@@ -428,25 +457,25 @@ with gr.Blocks(theme=gr.themes.Citrus(), css=css) as demo:
                 model_radio = gr.Radio(
                     choices=["Nano Banana", "Nano Banana PRO"],
                     value="Nano Banana",
-                    label="Model Selection",
-                    info="PRO: Higher quality, slower generation (50/day) | Standard: Faster generation (75/day)"
-                )
-                
-                # Resolution selection for PRO model
-                resolution_dropdown = gr.Dropdown(
-                    label="Resolution (Nano Banana PRO only)",
-                    choices=["1K (1 credit)", "2K (2 credits)", "4K (4 credits)"],
-                    value="1K (1 credit)",
-                    interactive=True,
-                    visible=False
                 )
                 
-                aspect_ratio_dropdown = gr.Dropdown(
-                    label="Aspect Ratio",
-                    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
-                )
+                with gr.Row():
+                    aspect_ratio_dropdown = gr.Dropdown(
+                        label="Aspect Ratio",
+                        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
+                    )
+                    
+                    resolution_dropdown = gr.Dropdown(
+                        label="Resolution (Nano Banana PRO only)",
+                        choices=["1K", "2K", "4K"],
+                        value="1K",
+                        interactive=True,
+                        visible=False
+                    )
+                    
+                    
                 
                 generate_button = gr.Button("Generate", variant="primary")
             with gr.Column(scale=1):
@@ -463,10 +492,20 @@ with gr.Blocks(theme=gr.themes.Citrus(), css=css) as demo:
 
     login_button = gr.LoginButton()
 
+    # Show/hide resolution dropdown based on model selection
+    def update_resolution_visibility(model_selection):
+        return gr.update(visible=(model_selection == "Nano Banana PRO"))
+    
+    model_radio.change(
+        fn=update_resolution_visibility,
+        inputs=[model_radio],
+        outputs=[resolution_dropdown]
+    )
+
     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, model_radio, manual_token],
+        inputs=[prompt_input, image_input_gallery, previous_video_state, last_frame_of_video_state, aspect_ratio_dropdown, model_radio, resolution_dropdown, manual_token],
         outputs=[output_image, create_video_button, extend_video_button, video_group],
         api_name=False
     )
@@ -513,4 +552,4 @@ with gr.Blocks(theme=gr.themes.Citrus(), css=css) as demo:
     demo.load(control_access, inputs=None, outputs=[main_interface, pro_message])
 
 if __name__ == "__main__":
-    demo.queue(max_size=None, default_concurrency_limit=None).launch(show_error=True)
+    demo.queue(max_size=None, default_concurrency_limit=None).launch(show_error=True)