gpu for cuda

handler.py

@@ -20,9 +20,7 @@ device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
 if device.type != 'cuda':
     raise ValueError("need to run on GPU")
 # set mixed precision dtype
-#dtype = torch.bfloat16 if torch.cuda.get_device_capability()[0] == 8 else torch.float16
-
-dtype = torch.float32
+dtype = torch.bfloat16 if torch.cuda.get_device_capability()[0] == 8 else torch.float16
     
 # controlnet mapping for controlnet id and control hinter
 CONTROLNET_MAPPING = {
@@ -87,7 +85,8 @@ class EndpointHandler():
 
                                                                       
         # Define Generator with seed
-        self.generator = torch.Generator(device="cpu").manual_seed(3)
+#        self.generator = torch.Generator(device="cpu").manual_seed(3)
+        self.generator = torch.Generator(device="cuda").manual_seed(3)
 
     def __call__(self, data: Any) -> List[List[Dict[str, float]]]:
         """