add authentication

app.py

@@ -20,6 +20,16 @@ from library.utils_prompt import *
 authorized_users_str = os.environ.get("AUTHORIZED_USER_IDS", "")
 AUTHORIZED_USER_IDS = set(authorized_users_str.split(',') if authorized_users_str and authorized_users_str.strip() else [])
 
+# Check if current user is authorized
+def check_authorization(profile: gr.OAuthProfile | None):
+    if profile is None:
+        return False #"Not logged in. Please Login for full access."
+    
+    if profile.username in AUTHORIZED_USER_IDS:
+        return True
+    else:
+        return False
+
 #Api sdk
 OR = OpenRouterAPI()
 
@@ -39,7 +49,8 @@ MODEL_PRICING = {
     "google/gemini-2.5-flash-preview:thinking": "$0.35",
     "gpt-4.1-nano": "$0.02",
     "openai/chatgpt-4o-latest": "$0.75",
-    "meta-llama/llama-4-maverick": "$0.04"
+    "meta-llama/llama-4-maverick": "$0.04",
+    "meta-llama/llama-4-maverick:free": "Free"
 }
 
 def get_sys_prompt(length="medium", photograph=False):
@@ -120,6 +131,23 @@ def create_demo():
     """
     # --- Pass css to gr.Blocks ---
     with gr.Blocks(theme=gr.themes.Monochrome(), css=custom_css) as demo:
+        # Store authentication state
+        is_authenticated = gr.State(False)
+        
+        # Initialize authentication on app load
+        def app_load():
+            # Get the OAuth profile (will be None if not logged in)
+            profile = getattr(gr.context.get('user', None), 'profile', None)
+            # Use the check_authorization function
+            return check_authorization(profile)
+        
+        # Update authentication state on load
+        demo.load(
+            fn=app_load,
+            inputs=None,
+            outputs=[is_authenticated]
+        )
+        
         with gr.Row():
             with gr.Column(scale=3):
                 gr.Markdown("# MATCHA: Museum Alt-Text for Cultural Heritage with AI 🍵 🌿")
@@ -139,7 +167,11 @@ def create_demo():
         # Store model choices and state
         show_all_models_state = gr.State(False)
 
-        # Define preferred and additional models directly in the function
+        # Define model lists
+        free_model = [
+            ("Llama 4 Maverick (Free)", "meta-llama/llama-4-maverick:free")
+        ]
+        
         preferred_models = [
             ("Gemini 2.0 Flash (cheap)", "google/gemini-2.0-flash-001"),
             ("GPT-4.1 Mini", "gpt-4.1-mini"),
@@ -155,11 +187,15 @@ def create_demo():
             ("Llama 4 Maverick", "meta-llama/llama-4-maverick")
         ]
         
-        # Calculate all models once
+        # Calculate all models
         all_models_list = preferred_models + additional_models
         
-        # Default model value
-        default_model = "google/gemini-2.0-flash-001"
+        # Function to update model dropdown based on auth status
+        def update_model_dropdown(is_auth):
+            if is_auth:
+                return gr.Dropdown(choices=preferred_models, value="google/gemini-2.0-flash-001")
+            else:
+                return gr.Dropdown(choices=free_model, value="meta-llama/llama-4-maverick:free")
         
         with gr.Row():
             # Left column: Controls and uploads
@@ -170,11 +206,25 @@ def create_demo():
                     file_count="multiple"
                 )
                 
-                # Model dropdown
+                # Model dropdown - initially empty, will be populated based on auth status
                 model_choice = gr.Dropdown(
-                    choices=preferred_models, 
+                    choices=[], 
                     label="Select Model",
-                    value=default_model
+                    value=None
+                )
+                
+                # Update the model dropdown when authentication state changes
+                is_authenticated.change(
+                    fn=update_model_dropdown,
+                    inputs=[is_authenticated],
+                    outputs=[model_choice]
+                )
+                
+                # Also update model dropdown on initial load
+                demo.load(
+                    fn=update_model_dropdown,
+                    inputs=[is_authenticated],
+                    outputs=[model_choice]
                 )
                 
                 length_choice = gr.Radio(
@@ -192,22 +242,37 @@ def create_demo():
                         info="Display additional model options in the dropdown above"
                     )
                     
+                    # Update checkbox interactivity based on auth status
+                    def update_checkbox_interactive(is_auth):
+                        return gr.Checkbox(
+                            label="Show Additional Models",
+                            value=False,
+                            info="Display additional model options in the dropdown above",
+                            interactive=is_auth
+                        )
+                    
+                    is_authenticated.change(
+                        fn=update_checkbox_interactive,
+                        inputs=[is_authenticated],
+                        outputs=[show_all_models]
+                    )
+                    
+                    # Also update checkbox interactivity on initial load
+                    demo.load(
+                        fn=update_checkbox_interactive,
+                        inputs=[is_authenticated],
+                        outputs=[show_all_models]
+                    )
+                    
                     content_type = gr.Radio(
                         choices=["Museum Object", "Photography"], 
                         label="Content Type",
                         value="Museum Object"
                     )
                 
-                # Find the default model's display name
-                default_model_name = "Unknown Model"
-                for name, value in preferred_models:
-                    if value == default_model:
-                        default_model_name = name
-                        break
-                
+                # Model info display
                 model_info = gr.Markdown(
-                    f"""**Current Model**: {default_model_name}  
-                    **Estimated cost per 100 Images**: {MODEL_PRICING[default_model]}""",
+                    "**Select a model to see information**",
                     elem_id="model-info-display"
                 )
                 
@@ -248,31 +313,34 @@ def create_demo():
                 all_images = gr.State([])
                 all_results = gr.State([])
         
-        # Handle checkbox change to update model dropdown - modern version
-        def toggle_models(show_all, current_model):
-            # Make a fresh copy of the models lists to avoid any reference issues
+        # Handle checkbox change to update model dropdown
+        def toggle_models(show_all, current_model, is_auth):
+            if not is_auth:
+                # Non-authenticated users only get free model
+                return gr.Dropdown(choices=free_model, value="meta-llama/llama-4-maverick:free")
+            
+            # For authenticated users, existing logic
             preferred_choices = list(preferred_models)
             all_choices = list(all_models_list)
             
             if show_all:
-                # When showing all models, use the fresh copy of all models
                 return gr.Dropdown(choices=all_choices, value=current_model)
             else:
-                # Check if current model is in preferred models list
                 preferred_values = [value for _, value in preferred_choices]
                 
                 if current_model in preferred_values:
-                    # Keep the current model if it's in preferred models
                     return gr.Dropdown(choices=preferred_choices, value=current_model)
                 else:
-                    # Reset to default model if current model is not in preferred models
-                    return gr.Dropdown(choices=preferred_choices, value=default_model)
+                    return gr.Dropdown(choices=preferred_choices, value="google/gemini-2.0-flash-001")
 
         # Update model info when model selection changes
         def update_model_info(model_value):
+            # Combine all model lists for lookup
+            combined_models = preferred_models + additional_models + free_model
+            
             # Find display name
             model_name = "Unknown Model"
-            for name, value in all_models_list:
+            for name, value in combined_models:
                 if value == model_value:
                     model_name = name
                     break
@@ -287,7 +355,7 @@ def create_demo():
         # Connect checkbox to toggle model choices
         show_all_models.change(
             fn=toggle_models,
-            inputs=[show_all_models, model_choice],
+            inputs=[show_all_models, model_choice, is_authenticated],
             outputs=[model_choice]
         )
         
@@ -410,37 +478,8 @@ def create_demo():
             """)
 
     return demo
-
-# New function to check if current user is authorized
-def check_authorization(profile: gr.OAuthProfile | None):
-    if profile is None:
-        return False #"Not logged in. Please Login for full access."
-    
-    if profile.username in AUTHORIZED_USER_IDS:
-        return True
-    else:
-        return False
-
-def hello(profile: gr.OAuthProfile | None) -> str:
-    if profile is None:
-        return "I don't know you."
-    return f"Hello {profile.name}"
-    
-def create_demo_free():
-    with gr.Blocks() as demo:
-        gr.LoginButton()
-        gr.Markdown("Please Login")
-        m2 = gr.Markdown()
-        m3 = gr.Markdown(f"Authorized User IDs: {', '.join(AUTHORIZED_USER_IDS)}")
-        demo.load(hello, inputs=None, outputs=m2)
-        m2 = gr.Markdown(str(check_authorization(None)))
-    return demo
     
 # Launch the app
 if __name__ == "__main__":
-    profile = gr.OAuthProfile()
-    if check_authorization(profile):
-        app = create_demo()
-    else:
-        app = create_demo_free()
+    app = create_demo()
     app.launch()