Update app.py

app.py

@@ -44,10 +44,11 @@ def init_db():
         import traceback
         traceback.print_exc()
 
-def check_and_update_usage(username: str, use_pro_model: bool) -> bool:
+def check_and_update_usage(username: str, use_pro_model: bool, credits_to_use: int = 1) -> bool:
     """
     Check if user has reached daily limit and update usage.
     Returns True if user can generate, False if limit reached.
+    credits_to_use: Number of credits to consume (1 for standard/1K, 2 for 2K, 4 for 4K)
     """
     # Exempted users bypass all checks
     if username in EXEMPTED_USERS:
@@ -72,12 +73,12 @@ def check_and_update_usage(username: str, use_pro_model: bool) -> bool:
                     # New user - create record
                     if use_pro_model:
                         cursor.execute("INSERT INTO usage (username, date, count_standard, count_pro) VALUES (?, ?, ?, ?)", 
-                                       (username, today, 0, 1))
+                                       (username, today, 0, credits_to_use))
                     else:
                         cursor.execute("INSERT INTO usage (username, date, count_standard, count_pro) VALUES (?, ?, ?, ?)", 
-                                       (username, today, 1, 0))
+                                       (username, today, credits_to_use, 0))
                     conn.commit()
-                    print(f"New user {username}: 1/{limit} ({model_name})")
+                    print(f"New user {username}: {credits_to_use}/{limit} ({model_name})")
                     return True
                 
                 user_date, user_count_standard, user_count_pro = result
@@ -87,25 +88,25 @@ def check_and_update_usage(username: str, use_pro_model: bool) -> bool:
                 if user_date != today:
                     if use_pro_model:
                         cursor.execute("UPDATE usage SET date = ?, count_standard = ?, count_pro = ? WHERE username = ?", 
-                                       (today, 0, 1, username))
+                                       (today, 0, credits_to_use, username))
                     else:
                         cursor.execute("UPDATE usage SET date = ?, count_standard = ?, count_pro = ? WHERE username = ?", 
-                                       (today, 1, 0, username))
+                                       (today, credits_to_use, 0, username))
                     conn.commit()
-                    print(f"User {username} reset for new day: 1/{limit} ({model_name})")
+                    print(f"User {username} reset for new day: {credits_to_use}/{limit} ({model_name})")
                     return True
                 
-                # Check limit
-                if user_count >= limit:
-                    print(f"User {username} reached limit: {user_count}/{limit} ({model_name})")
+                # Check if user has enough credits remaining
+                if user_count + credits_to_use > limit:
+                    print(f"User {username} insufficient credits: needs {credits_to_use}, has {limit - user_count}/{limit} remaining ({model_name})")
                     return False
                 
-                # Increment count
-                new_count = user_count + 1
+                # Increment count by credits used
+                new_count = user_count + credits_to_use
                 cursor.execute(f"UPDATE usage SET {count_column} = ? WHERE username = ?", 
                                (new_count, username))
                 conn.commit()
-                print(f"User {username} usage: {new_count}/{limit} ({model_name})")
+                print(f"User {username} usage: {new_count}/{limit} (used {credits_to_use} credits) ({model_name})")
                 return True
                 
         except Exception as e:
@@ -201,6 +202,24 @@ def get_username(token: Optional[Union[gr.OAuthToken, str]]) -> Optional[str]:
         print(f"Could not get username: {e}")
         return None
 
+def get_credit_cost(resolution: str) -> int:
+    """Get the credit cost for a given resolution."""
+    if "4K" in resolution:
+        return 4
+    elif "2K" in resolution:
+        return 2
+    else:  # 1K
+        return 1
+
+def get_resolution_value(resolution: str) -> str:
+    """Extract the resolution value from the dropdown selection."""
+    if "4K" in resolution:
+        return "4K"
+    elif "2K" in resolution:
+        return "2K"
+    else:
+        return "1K"
+
 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:
@@ -413,6 +432,15 @@ with gr.Blocks(theme=gr.themes.Citrus(), css=css) as demo:
                     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"],
@@ -485,4 +513,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)