Add HuggingFace login support for private/gated models

app.py

@@ -10,10 +10,45 @@ import sys
 
 import gradio as gr
 import torch
-from huggingface_hub import hf_hub_download, scan_cache_dir
+from huggingface_hub import hf_hub_download, scan_cache_dir, login, whoami
 from safetensors import safe_open
 
 
+def hf_login(token: str):
+    """Login to Hugging Face with provided token."""
+    if not token:
+        return "❌ Please provide a token", "Not logged in"
+    
+    try:
+        login(token=token, add_to_git_credential=False)
+        user_info = whoami()
+        username = user_info.get('name', 'Unknown')
+        return f"✅ Successfully logged in as: {username}", f"✅ Logged in as {username}"
+    except Exception as e:
+        return f"❌ Login failed: {str(e)}", "❌ Not logged in"
+
+
+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 check_hf_status():
+    """Check current HF login status."""
+    try:
+        user_info = whoami()
+        username = user_info.get('name', 'Unknown')
+        return f"✅ Currently logged in as: {username}", f"✅ Logged in as {username}"
+    except Exception:
+        return "ℹ️ Not logged in", "Not logged in"
+
+
 def get_param(model_id: str, param_key: str, log_buffer: io.StringIO, progress: gr.Progress):
     """
     Download and return a specific parameter tensor from a Hugging Face model.
@@ -345,6 +380,27 @@ custom_css = """
 
 with gr.Blocks(title="Hugging Face Model Weight Inspector") as demo:
     gr.Markdown("# 🔍 Hugging Face Model Weight Inspector")
+    
+    # HF Login section
+    with gr.Accordion("🔐 Hugging Face Login (for private/gated models)", open=False):
+        with gr.Row():
+            with gr.Column(scale=3):
+                hf_token_input = gr.Textbox(
+                    label="HF Token",
+                    placeholder="hf_...",
+                    type="password",
+                )
+            with gr.Column(scale=2):
+                hf_status = gr.Textbox(
+                    label="Status",
+                    value="Not logged in",
+                    interactive=False,
+                )
+        with gr.Row():
+            login_btn = gr.Button("🔑 Login", variant="primary", scale=1)
+            logout_btn = gr.Button("🚪 Logout", variant="secondary", scale=1)
+            check_status_btn = gr.Button("ℹ️ Check Status", variant="secondary", scale=1)
+        login_output = gr.Textbox(label="Login Status", interactive=False, lines=2)
 
     with gr.Row():
         with gr.Column(scale=1):
@@ -393,6 +449,24 @@ with gr.Blocks(title="Hugging Face Model Weight Inspector") as demo:
                 label="Status", interactive=False, lines=6)
 
     # Event handlers
+    login_btn.click(
+        fn=hf_login,
+        inputs=[hf_token_input],
+        outputs=[login_output, hf_status],
+    )
+    
+    logout_btn.click(
+        fn=hf_logout,
+        inputs=[],
+        outputs=[login_output, hf_status],
+    )
+    
+    check_status_btn.click(
+        fn=check_hf_status,
+        inputs=[],
+        outputs=[login_output, hf_status],
+    )
+    
     list_keys_btn.click(
         fn=list_keys,
         inputs=[model_id_input],