Update app.py

app.py

@@ -49,6 +49,7 @@ class Preprocessor:
         elif name == "NormalBae":
             print("Loading NormalBae")
             device = "cuda" if torch.cuda.is_available() else "cpu"
+            # The .to(device) call should add the .device attribute
             self.model = NormalBaeDetector.from_pretrained(self.MODEL_ID).to(device)
             if torch.cuda.is_available():
                 torch.cuda.empty_cache()
@@ -57,11 +58,35 @@ class Preprocessor:
             raise ValueError
         return
 
+    # Modified __call__ method to safely check device
     def __call__(self, image: Image.Image, **kwargs) -> Image.Image:
         device = "cuda" if torch.cuda.is_available() else "cpu"
-        if self.model.device.type != device:
-             self.model.to(device)
-        return self.model(image, **kwargs)
+        
+        # Check if the model has the 'device' attribute before accessing its 'type'
+        if hasattr(self.model, 'device'):
+            if self.model.device.type != device:
+                 print(f"Moving preprocessor model to {device}") # Debug print
+                 # Attempt to move the model if it's not on the correct device
+                 try:
+                     self.model.to(device)
+                 except Exception as e:
+                     print(f"Error moving preprocessor model to {device}: {e}")
+                     # If moving fails, the next call might fail with a device mismatch
+                     pass # Continue and let the next line potentially raise an error
+        else:
+            # If .device attribute is missing, assume it's not on the correct device and try to move it
+            # This might happen if NormalBaeDetector is not a standard torch.nn.Module
+            print("Warning: Preprocessor model has no .device attribute. Attempting to move to correct device.")
+            try:
+                self.model.to(device) # Attempt to move even if .device is missing
+            except Exception as e:
+                print(f"Error attempting to move preprocessor model without .device attribute: {e}")
+                # If moving fails, the next call will likely fail with a device mismatch
+                pass # Continue and let the next line potentially raise an error
+
+
+        return self.model(image, **kwargs) # This call might fail if the model isn't on the correct device
+
 
 # Load models and preprocessor when the script starts
 # Controlnet Normal