Upload inference.py with huggingface_hub

inference.py

@@ -5,24 +5,50 @@ Simple inference script for TFv6 NavSim model.
 import torch
 import numpy as np
 import cv2
-from stand_alone_model import load_tf
+from huggingface_hub import hf_hub_download
+
+# Lazy import to avoid errors if not downloaded yet
+try:
+    from stand_alone_model import load_tf
+except ImportError:
+    load_tf = None
 
 
 class TFv6NavSimInference:
     """Easy-to-use inference wrapper for TFv6 NavSim model."""
     
-    def __init__(self, model_path="model_0060.pth", device=None):
+    def __init__(self, model_path=None, device=None, repo_id="longpollehn/tfv6_navsim"):
         """
         Initialize the model.
         
         Args:
-            model_path: Path to the model checkpoint (.pth file)
+            model_path: Path to the model checkpoint (.pth file). If None, downloads from HF.
             device: torch.device or None (will auto-detect CUDA)
+            repo_id: Hugging Face repo to download from if model_path is None
         """
         if device is None:
             self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
         else:
             self.device = device
+        
+        # Auto-download from Hugging Face if no path provided
+        if model_path is None:
+            print(f"Downloading model from {repo_id}...")
+            model_path = hf_hub_download(repo_id=repo_id, filename="model_0060.pth")
+            print(f"Downloaded to {model_path}")
+        
+        # Import stand_alone_model here if not already imported
+        if load_tf is None:
+            import sys
+            import os
+            # Download stand_alone_model.py if needed
+            model_code_path = hf_hub_download(repo_id=repo_id, filename="stand_alone_model.py")
+            config_path = hf_hub_download(repo_id=repo_id, filename="config.json")
+            # Add to path and import
+            sys.path.insert(0, os.path.dirname(model_code_path))
+            from stand_alone_model import load_tf as _load_tf
+            global load_tf
+            load_tf = _load_tf
             
         print(f"Loading model on {self.device}...")
         self.model = load_tf(model_path, self.device)
@@ -120,9 +146,10 @@ class TFv6NavSimInference:
 
 
 def main():
-    """Example usage."""
-    # Initialize model
-    model = TFv6NavSimInference()
+    """Example usage - automatically downloads from HuggingFace."""
+    # Initialize model (will auto-download if not present)
+    print("Example: Using model directly from HuggingFace")
+    model = TFv6NavSimInference()  # No arguments needed!
     
     # Example: Create dummy input
     rgb = np.random.randint(0, 255, (900, 1600, 3), dtype=np.uint8)
@@ -133,7 +160,7 @@ def main():
     # Run prediction
     result = model.predict(rgb, command, speed, acceleration)
     
-    print("Prediction results:")
+    print("\nPrediction results:")
     if result['waypoints'] is not None:
         print(f"  Waypoints shape: {result['waypoints'].shape}")
         print(f"  First 3 waypoints:\n{result['waypoints'][:3]}")