Update app.py

app.py

@@ -19,7 +19,8 @@ DATA_DIR = Path("/data")
 DATA_DIR.mkdir(exist_ok=True)
 DB_PATH = DATA_DIR / "usage_limits.db"
 
-DAILY_LIMIT = 50
+DAILY_LIMIT_STANDARD = 75
+DAILY_LIMIT_PRO = 50
 EXEMPTED_USERS = ["multimodalart"]
 db_lock = Lock()
 
@@ -32,7 +33,8 @@ def init_db():
                 CREATE TABLE IF NOT EXISTS usage (
                     username TEXT PRIMARY KEY,
                     date TEXT NOT NULL,
-                    count INTEGER NOT NULL
+                    count_standard INTEGER NOT NULL DEFAULT 0,
+                    count_pro INTEGER NOT NULL DEFAULT 0
                 )
             ''')
             conn.commit()
@@ -42,7 +44,7 @@ def init_db():
         import traceback
         traceback.print_exc()
 
-def check_and_update_usage(username: str) -> bool:
+def check_and_update_usage(username: str, use_pro_model: bool) -> bool:
     """
     Check if user has reached daily limit and update usage.
     Returns True if user can generate, False if limit reached.
@@ -52,6 +54,10 @@ def check_and_update_usage(username: str) -> bool:
         print(f"User {username} is exempted from rate limits")
         return True
     
+    limit = DAILY_LIMIT_PRO if use_pro_model else DAILY_LIMIT_STANDARD
+    count_column = "count_pro" if use_pro_model else "count_standard"
+    model_name = "PRO" if use_pro_model else "Standard"
+    
     with db_lock:
         try:
             with sqlite3.connect(DB_PATH) as conn:
@@ -59,38 +65,47 @@ def check_and_update_usage(username: str) -> bool:
                 cursor = conn.cursor()
                 
                 # Get user record
-                cursor.execute("SELECT date, count FROM usage WHERE username = ?", (username,))
+                cursor.execute("SELECT date, count_standard, count_pro FROM usage WHERE username = ?", (username,))
                 result = cursor.fetchone()
                 
                 if result is None:
                     # New user - create record
-                    cursor.execute("INSERT INTO usage (username, date, count) VALUES (?, ?, ?)", 
-                                   (username, today, 1))
+                    if use_pro_model:
+                        cursor.execute("INSERT INTO usage (username, date, count_standard, count_pro) VALUES (?, ?, ?, ?)", 
+                                       (username, today, 0, 1))
+                    else:
+                        cursor.execute("INSERT INTO usage (username, date, count_standard, count_pro) VALUES (?, ?, ?, ?)", 
+                                       (username, today, 1, 0))
                     conn.commit()
-                    print(f"New user {username}: 1/{DAILY_LIMIT}")
+                    print(f"New user {username}: 1/{limit} ({model_name})")
                     return True
                 
-                user_date, user_count = result
+                user_date, user_count_standard, user_count_pro = result
+                user_count = user_count_pro if use_pro_model else user_count_standard
                 
                 # Reset if new day
                 if user_date != today:
-                    cursor.execute("UPDATE usage SET date = ?, count = ? WHERE username = ?", 
-                                   (today, 1, username))
+                    if use_pro_model:
+                        cursor.execute("UPDATE usage SET date = ?, count_standard = ?, count_pro = ? WHERE username = ?", 
+                                       (today, 0, 1, username))
+                    else:
+                        cursor.execute("UPDATE usage SET date = ?, count_standard = ?, count_pro = ? WHERE username = ?", 
+                                       (today, 1, 0, username))
                     conn.commit()
-                    print(f"User {username} reset for new day: 1/{DAILY_LIMIT}")
+                    print(f"User {username} reset for new day: 1/{limit} ({model_name})")
                     return True
                 
                 # Check limit
-                if user_count >= DAILY_LIMIT:
-                    print(f"User {username} reached limit: {user_count}/{DAILY_LIMIT}")
+                if user_count >= limit:
+                    print(f"User {username} reached limit: {user_count}/{limit} ({model_name})")
                     return False
                 
                 # Increment count
                 new_count = user_count + 1
-                cursor.execute("UPDATE usage SET count = ? WHERE username = ?", 
+                cursor.execute(f"UPDATE usage SET {count_column} = ? WHERE username = ?", 
                                (new_count, username))
                 conn.commit()
-                print(f"User {username} usage: {new_count}/{DAILY_LIMIT}")
+                print(f"User {username} usage: {new_count}/{limit} ({model_name})")
                 return True
                 
         except Exception as e:
@@ -100,34 +115,37 @@ def check_and_update_usage(username: str) -> bool:
             # On error, allow the request (fail open)
             return True
 
-def get_remaining_generations(username: str) -> int:
+def get_remaining_generations(username: str, use_pro_model: bool) -> int:
     """Get the number of remaining generations for today."""
     # Exempted users have unlimited generations
     if username in EXEMPTED_USERS:
         return 999999  # Return a large number to indicate unlimited
     
+    limit = DAILY_LIMIT_PRO if use_pro_model else DAILY_LIMIT_STANDARD
+    
     with db_lock:
         try:
             with sqlite3.connect(DB_PATH) as conn:
                 today = str(date.today())
                 cursor = conn.cursor()
                 
-                cursor.execute("SELECT date, count FROM usage WHERE username = ?", (username,))
+                cursor.execute("SELECT date, count_standard, count_pro FROM usage WHERE username = ?", (username,))
                 result = cursor.fetchone()
                 
                 if result is None:
-                    return DAILY_LIMIT
+                    return limit
                 
-                user_date, user_count = result
+                user_date, user_count_standard, user_count_pro = result
+                user_count = user_count_pro if use_pro_model else user_count_standard
                 
                 # Reset if new day
                 if user_date != today:
-                    return DAILY_LIMIT
+                    return limit
                 
-                return max(0, DAILY_LIMIT - user_count)
+                return max(0, limit - user_count)
         except Exception as e:
             print(f"Error getting remaining generations for {username}: {e}")
-            return DAILY_LIMIT
+            return limit
 
 # Initialize database on module load
 init_db()
@@ -262,23 +280,37 @@ 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, use_pro_model: bool, 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, 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")
+    
     # Check rate limit
     username = get_username(oauth_token) or get_username(manual_token)
     if not username:
         raise gr.Error("Could not identify user.")
     
-    if not check_and_update_usage(username):
-        raise gr.Error(f"This demo is made for interactive generations, not automated workflows. You have generated {DAILY_LIMIT} images today, come back tomorrow for more.")
+    can_generate = check_and_update_usage(username, use_pro_model)
+    
+    if not can_generate:
+        # Check if user has quota on the other model
+        remaining_other = get_remaining_generations(username, not use_pro_model)
+        limit_current = DAILY_LIMIT_PRO if use_pro_model else DAILY_LIMIT_STANDARD
+        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"
+        
+        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}!")
+        
+        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.")
     
     try:
         contents = [Image.open(image_path[0]) for image_path in images] if images else []
         contents.append(prompt)
         
-        # Select model based on checkbox
+        # 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 (omit image_config if Auto is selected)
@@ -384,22 +416,18 @@ with gr.Blocks(theme=gr.themes.Citrus(), css=css) as demo:
                 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
+                model_radio = gr.Radio(
+                    choices=["Nano Banana", "Nano Banana PRO"],
+                    value="Nano Banana PRO",
                 )
                 
                 aspect_ratio_dropdown = gr.Dropdown(
-                    label="Aspect Ratio (only for standard model)",
+                    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
                 )
                 
-                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")
@@ -418,7 +446,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, use_pro_checkbox, manual_token],
+        inputs=[prompt_input, image_input_gallery, previous_video_state, last_frame_of_video_state, aspect_ratio_dropdown, model_radio, manual_token],
         outputs=[output_image, create_video_button, extend_video_button, video_group],
         api_name=False
     )