Update handler.py

handler.py

@@ -15,8 +15,8 @@ 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=""):
@@ -35,7 +35,6 @@ class EndpointHandler():
 
         self.pipe = StableDiffusionUpscalePipeline.from_pretrained(self.path, torch_dtype=torch.float16)
         self.pipe = self.pipe.to(device)
-        self.pipe.unet.in_channels = 8
         logger.info('data received %s', data)
         inputs = data.get("inputs")
         logger.info('inputs received %s', inputs)
@@ -44,7 +43,7 @@ class EndpointHandler():
         logger.info('image_base64')
         image_bytes = BytesIO(image_base64)
         logger.info('image_bytes')
-        image = Image.open(image_bytes)
+        image = Image.open(image_bytes).convert("RGB")
         prompt = inputs['prompt']
         logger.info('image')