fix

.gitattributes

@@ -1,39 +1,37 @@
-*.7z filter=lfs diff=lfs merge=lfs -text
-*.arrow filter=lfs diff=lfs merge=lfs -text
-*.bin filter=lfs diff=lfs merge=lfs -text
-*.bz2 filter=lfs diff=lfs merge=lfs -text
-*.ckpt filter=lfs diff=lfs merge=lfs -text
-*.ftz filter=lfs diff=lfs merge=lfs -text
-*.gz filter=lfs diff=lfs merge=lfs -text
-*.h5 filter=lfs diff=lfs merge=lfs -text
-*.joblib filter=lfs diff=lfs merge=lfs -text
-*.lfs.* filter=lfs diff=lfs merge=lfs -text
-*.mlmodel filter=lfs diff=lfs merge=lfs -text
-*.model filter=lfs diff=lfs merge=lfs -text
-*.msgpack filter=lfs diff=lfs merge=lfs -text
-*.npy filter=lfs diff=lfs merge=lfs -text
-*.npz filter=lfs diff=lfs merge=lfs -text
-*.onnx filter=lfs diff=lfs merge=lfs -text
-*.ot filter=lfs diff=lfs merge=lfs -text
-*.parquet filter=lfs diff=lfs merge=lfs -text
-*.pb filter=lfs diff=lfs merge=lfs -text
-*.pickle filter=lfs diff=lfs merge=lfs -text
-*.pkl filter=lfs diff=lfs merge=lfs -text
-*.pt filter=lfs diff=lfs merge=lfs -text
-*.pth filter=lfs diff=lfs merge=lfs -text
-*.rar filter=lfs diff=lfs merge=lfs -text
-*.safetensors filter=lfs diff=lfs merge=lfs -text
-saved_model/**/* filter=lfs diff=lfs merge=lfs -text
-*.tar.* filter=lfs diff=lfs merge=lfs -text
-*.tar filter=lfs diff=lfs merge=lfs -text
-*.tflite filter=lfs diff=lfs merge=lfs -text
-*.tgz filter=lfs diff=lfs merge=lfs -text
-*.wasm filter=lfs diff=lfs merge=lfs -text
-*.xz filter=lfs diff=lfs merge=lfs -text
-*.zip filter=lfs diff=lfs merge=lfs -text
-*.zst filter=lfs diff=lfs merge=lfs -text
-*tfevents* filter=lfs diff=lfs merge=lfs -text
-checkpoints/dora.pt filter=lfs diff=lfs merge=lfs -text
-checkpoints/mtlora.pt filter=lfs diff=lfs merge=lfs -text
-utils/res10_300x300_ssd_iter_140000_fp16.caffemodel filter=lfs diff=lfs merge=lfs -text
-core/vision_encoder/bpe_simple_vocab_16e6.txt.gz filter=lfs diff=lfs merge=lfs -text
+*.7z filter=lfs diff=lfs merge=lfs -text
+*.arrow filter=lfs diff=lfs merge=lfs -text
+*.bin filter=lfs diff=lfs merge=lfs -text
+*.bz2 filter=lfs diff=lfs merge=lfs -text
+*.ckpt filter=lfs diff=lfs merge=lfs -text
+*.ftz filter=lfs diff=lfs merge=lfs -text
+*.gz filter=lfs diff=lfs merge=lfs -text
+*.h5 filter=lfs diff=lfs merge=lfs -text
+*.joblib filter=lfs diff=lfs merge=lfs -text
+*.lfs.* filter=lfs diff=lfs merge=lfs -text
+*.mlmodel filter=lfs diff=lfs merge=lfs -text
+*.model filter=lfs diff=lfs merge=lfs -text
+*.msgpack filter=lfs diff=lfs merge=lfs -text
+*.npy filter=lfs diff=lfs merge=lfs -text
+*.npz filter=lfs diff=lfs merge=lfs -text
+*.onnx filter=lfs diff=lfs merge=lfs -text
+*.ot filter=lfs diff=lfs merge=lfs -text
+*.parquet filter=lfs diff=lfs merge=lfs -text
+*.pb filter=lfs diff=lfs merge=lfs -text
+*.pickle filter=lfs diff=lfs merge=lfs -text
+*.pkl filter=lfs diff=lfs merge=lfs -text
+*.pt filter=lfs diff=lfs merge=lfs -text
+*.pth filter=lfs diff=lfs merge=lfs -text
+*.rar filter=lfs diff=lfs merge=lfs -text
+*.safetensors filter=lfs diff=lfs merge=lfs -text
+saved_model/**/* filter=lfs diff=lfs merge=lfs -text
+*.tar.* filter=lfs diff=lfs merge=lfs -text
+*.tar filter=lfs diff=lfs merge=lfs -text
+*.tflite filter=lfs diff=lfs merge=lfs -text
+*.tgz filter=lfs diff=lfs merge=lfs -text
+*.wasm filter=lfs diff=lfs merge=lfs -text
+*.xz filter=lfs diff=lfs merge=lfs -text
+*.zip filter=lfs diff=lfs merge=lfs -text
+*.zst filter=lfs diff=lfs merge=lfs -text
+*tfevents* filter=lfs diff=lfs merge=lfs -text
+utils/res10_300x300_ssd_iter_140000_fp16.caffemodel filter=lfs diff=lfs merge=lfs -text
+core/vision_encoder/bpe_simple_vocab_16e6.txt.gz filter=lfs diff=lfs merge=lfs -text

