Upload handler.py with huggingface_hub

handler.py

@@ -15,6 +15,30 @@ import numpy as np
 from PIL import Image
 import cv2
 
+# CRITICAL: Patch torch.autocast BEFORE any SAM3 imports
+# SAM3 uses @torch.autocast decorators that get applied at import time
+# We must patch torch.autocast before the decorators are evaluated
+class Float32Autocast:
+    """No-op autocast that forces float32."""
+    def __init__(self, device_type, dtype=None, enabled=True):
+        self.device_type = device_type
+        self.dtype = torch.float32
+        self.enabled = False
+
+    def __enter__(self):
+        return self
+
+    def __exit__(self, *args):
+        pass
+
+# Store original and replace globally
+_ORIGINAL_AUTOCAST = torch.autocast
+torch.autocast = Float32Autocast
+if hasattr(torch.cuda, 'amp'):
+    torch.cuda.amp.autocast = Float32Autocast
+if hasattr(torch, 'amp'):
+    torch.amp.autocast = Float32Autocast
+
 # Configure logging
 logging.basicConfig(
     level=logging.INFO,
@@ -24,7 +48,10 @@ logging.basicConfig(
 )
 logger = logging.getLogger(__name__)
 
+logger.info("✓ Patched torch.autocast globally before SAM3 import")
+
 # SAM3 imports - using local sam3 package in repository
+# This will now use our patched autocast for all @torch.autocast decorators
 from sam3.model_builder import build_sam3_video_predictor
 
 # HuggingFace Hub for uploads
@@ -67,39 +94,9 @@ class EndpointHandler:
             logger.info(f"GPU Device: {torch.cuda.get_device_name(0)}")
             logger.info(f"CUDA Version: {torch.version.cuda}")
             logger.info(f"Total GPU Memory: {torch.cuda.get_device_properties(0).total_memory / 1e9:.2f} GB")
-        
-        # CRITICAL FIX: Patch torch.autocast BEFORE building the predictor
-        # SAM3 has @torch.autocast decorators hardcoded to use BFloat16
-        # We need to override the autocast context manager to be a no-op
-        logger.info("Patching torch.autocast to disable BFloat16 (before model loading)...")
-
-        # Store the original autocast
-        self._original_autocast = torch.autocast
-
-        # Create a no-op autocast that always disables mixed precision
-        class Float32Autocast:
-            def __init__(self, device_type, dtype=None, enabled=True):
-                # Completely disable autocast
-                self.device_type = device_type
-                self.dtype = torch.float32
-                self.enabled = False
-
-            def __enter__(self):
-                return self
-
-            def __exit__(self, *args):
-                pass
-
-        # Monkey-patch torch.autocast globally BEFORE importing/building
-        torch.autocast = Float32Autocast
-        if hasattr(torch.cuda.amp, 'autocast'):
-            torch.cuda.amp.autocast = Float32Autocast
-        if hasattr(torch.amp, 'autocast'):
-            torch.amp.autocast = Float32Autocast
-
-        logger.info("✓ Patched torch.autocast to be a no-op (forces float32)")
 
         # Build SAM3 video predictor
+        # Note: torch.autocast was already patched at module import time
         try:
             logger.info("Building SAM3 video predictor...")
             start_time = time.time()