Update handler.py

handler.py

@@ -16,15 +16,16 @@ logger.warning("WARN")
 # set device
 device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
 
-if device.type != 'cuda':
-    raise ValueError("need to run on GPU")
+#if device.type != 'cuda':
+    #raise ValueError("need to run on GPU")
 
 class EndpointHandler():
     def __init__(self, path=""):
+        self.path = path
         # load the optimized model
         #model_id = "stabilityai/stable-diffusion-x4-upscaler"
-        self.pipe = StableDiffusionPipeline.from_pretrained(path, torch_dtype=torch.float16)
-        self.pipe = self.pipe.to(device)
+        #self.pipe = StableDiffusionPipeline.from_pretrained(path, torch_dtype=torch.float16)
+        #self.pipe = self.pipe.to(device)
 
 
     def __call__(self, *args, **kwargs) -> List[Dict[str, Any]]:
@@ -34,8 +35,8 @@ class EndpointHandler():
         Return:
             A :obj:`dict`:. base64 encoded image
         """
-        print(args)
-        print(kwargs)
+        logger.info('args received %s', args)
+        logger.info('kwargs received %s', kwargs)
 
         # postprocess the prediction
         return {"image": "nothing"}