Add GPU mode

handler.py

@@ -7,9 +7,6 @@ from PIL import Image
 from diffusers import StableDiffusionPipeline
 
 
-REPO_ID = "runwayml/stable-diffusion-v1-5"
-
-
 # helper decoder
 def decode_base64_image(image_string):
     base64_image = base64.b64decode(image_string)
@@ -20,7 +17,7 @@ def decode_base64_image(image_string):
 class EndpointHandler:
     def __init__(self, path=""):
         self.pipe = StableDiffusionPipeline.from_pretrained("./stable-diffusion-v1-5")
-        # self.pipe = self.pipe.to("cuda")
+        self.pipe = self.pipe.to("cuda")
 
     def __call__(self, data: Any) -> List[List[Dict[str, float]]]:
         """