Update handler.py

handler.py

@@ -2,7 +2,7 @@ from typing import  Dict, List, Any
 from PIL import Image
 import torch
 from torch import autocast
-from diffusers import StableDiffusionPipeline
+from diffusers import StableDiffusionUpscalePipeline
 import base64
 from io import BytesIO
 from transformers.utils import logging
@@ -24,7 +24,7 @@ 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 = StableDiffusionUpscalePipeline.from_pretrained(path, torch_dtype=torch.float16)
         self.pipe = self.pipe.to(device)
 
 
@@ -44,9 +44,10 @@ class EndpointHandler():
         image_bytes = BytesIO(image_base64)
         logger.info('image_bytes')
         image = Image.open(image_bytes)
+        prompt = inputs['prompt']
         logger.info('image')
         with autocast(device.type):
-            upscaled_image = self.pipe(prompt="", image = image).images[0]
+            upscaled_image = self.pipe(prompt, image).images[0]
 
         #buffered = BytesIO()
         #upscaled_image.save(buffered, format="JPEG")