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 DiffusionPipeline
+from diffusers import StableDiffusionUpscalePipeline
 import base64
 from io import BytesIO
 from transformers.utils import logging
@@ -16,15 +16,15 @@ 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 = DiffusionPipeline.from_pretrained(model_id, torch_dtype=torch.float16)
+        #model_id = "stabilityai/stable-diffusion-x4-upscaler"
+        self.pipe = StableDiffusionUpscalePipeline.from_pretrained(path, torch_dtype=torch.float16)
         self.pipe = self.pipe.to(device)