Fix weights downloader

README.md

@@ -4,7 +4,7 @@ emoji: 🔍
 colorFrom: blue
 colorTo: purple
 sdk: gradio
-sdk_version: 4.0.0
+sdk_version: 6.0.2
 app_file: app.py
 pinned: false
 license: mit

app.py

@@ -204,21 +204,7 @@ def create_interface():
     """Create and configure the Gradio interface."""
     
     with gr.Blocks(
-        title="Forensic Image Analysis Agent",
-        theme=gr.themes.Soft(),
-        css="""
-        .gradio-container {
-            font-family: 'Segoe UI', Tahoma, Geneva, Verdana, sans-serif;
-        }
-        .main-header {
-            text-align: center;
-            padding: 20px;
-            background: linear-gradient(135deg, #667eea 0%, #764ba2 100%);
-            color: white;
-            border-radius: 10px;
-            margin-bottom: 20px;
-        }
-        """
+        title="Forensic Image Analysis Agent"
     ) as demo:
         gr.HTML("""
         <div class="main-header">
@@ -337,6 +323,20 @@ if __name__ == "__main__":
     demo.launch(
         server_name="0.0.0.0",
         server_port=7860,
-        share=False
+        share=False,
+        theme=gr.themes.Soft(),
+        css="""
+        .gradio-container {
+            font-family: 'Segoe UI', Tahoma, Geneva, Verdana, sans-serif;
+        }
+        .main-header {
+            text-align: center;
+            padding: 20px;
+            background: linear-gradient(135deg, #667eea 0%, #764ba2 100%);
+            color: white;
+            border-radius: 10px;
+            margin-bottom: 20px;
+        }
+        """
     )
 

app_requirements.txt

@@ -1,5 +1,5 @@
 # Gradio for UI
-gradio>=4.0.0
+gradio>=6.0.2
 
 # Core LLM and agent dependencies
 langchain>=0.1.0

src/utils/weight_downloader.py

@@ -4,6 +4,7 @@ Utility to download and verify TruFor model weights automatically.
 
 import hashlib
 import os
+import shutil
 import zipfile
 from pathlib import Path
 from typing import Optional, Tuple
@@ -64,6 +65,13 @@ def ensure_trufor_weights(workspace_root: Optional[Path] = None, auto_download:
     """
     Ensure TruFor weights are available, downloading if necessary.
     
+    Downloads TruFor_weights.zip from the official source, verifies MD5,
+    and extracts trufor.pth.tar to weights/trufor/trufor.pth.tar.
+    
+    Zip structure: weights/trufor.pth.tar
+    Final path: projectroot/weights/trufor/trufor.pth.tar
+    MD5 is verified on the zip file (not the tar).
+    
     Args:
         workspace_root: Root directory of the workspace. If None, tries to detect it.
         auto_download: If True, automatically download weights if missing.
@@ -82,19 +90,14 @@ def ensure_trufor_weights(workspace_root: Optional[Path] = None, auto_download:
     
     # Check if weights already exist
     if weights_path.exists():
-        # Verify MD5 if possible
-        try:
-            md5_hash = _calculate_md5(weights_path)
-            if md5_hash.lower() == TRUFOR_WEIGHTS_MD5.lower():
-                return True, f"✅ TruFor weights found and verified at {weights_path}"
-            else:
-                print(f"⚠️  MD5 mismatch. Expected: {TRUFOR_WEIGHTS_MD5}, Got: {md5_hash}")
-                print("   Weights file exists but MD5 doesn't match. Consider re-downloading.")
-                # Still return True - file exists, just not verified
-                return True, f"⚠️  TruFor weights found at {weights_path} but MD5 verification failed"
-        except Exception as e:
-            # If MD5 check fails, still return True if file exists
-            return True, f"✅ TruFor weights found at {weights_path} (MD5 check skipped: {e})"
+        # File exists - we can't verify MD5 since it's for the zip, not the tar
+        file_size = weights_path.stat().st_size
+        if file_size > 0:
+            return True, f"✅ TruFor weights found at {weights_path} ({file_size / 1024 / 1024:.1f} MB)"
+        else:
+            # Empty file - delete and re-download
+            weights_path.unlink()
+            print("⚠️  Found empty weights file, re-downloading...")
     
     # Weights don't exist
     if not auto_download:
@@ -118,73 +121,81 @@ def ensure_trufor_weights(workspace_root: Optional[Path] = None, auto_download:
         if not _download_file(TRUFOR_WEIGHTS_URL, zip_path):
             return False, f"❌ Failed to download weights from {TRUFOR_WEIGHTS_URL}"
         
+        # Verify MD5 of the zip file immediately after download
+        print(f"🔐 Verifying download integrity (MD5)...")
+        try:
+            zip_md5 = _calculate_md5(zip_path)
+            if zip_md5.lower() != TRUFOR_WEIGHTS_MD5.lower():
+                zip_path.unlink()
+                return False, (
+                    f"❌ Downloaded zip MD5 mismatch!\n"
+                    f"   Expected: {TRUFOR_WEIGHTS_MD5}\n"
+                    f"   Got: {zip_md5}\n"
+                    f"   The download may be corrupted. Please try again."
+                )
+            print(f"✅ MD5 verified: {zip_md5}")
+        except Exception as e:
+            print(f"⚠️  MD5 verification failed: {e}. Continuing with extraction...")
+        
         # Extract zip file
+        # Zip structure: weights/trufor.pth.tar
         print(f"📦 Extracting weights...")
         with zipfile.ZipFile(zip_path, 'r') as zip_ref:
             # Find the weights file in the zip
             members = zip_ref.namelist()
             weights_member = None
             
-            # Look for trufor.pth.tar in the zip
+            # Look for trufor.pth.tar in the zip (expected: weights/trufor.pth.tar)
             for member in members:
-                if member.endswith(TRUFOR_WEIGHTS_FILENAME) or member.endswith(f"/{TRUFOR_WEIGHTS_FILENAME}"):
+                if member.endswith(TRUFOR_WEIGHTS_FILENAME):
                     weights_member = member
                     break
             
-            if weights_member:
-                # Extract just the weights file
-                zip_ref.extract(weights_member, weights_dir)
-                
-                # Move to final location if needed
-                extracted_path = weights_dir / weights_member
-                if extracted_path != weights_path:
-                    if weights_path.exists():
-                        weights_path.unlink()
-                    extracted_path.rename(weights_path)
-                
-                # Clean up zip file
+            if not weights_member:
                 zip_path.unlink()
+                return False, (
+                    f"❌ Could not find {TRUFOR_WEIGHTS_FILENAME} in downloaded zip.\n"
+                    f"   Zip contents: {members}"
+                )
+            
+            # Extract just the weights file to a temp location
+            # zip_ref.extract will create the nested directory structure
+            zip_ref.extract(weights_member, weights_dir)
+            
+            # Move from extracted location to final location
+            # e.g., weights/trufor/weights/trufor.pth.tar -> weights/trufor/trufor.pth.tar
+            extracted_path = weights_dir / weights_member
+            
+            if extracted_path != weights_path:
+                # Move to final location
+                if weights_path.exists():
+                    weights_path.unlink()
+                shutil.move(str(extracted_path), str(weights_path))
                 
-                # Verify MD5
+                # Clean up any empty directories left from extraction
                 try:
-                    md5_hash = _calculate_md5(weights_path)
-                    if md5_hash.lower() == TRUFOR_WEIGHTS_MD5.lower():
-                        return True, f"✅ TruFor weights downloaded and verified at {weights_path}"
-                    else:
-                        return False, (
-                            f"❌ Downloaded weights MD5 mismatch!\n"
-                            f"   Expected: {TRUFOR_WEIGHTS_MD5}\n"
-                            f"   Got: {md5_hash}\n"
-                            f"   File at: {weights_path}"
-                        )
-                except Exception as e:
-                    return True, f"✅ TruFor weights downloaded at {weights_path} (MD5 check failed: {e})"
+                    # Remove the 'weights' folder if it was created inside weights_dir
+                    extracted_parent = extracted_path.parent
+                    while extracted_parent != weights_dir and extracted_parent.exists():
+                        if not any(extracted_parent.iterdir()):
+                            extracted_parent.rmdir()
+                        extracted_parent = extracted_parent.parent
+                except Exception:
+                    pass  # Ignore cleanup errors
+        
+        # Clean up zip file
+        zip_path.unlink()
+        
+        # Verify final file exists and has content
+        if weights_path.exists():
+            file_size = weights_path.stat().st_size
+            if file_size > 0:
+                return True, f"✅ TruFor weights downloaded successfully to {weights_path} ({file_size / 1024 / 1024:.1f} MB)"
             else:
-                # Extract all files and look for weights
-                zip_ref.extractall(weights_dir)
-                zip_path.unlink()
-                
-                # Look for the weights file in extracted files
-                for root, dirs, files in os.walk(weights_dir):
-                    for file in files:
-                        if file == TRUFOR_WEIGHTS_FILENAME:
-                            found_path = Path(root) / file
-                            if found_path != weights_path:
-                                if weights_path.exists():
-                                    weights_path.unlink()
-                                found_path.rename(weights_path)
-                            
-                            # Verify MD5
-                            try:
-                                md5_hash = _calculate_md5(weights_path)
-                                if md5_hash.lower() == TRUFOR_WEIGHTS_MD5.lower():
-                                    return True, f"✅ TruFor weights downloaded and verified at {weights_path}"
-                                else:
-                                    return False, f"❌ Downloaded weights MD5 mismatch: {md5_hash}"
-                            except Exception as e:
-                                return True, f"✅ TruFor weights downloaded at {weights_path} (MD5 check failed: {e})"
-                
-                return False, f"❌ Could not find {TRUFOR_WEIGHTS_FILENAME} in downloaded zip"
+                weights_path.unlink()
+                return False, f"❌ Extracted weights file is empty"
+        else:
+            return False, f"❌ Failed to extract weights to {weights_path}"
         
     except Exception as e:
         # Clean up on error