Implement per-user session-based authentication

app.py

@@ -10,58 +10,55 @@ import sys
 
 import gradio as gr
 import torch
-from huggingface_hub import hf_hub_download, scan_cache_dir, login, whoami
+from huggingface_hub import hf_hub_download, scan_cache_dir, whoami
 from safetensors import safe_open
 
-# Auto-login with HF_TOKEN if available (for HuggingFace Spaces)
-HF_TOKEN = os.environ.get("HF_TOKEN")
-if HF_TOKEN:
-    try:
-        login(token=HF_TOKEN, add_to_git_credential=False)
-        print("✅ Automatically logged in using HF_TOKEN")
-    except Exception as e:
-        print(f"⚠️ Auto-login failed: {str(e)}")
+# Default token from HF_TOKEN environment variable (for HuggingFace Spaces)
+DEFAULT_HF_TOKEN = os.environ.get("HF_TOKEN")
 
 
-def hf_login(token: str):
-    """Login to Hugging Face with provided token."""
+def hf_login(token: str, session_token: str):
+    """Login to Hugging Face with provided token (per-user session)."""
     if not token:
-        return "❌ Please provide a token", "Not logged in"
+        return "❌ Please provide a token", "Not logged in", session_token
     
     try:
-        login(token=token, add_to_git_credential=False)
-        user_info = whoami()
+        user_info = whoami(token=token)
         username = user_info.get('name', 'Unknown')
-        return f"✅ Successfully logged in as: {username}", f"✅ Logged in as {username}"
+        return f"✅ Successfully logged in as: {username}", f"✅ Logged in as {username}", token
     except Exception as e:
-        return f"❌ Login failed: {str(e)}", "❌ Not logged in"
+        return f"❌ Login failed: {str(e)}", "❌ Not logged in", session_token
 
 
-def hf_logout():
-    """Logout from Hugging Face."""
-    try:
-        # Clear token by logging in with empty token
-        from huggingface_hub import logout
-        logout()
-        return "✅ Successfully logged out", "Not logged in"
-    except Exception as e:
-        return f"❌ Logout failed: {str(e)}", "Status unknown"
+def hf_logout(session_token: str):
+    """Logout from Hugging Face (clear session token)."""
+    return "✅ Successfully logged out", "Not logged in", None
 
 
-def check_hf_status():
-    """Check current HF login status."""
+def check_hf_status(session_token: str):
+    """Check current HF login status for this session."""
+    # Check session token first, then fall back to default token
+    token = session_token or DEFAULT_HF_TOKEN
+    
+    if not token:
+        return "ℹ️ Not logged in", "Not logged in", session_token
+    
     try:
-        user_info = whoami()
+        user_info = whoami(token=token)
         username = user_info.get('name', 'Unknown')
-        return f"✅ Currently logged in as: {username}", f"✅ Logged in as {username}"
+        source = "(session)" if session_token else "(default HF_TOKEN)"
+        return f"✅ Currently logged in as: {username} {source}", f"✅ Logged in as {username}", session_token
     except Exception:
-        return "ℹ️ Not logged in", "Not logged in"
+        return "ℹ️ Not logged in", "Not logged in", session_token
 
 
-def get_param(model_id: str, param_key: str, log_buffer: io.StringIO, progress: gr.Progress):
+def get_param(model_id: str, param_key: str, log_buffer: io.StringIO, progress: gr.Progress, token: str = None):
     """
     Download and return a specific parameter tensor from a Hugging Face model.
     """
+    # Use session token or fall back to default token
+    auth_token = token or DEFAULT_HF_TOKEN
+    
     # Redirect stderr to log buffer for real-time tqdm updates
     original_stderr = sys.stderr
     sys.stderr = log_buffer
@@ -73,7 +70,7 @@ def get_param(model_id: str, param_key: str, log_buffer: io.StringIO, progress:
             progress(0.1, desc="Downloading index...")
 
             index_path = hf_hub_download(
-                model_id, "model.safetensors.index.json")
+                model_id, "model.safetensors.index.json", token=auth_token)
 
             log_buffer.write(f"✓ Index file found: {index_path}\n")
 
