Update app.py

app.py

@@ -1,10 +1,10 @@
 import gradio as gr
 from huggingface_hub import InferenceClient
 import os
-import json
 import base64
 from PIL import Image
 import io
+from typing import Optional
 
 ACCESS_TOKEN = os.getenv("HF_TOKEN")
 print("Access token loaded.")
@@ -39,6 +39,15 @@ def encode_image(image_path):
         print(f"Error encoding image: {e}")
         return None
 
+
+def render_profile(profile: Optional[gr.OAuthProfile]) -> str:
+    """Display the current authentication status in the UI."""
+    if profile is None:
+        return "Not signed in. Use the button above to sign in with Hugging Face and run inference with your own account."
+
+    display_name = getattr(profile, "name", None) or getattr(profile, "username", "Hugging Face user")
+    return f"Signed in as **{display_name}**."
+
 def respond(
     message,
     image_files,  # Changed parameter name and structure
@@ -51,7 +60,8 @@ def respond(
     seed,
     custom_model,    
     model_search_term,
-    selected_model
+    selected_model,
+    oauth_token: Optional[gr.OAuthToken] = None
 ):
     print(f"Received message: {message}")
     print(f"Received {len(image_files) if image_files else 0} images")
@@ -63,9 +73,22 @@ def respond(
     print(f"Model search term: {model_search_term}")
     print(f"Selected model from radio: {selected_model}")
 
+    api_token: Optional[str] = None
+
+    if oauth_token is not None and getattr(oauth_token, "token", None):
+        api_token = oauth_token.token
+        print("Using OAuth token from signed-in user for inference.")
+    elif ACCESS_TOKEN:
+        api_token = ACCESS_TOKEN
+        print("Using server-configured Hugging Face token for inference.")
+    else:
+        raise gr.Error(
+            "No Hugging Face session detected. Please sign in with your Hugging Face account before running the chat."
+        )
+
     # Initialize the Inference Client with default HF inference
-    client = InferenceClient(token=ACCESS_TOKEN)
-    print(f"Hugging Face Inference Client initialized with standard HF inference.")
+    client = InferenceClient(token=api_token)
+    print("Hugging Face Inference Client initialized with available token.")
 
     # Convert seed to None if -1 (meaning random)
     if seed == -1:
@@ -201,6 +224,10 @@ def respond(
 
 # GRADIO UI
 with gr.Blocks(theme="Nymbo/Nymbo_Theme") as demo:
+    with gr.Row(elem_id="oauth-row"):
+        login_button = gr.LoginButton()
+        auth_status = gr.Markdown(render_profile(None), elem_id="oauth-status")
+
     # Create the chatbot component
     chatbot = gr.Chatbot(
         height=600, 
@@ -386,7 +413,19 @@ with gr.Blocks(theme="Nymbo/Nymbo_Theme") as demo:
             return history
     
     # Define bot response function
-    def bot(history, system_msg, max_tokens, temperature, top_p, freq_penalty, seed, custom_model, search_term, selected_model):
+    def bot(
+        history,
+        system_msg,
+        max_tokens,
+        temperature,
+        top_p,
+        freq_penalty,
+        seed,
+        custom_model,
+        search_term,
+        selected_model,
+        oauth_token: Optional[gr.OAuthToken] = None,
+    ):
         # Check if history is valid
         if not history or len(history) == 0:
             print("No history to process")
@@ -435,7 +474,8 @@ with gr.Blocks(theme="Nymbo/Nymbo_Theme") as demo:
                 seed,
                 custom_model,
                 search_term,
-                selected_model
+                selected_model,
+                oauth_token
             ):
                 history[-1][1] = response
                 yield history
@@ -453,7 +493,8 @@ with gr.Blocks(theme="Nymbo/Nymbo_Theme") as demo:
                 seed,
                 custom_model,
                 search_term,
-                selected_model
+                selected_model,
+                oauth_token
             ):
                 history[-1][1] = response
                 yield history
@@ -492,6 +533,8 @@ with gr.Blocks(theme="Nymbo/Nymbo_Theme") as demo:
     )
     print("Featured model radio button change event linked.")
 
+    demo.load(render_profile, inputs=None, outputs=auth_status)
+
 print("Gradio interface initialized.")
 
 if __name__ == "__main__":