Update handler.py

handler.py

@@ -24,8 +24,8 @@ class EndpointHandler():
         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, data) -> List[Dict[str, Any]]:
@@ -39,13 +39,19 @@ class EndpointHandler():
         inputs = data.get("inputs")
         logger.info('inputs received %s', inputs)
 
-        decoded_image = Image.open(BytesIO(base64.b64decode(inputs['image'])))
-        with autocast(device.type):
-            upscaled_image = self.pipe(prompt="", image = decoded_image).images[0]
+        image_base64 = base64.b64decode(inputs['image'])
+        logger.info('image_base64')
+        image_bytes = BytesIO(image_base64)
+        logger.info('image_bytes')
+        image = Image.open(image_bytes)
+        logger.info('image')
+        #with autocast(device.type):
+        #    upscaled_image = self.pipe(prompt="", image = decoded_image).images[0]
 
-        buffered = BytesIO()
-        upscaled_image.save(buffered, format="JPEG")
-        img_str = base64.b64encode(buffered.getvalue())
+        #buffered = BytesIO()
+        #upscaled_image.save(buffered, format="JPEG")
+        #img_str = base64.b64encode(buffered.getvalue())
 
         # postprocess the prediction
-        return {"image": img_str}
+        #return {"image": img_str}
+        return {"image": "test"}