cuda fix

instruction_classifier.py

@@ -76,8 +76,11 @@ class InstructionClassifierSanitizer:
         self.model_repo_id = model_repo_id
         self.model_filename = model_filename
         
-        # Initialize device
-        self.device = get_device()
+        # Initialize device - always use CPU for initialization in ZeroGPU environments
+        # GPU operations will be handled within @spaces.GPU decorated methods
+        self.device = torch.device('cpu')
+        self.target_device = get_device()  # Store target device for later use
+        print(f"🔧 Device configuration: init_device={self.device}, target_device={self.target_device}")
         
         # Map friendly names to actual model names
         model_mapping = {
@@ -103,9 +106,10 @@ class InstructionClassifierSanitizer:
                 model_path = "models/best_instruction_classifier.pth"
             
             if os.path.exists(model_path):
-                checkpoint = torch.load(model_path, map_location=self.device)
+                checkpoint = torch.load(model_path, map_location='cpu')
                 self._load_model_weights(checkpoint)
                 print(f"✅ Loaded instruction classifier model from {model_path}")
+                print(f"   Model loaded on {self.device} for ZeroGPU compatibility")
             else:
                 raise FileNotFoundError(f"Model file not found: {model_path}")
         else:
@@ -141,14 +145,15 @@ class InstructionClassifierSanitizer:
                 file_size = os.path.getsize(model_path) / (1024**3)  # GB
                 print(f"   File size: {file_size:.2f} GB")
                 
-                # Load the checkpoint from the downloaded file
+                # Load the checkpoint from the downloaded file - always use CPU for ZeroGPU compatibility
                 print("🔄 Loading checkpoint into memory...")
-                checkpoint = torch.load(model_path, map_location=self.device)
+                checkpoint = torch.load(model_path, map_location='cpu')
                 print(f"   Checkpoint keys: {len(checkpoint.keys())}")
                 
                 self._load_model_weights(checkpoint)
                 print(f"✅ Model weights loaded from {self.model_repo_id}")
                 print(f"   Model parameter count: {sum(p.numel() for p in self.model.parameters())}")
+                print(f"   Model loaded on {self.device} for ZeroGPU compatibility")
                 
             except Exception as e:
                 print(f"❌ CRITICAL ERROR: Failed to download model from {self.model_repo_id}")
@@ -171,9 +176,9 @@ class InstructionClassifierSanitizer:
             if not key.startswith('loss_fct'):  # Skip loss function weights
                 model_state_dict[key] = value
         
-        # Load the filtered state dict
+        # Load the filtered state dict - keep on CPU for ZeroGPU compatibility
         self.model.load_state_dict(model_state_dict, strict=False)
-        self.model.to(self.device)
+        self.model.to(self.device)  # Keep on CPU during initialization
         self.model.eval()
     
     @spaces.GPU
@@ -190,6 +195,12 @@ class InstructionClassifierSanitizer:
         if not tool_output or not tool_output.strip():
             return tool_output
         
+        # Move model to target device (GPU) within @spaces.GPU decorated method
+        if self.device != self.target_device:
+            print(f"🚀 Moving model from {self.device} to {self.target_device} within @spaces.GPU context")
+            self.model.to(self.target_device)
+            self.device = self.target_device
+        
         try:
             # Step 1: Detect if the tool output contains instructions
             is_injection, confidence_score, tagged_text = self._detect_injection(tool_output)
@@ -219,6 +230,7 @@ class InstructionClassifierSanitizer:
             # Return original output if sanitization fails
             return tool_output
     
+    @spaces.GPU
     def sanitize_with_annotations(self, tool_output: str) -> Tuple[str, List[Dict[str, any]]]:
         """
         Sanitization function that also returns annotation data for flagged content.
@@ -233,6 +245,12 @@ class InstructionClassifierSanitizer:
         if not tool_output or not tool_output.strip():
             return tool_output, []
         
+        # Move model to target device (GPU) within @spaces.GPU decorated method
+        if self.device != self.target_device:
+            print(f"🚀 Moving model from {self.device} to {self.target_device} within @spaces.GPU context")
+            self.model.to(self.target_device)
+            self.device = self.target_device
+        
         try:
             # Step 1: Detect if the tool output contains instructions
             is_injection, confidence_score, tagged_text = self._detect_injection(tool_output)