fix: hf adjustments

app.py

@@ -5,73 +5,152 @@ import tempfile
 import os
 import requests
 import json
+import logging
+import traceback
+
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
 
 def inspect_model(model_id, config_file=None):
+    logger.info(f"Processing model ID: {model_id}")
+    
     if not model_id or '/' not in model_id:
         return "Please provide a valid model ID in the format username/modelname", "No config loaded."
     
     username, modelname = model_id.split('/', 1)
+    logger.info(f"Username: {username}, Model name: {modelname}")
     
-    model_url = f"https://huggingface.co/{username}/{modelname}/resolve/main/model.safetensors"
+    model_summary = "Processing..."
+    config_str = "No config loaded."
     
     try:
-        response = requests.get(model_url)
+        model_filename = "model.safetensors"
+        if "/" in modelname:
+            parts = modelname.split("/")
+            modelname = parts[0]
+            if len(parts) > 1 and parts[1].strip():
+                model_filename = parts[1]
+        
+        model_url = f"https://huggingface.co/{username}/{modelname}/resolve/main/{model_filename}"
+        logger.info(f"Attempting to download model from: {model_url}")
+        
+        response = requests.get(model_url, stream=True)
         response.raise_for_status()  
         
-
+        total_size = int(response.headers.get('content-length', 0))
+        logger.info(f"Model file size: {total_size/1024/1024:.2f} MB")
+        
         with tempfile.NamedTemporaryFile(delete=False, suffix=".safetensors") as tmp:
-            tmp.write(response.content)
+            if total_size > 0:
+                downloaded = 0
+                for chunk in response.iter_content(chunk_size=8192):
+                    if chunk:
+                        tmp.write(chunk)
+                        downloaded += len(chunk)
+                        if downloaded % (100 * 1024 * 1024) == 0: 
+                            logger.info(f"Downloaded {downloaded/1024/1024:.2f} MB / {total_size/1024/1024:.2f} MB")
+            else:
+                tmp.write(response.content)
+            
             model_path = tmp.name
+            logger.info(f"Model downloaded to temporary file: {model_path}")
         
-
+        logger.info("Loading model summary...")
         summary = load_model_summary(model_path)
+        logger.info(f"Loading state dictionary... (This may take time for large models)")
         state_dict = load_file(model_path)
+        logger.info("Analyzing top layers...")
         top_layers = get_top_layers(state_dict, summary["total_params"])
         top_layers_str = "\n".join([
             f"{layer['name']}: shape={layer['shape']}, params={layer['params']:,} ({layer['percent']}%)"
             for layer in top_layers
         ])
         
-
         config_data = {}
         if config_file is not None:
+            logger.info("Processing uploaded config file")
             with tempfile.NamedTemporaryFile(delete=False, suffix=".json") as tmp_cfg:
                 tmp_cfg.write(config_file.read())
-                config_data = load_config(tmp_cfg.name)
-                os.unlink(tmp_cfg.name)  
+                config_path = tmp_cfg.name
+                
+            logger.info(f"Loading config from uploaded file: {config_path}")
+            config_data = load_config(config_path)
+            os.unlink(config_path)  
         else:
             config_url = f"https://huggingface.co/{username}/{modelname}/resolve/main/config.json"
+            logger.info(f"Attempting to download config from: {config_url}")
             try:
                 config_response = requests.get(config_url)
                 config_response.raise_for_status()
                 config_data = json.loads(config_response.content)
+                logger.info("Config file downloaded and parsed successfully")
             except Exception as e:
-                pass  
+                logger.warning(f"Could not download or parse config file: {str(e)}")
         
         config_str = "\n".join([f"{k}: {v}" for k, v in config_data.items()]) if config_data else "No config loaded."
         
+        # Clean up temporary file
+        logger.info(f"Cleaning up temporary file: {model_path}")
         os.unlink(model_path)
         
-        return (
+        model_summary = (
             f" Total tensors: {summary['num_tensors']}\n"
             f" Total parameters: {summary['total_params']:,}\n\n"
-            f" Top Layers:\n{top_layers_str}",
-            config_str
+            f" Top Layers:\n{top_layers_str}"
         )
+        logger.info("Model inspection completed successfully")
+        
+        return model_summary, config_str
     
     except Exception as e:
-        return f"Error: {str(e)}", "No config loaded."
+        error_msg = f"Error: {str(e)}\n\nTraceback:\n{traceback.format_exc()}"
+        logger.error(error_msg)
+        return error_msg, "No config loaded."
+
+with gr.Blocks(title="Model Inspector") as demo:
+    gr.Markdown("# Model Inspector")
+    gr.Markdown("Enter a HuggingFace model ID in the format username/modelname to analyze its structure, parameter count, and configuration.")
+    gr.Markdown("You can specify a custom safetensors file by using username/modelname/filename.safetensors")
+    
+    with gr.Row():
+        with gr.Column():
+            model_id = gr.Textbox(
+                label="Model ID from HuggingFace", 
+                placeholder="username/modelname", 
+                lines=1
+            )
+            config_file = gr.File(
+                label="Upload config.json (optional)", 
+                type="binary"
+            )
+            submit_btn = gr.Button("Analyze Model", variant="primary")
+            status = gr.Markdown("Ready. Enter a model ID and click 'Analyze Model'")
+        
+        with gr.Column():
+            model_summary = gr.Textbox(label="Model Summary", lines=15)
+            config_output = gr.Textbox(label="Config", lines=10)
+    
+    def update_status(text):
+        return text
+    
+    def on_submit(model_id, config_file):
+        status_update = update_status("Processing... This may take some time for large models.")
+        yield status_update, None, None
+        
+        try:
+            summary, config = inspect_model(model_id, config_file)
+            status_update = update_status("Analysis complete!")
+            yield status_update, summary, config
+        except Exception as e:
+            error_msg = f"Error during analysis: {str(e)}"
+            status_update = update_status(f"❌ {error_msg}")
+            yield status_update, error_msg, "No config loaded."
+    
+    submit_btn.click(
+        fn=on_submit,
+        inputs=[model_id, config_file],
+        outputs=[status, model_summary, config_output],
+        show_progress=True
+    )
 
-gr.Interface(
-    fn=inspect_model,
-    inputs=[
-        gr.Textbox(label="Model ID from HuggingFace", placeholder="username/modelname", lines=1),
-        gr.File(label="Upload config.json (optional)", type="binary")
-    ],
-    outputs=[
-        gr.Textbox(label="Model Summary", lines=15),
-        gr.Textbox(label="Config", lines=10)
-    ],
-    title="Model Inspector",
-    description="Enter a HuggingFace model ID in the format username/modelname to analyze its structure, parameter count, and configuration."
-).launch()
+demo.launch()