app.py

@@ -11,7 +11,7 @@ from PIL import Image, ImageDraw, ImageFont
 import base64
 from io import BytesIO
 import traceback  # Import traceback at the top
-
+from huggingface_hub import snapshot_download # Use snapshot_download for startup
 from utils.face_detector import FaceDetector
 
 # Class definitions
@@ -36,8 +36,9 @@ model = None
 transform = None
 detector = None
 device = None
-current_ckpt_dir = None 
+current_ckpt_dir = None # This will now store the full path to the loaded model
 CHECKPOINTS_DIR = './checkpoints/'
+MODEL_REPO_ID = "Antuke/FaR-FT-PE"
 
 def scan_checkpoints(ckpt_dir):
     """Scans a directory for .pt or .pth files."""
@@ -49,6 +50,7 @@ def scan_checkpoints(ckpt_dir):
         ckpt_files = [
             os.path.join(ckpt_dir, f)
             for f in sorted(os.listdir(ckpt_dir))
+            if f.endswith(('.pt', '.pth')) # Ensure we only scan for model files
         ]
     except Exception as e:
         print(f"Error scanning checkpoint directory {ckpt_dir}: {e}")
@@ -78,33 +80,41 @@ def load_model(device,ckpt_dir='./checkpoints/mtlora.pt', pe_vision_config="PE-C
     model.load_model(filepath=ckpt_dir,map_location=device)
     return model,transform
 
-def load_model_and_update_status(ckpt_dir):
+def load_model_and_update_status(model_filepath):
     """
-    Wrapper function to load a model and return a status message.
-    This is used by the dropdown's 'change' event.
+    Wrapper function to load a model *from a local file path* and return a status.
+    The file path is provided by the dropdown.
     """
     global model, current_ckpt_dir
     
-    if ckpt_dir is None or ckpt_dir == "":
+    if model_filepath is None or model_filepath == "":
         return "No checkpoint selected."
 
-    if model is not None and ckpt_dir == current_ckpt_dir:
-        status = f"Model already loaded: {os.path.basename(ckpt_dir)}"
+    # Check if this model *filepath* is already loaded
+    if model is not None and model_filepath == current_ckpt_dir:
+        status = f"Model already loaded: {os.path.basename(model_filepath)}"
         print(status)
         return status
         
-    gr.Info(f"Loading model: {os.path.basename(ckpt_dir)}...")
+    gr.Info(f"Loading model: {os.path.basename(model_filepath)}...")
     try:
-        init_model(ckpt_dir=ckpt_dir, detection_confidence=0.5)
-        current_ckpt_dir = ckpt_dir  # Set global directory on successful load
-        status = f"Successfully loaded: {os.path.basename(ckpt_dir)}"
+        # --- This is the new logic ---
+        # The file is already local. Just initialize it.
+        # The 'model_filepath' *is* the ckpt_dir for init_model
+        init_model(ckpt_dir=model_filepath, detection_confidence=0.5)
+        # --- End of new logic ---
+        
+        current_ckpt_dir = model_filepath  # Set global path on successful load
+        status = f"Successfully loaded: {os.path.basename(model_filepath)}"
         gr.Info("Model loaded successfully!")
         print(status)
         return status
+
     except Exception as e:
-        status = f"Failed to load {os.path.basename(ckpt_dir)}: {str(e)}"
-        print(status)
         traceback.print_exc()
+        status = f"Failed to load {os.path.basename(model_filepath)}: {e}"
+        gr.Info(f"Error: {status}")
+        print(f"ERROR: {status}")
         return status
 
 def predict(model, image):
@@ -182,10 +192,10 @@ def init_model(ckpt_dir="./checkpoints/mtlora.pt", detection_confidence=0.5):
 
     # Verify model weights exist
     if not os.path.exists(ckpt_dir):
-       error_msg = f"Model weights not found: {ckpt_dir}."
-       print(f"ERROR: {error_msg}")
-       raise FileNotFoundError(error_msg)
-       
+        error_msg = f"Model weights not found: {ckpt_dir}."
+        print(f"ERROR: {error_msg}")
+        raise FileNotFoundError(error_msg)
+        
     print(f"Model weights found: {ckpt_dir}")
 
     # Load the perception encoder
@@ -214,7 +224,11 @@ def process_image(image, selected_checkpoint_path):
         return None, "<p style='color: red;'>Please upload an image</p>"
 
     # Ensure model is initialized
+    # This check is crucial. If the user changes the dropdown, it triggers
+    # load_model_and_update_status. If they just hit "Classify",
+    # this check ensures the selected model is loaded.
     if model is None or selected_checkpoint_path != current_ckpt_dir:
+        print(f"Model mismatch or not loaded. Selected: {selected_checkpoint_path}, Current: {current_ckpt_dir}")
         status = load_model_and_update_status(selected_checkpoint_path)
         if "Failed" in status or "Error" in status:
             return image, f"<p style'color: red;'>Model Error: {status}</p>"
@@ -621,8 +635,8 @@ def create_interface(checkpoint_list, default_checkpoint, initial_status):
             with gr.Column(scale=3):
                 checkpoint_dropdown = gr.Dropdown(
                     label="Select Model Checkpoint",
-                    choices=checkpoint_list,
-                    value=default_checkpoint,
+                    choices=checkpoint_list, # This is now a list of (label, path) tuples
+                    value=default_checkpoint, # This is the full path to the default
                 )
             with gr.Column(scale=2):
                 model_status_text = gr.Textbox(
@@ -710,13 +724,13 @@ def create_interface(checkpoint_list, default_checkpoint, initial_status):
         # Event handlers
         analyze_btn.click(
             fn=process_image,
-            inputs=[input_image, checkpoint_dropdown], # Pass dropdown value
+            inputs=[input_image, checkpoint_dropdown], # Pass dropdown value (which is the path)
             outputs=[output_image, output_html]
         )
         
         checkpoint_dropdown.change(
             fn=load_model_and_update_status,
-            inputs=[checkpoint_dropdown],
+            inputs=[checkpoint_dropdown], # Pass dropdown value (which is the path)
             outputs=[model_status_text]
         )
 
@@ -725,12 +739,29 @@ def create_interface(checkpoint_list, default_checkpoint, initial_status):
 
 # === Main Application Startup ===
 
-# Initialize for Hugging Face Spaces (module-level)
 print("="*60)
 print("VLM SOFT BIOMETRICS - GRADIO INTERFACE")
 print("="*60)
 
+# --- NEW: Download models BEFORE scanning ---
+print(f"Downloading model weights from {MODEL_REPO_ID} to {CHECKPOINTS_DIR}...")
+os.makedirs(CHECKPOINTS_DIR, exist_ok=True)
+try:
+    snapshot_download(
+        repo_id=MODEL_REPO_ID,
+        local_dir=CHECKPOINTS_DIR,
+        allow_patterns=["*.pt", "*.pth"], # Grabs all weight files
+        local_dir_use_symlinks=False,
+    )
+    print("Model download complete.")
+except Exception as e:
+    print(f"CRITICAL: Failed to download models from Hub. {e}")
+    traceback.print_exc()
+# --- End of NEW code ---
+
+
 # --- 1. Scan for models first ---
+# This will now find the files you just downloaded
 checkpoint_list, default_checkpoint = scan_checkpoints(CHECKPOINTS_DIR)
 
 if not checkpoint_list:
@@ -744,6 +775,7 @@ initial_status_msg = "No default model found. Please select one."
 if default_checkpoint:
     print(f"\nInitializing default model: {default_checkpoint}")
     # This will load the model AND set current_ckpt_dir
+    # It now correctly uses the local file path
     initial_status_msg = load_model_and_update_status(default_checkpoint)
     print(initial_status_msg)
 else:
@@ -760,8 +792,9 @@ if __name__ == "__main__":
     import argparse
 
     parser = argparse.ArgumentParser(description="VLM Soft Biometrics - Gradio Interface")
+    # ckpt_dir is now managed by the startup download, so this arg is less relevant
     parser.add_argument("--ckpt_dir", type=str, default="./checkpoints/", 
-                        help="Path to the checkpoint directory (overridden by UI)")
+                        help="Path to the checkpoint directory (will be populated from HF Hub)")
     parser.add_argument("--detection_confidence", type=float, default=0.5,
                         help="Confidence threshold for face detection")
     parser.add_argument("--port", type=int, default=7860,
@@ -772,9 +805,9 @@ if __name__ == "__main__":
                         help="Server name/IP to bind to")
     args = parser.parse_args()
     
-    # Update global config if args are provided (though UI dropdown is primary)
+    # Update global config if args are provided
     CHECKPOINTS_DIR = args.ckpt_dir
-    # Note: detection_confidence is passed to init_model, so it's handled.
+    # Note: detection_confidence is passed to init_model during load, so it's handled.
 
     print(f"\nLaunching server on {args.server_name}:{args.port}")
     print(f"Monitoring checkpoint directory: {CHECKPOINTS_DIR}")

requirements.txt

@@ -8,4 +8,5 @@ peft
 python-dotenv
 tqdm
 gradio
-timm
+timm
+huggingface-hub