@@ -96,7 +93,7 @@ def get_param(model_id: str, param_key: str, log_buffer: io.StringIO, progress:
         log_buffer.write(f"📥 Downloading shard: {shard_file}...\n")
         progress(0.3, desc=f"Downloading {shard_file}...")
 
-        shard_path = hf_hub_download(model_id, shard_file)
+        shard_path = hf_hub_download(model_id, shard_file, token=auth_token)
 
         log_buffer.write(f"\n✓ Shard downloaded: {shard_path}\n")
         progress(0.7, desc="Loading tensor...")
@@ -113,17 +110,20 @@ def get_param(model_id: str, param_key: str, log_buffer: io.StringIO, progress:
         sys.stderr = original_stderr
 
 
-def get_available_keys(model_id: str):
+def get_available_keys(model_id: str, token: str = None):
     """Get all available parameter keys from a model."""
+    # Use session token or fall back to default token
+    auth_token = token or DEFAULT_HF_TOKEN
+    
     try:
-        index_path = hf_hub_download(model_id, "model.safetensors.index.json")
+        index_path = hf_hub_download(model_id, "model.safetensors.index.json", token=auth_token)
         with open(index_path, "r", encoding="utf-8") as f:
             index = json.load(f)
         return sorted(index["weight_map"].keys())
     except Exception:
         # Try single file
         try:
-            shard_path = hf_hub_download(model_id, "model.safetensors")
+            shard_path = hf_hub_download(model_id, "model.safetensors", token=auth_token)
             with safe_open(shard_path, framework="pt") as f:
                 return sorted(f.keys())
         except Exception as e:
@@ -156,7 +156,7 @@ def format_tensor_info(tensor: torch.Tensor) -> str:
     return "<br>".join(info)
 
 
-def fetch_param(model_id: str, param_key: str, progress=gr.Progress()):
+def fetch_param(model_id: str, param_key: str, session_token: str, progress=gr.Progress()):
     """Fetch parameter and return formatted info and tensor preview."""
     log_buffer = io.StringIO()
     last_log_value = ""
@@ -179,7 +179,7 @@ def fetch_param(model_id: str, param_key: str, progress=gr.Progress()):
         
         def download_thread():
             try:
-                result_tensor[0] = get_param(model_id, param_key, log_buffer, progress)
+                result_tensor[0] = get_param(model_id, param_key, log_buffer, progress, session_token)
             except Exception as e:
                 download_error[0] = e
             finally:
@@ -260,13 +260,13 @@ def fetch_param(model_id: str, param_key: str, progress=gr.Progress()):
         yield f"**Error:** {str(e)}", "", None, log_buffer.getvalue()
 
 
-def list_keys(model_id: str):
+def list_keys(model_id: str, session_token: str):
     """List all available keys for a model."""
     if not model_id:
         return "Please provide a model ID."
 
     try:
-        keys = get_available_keys(model_id)
+        keys = get_available_keys(model_id, session_token)
         if not keys:
             return "No keys found or failed to load model."
         return "\n".join(keys)
@@ -390,11 +390,14 @@ custom_css = """
 with gr.Blocks(title="Hugging Face Model Weight Inspector") as demo:
     gr.Markdown("# 🔍 Hugging Face Model Weight Inspector")
     
+    # Session state for per-user token
+    session_token = gr.State(None)
+    
     # HF Login section
-    with gr.Accordion("🔐 Hugging Face Login (Optional - auto-login with HF_TOKEN)", open=False):
+    with gr.Accordion("🔐 Hugging Face Login (Per-User Session)", open=False):
         gr.Markdown("""
-        **Note:** This Space automatically uses the `HF_TOKEN` secret if configured.  
-        Manual login below is only needed if you want to use a different token.
+        **Note:** This Space uses the default `HF_TOKEN` secret for all users if no session token is provided.  
+        Login below with your own token for per-user authentication (affects only your session).
         """)
         with gr.Row():
             with gr.Column(scale=3):
@@ -404,9 +407,10 @@ with gr.Blocks(title="Hugging Face Model Weight Inspector") as demo:
                     type="password",
                 )
             with gr.Column(scale=2):
+                initial_status = "✅ Using default HF_TOKEN" if DEFAULT_HF_TOKEN else "Not logged in"
                 hf_status = gr.Textbox(
                     label="Status",
-                    value="Not logged in",
+                    value=initial_status,
                     interactive=False,
                 )
         with gr.Row():
@@ -464,31 +468,31 @@ with gr.Blocks(title="Hugging Face Model Weight Inspector") as demo:
     # Event handlers
     login_btn.click(
         fn=hf_login,
-        inputs=[hf_token_input],
-        outputs=[login_output, hf_status],
+        inputs=[hf_token_input, session_token],
+        outputs=[login_output, hf_status, session_token],
     )
     
     logout_btn.click(
         fn=hf_logout,
-        inputs=[],
-        outputs=[login_output, hf_status],
+        inputs=[session_token],
+        outputs=[login_output, hf_status, session_token],
     )
     
     check_status_btn.click(
         fn=check_hf_status,
-        inputs=[],
-        outputs=[login_output, hf_status],
+        inputs=[session_token],
+        outputs=[login_output, hf_status, session_token],
     )
     
     list_keys_btn.click(
         fn=list_keys,
-        inputs=[model_id_input],
+        inputs=[model_id_input, session_token],
         outputs=[keys_output],
     )
 
     fetch_btn.click(
         fn=fetch_param,
-        inputs=[model_id_input, param_key_input],
+        inputs=[model_id_input, param_key_input, session_token],
         outputs=[info_output, preview_output, download_output, log_output],
     )
 
@@ -509,6 +513,13 @@ with gr.Blocks(title="Hugging Face Model Weight Inspector") as demo:
         inputs=[],
         outputs=[clear_status],
     )
+    
+    # Auto-check status on load
+    demo.load(
+        fn=check_hf_status,
+        inputs=[session_token],
+        outputs=[login_output, hf_status, session_token],
+    )
 
 
 if __name__ == "__main__":