Update handler.py

handler.py

@@ -7,8 +7,11 @@ class InferenceHandler:
     def __init__(self):
         self.device = "cuda" if torch.cuda.is_available() else "cpu"
         model_name = "colt12/maxcushion"
-        self.pipe = StableDiffusionXLPipeline.from_pretrained(model_name, torch_dtype=torch.float16, use_safetensors=True, variant="fp16")
-        self.pipe = self.pipe.to(self.device)
+        self.pipe = StableDiffusionXLPipeline.from_pretrained(
+            model_name,
+            torch_dtype=torch.float16,
+            use_safetensors=True
+        ).to(self.device)  # Moved to the device here
 
     def __call__(self, inputs):
         prompt = inputs.get("prompt", "")