cuda

app.py

@@ -69,8 +69,10 @@ def segmentation(img):
     
 
 def upscale(image,prompt):
+    print("upscale")
+
     # pipeline = DiffusionPipeline.from_pretrained("stabilityai/stable-diffusion-x4-upscaler")
-    pipeline = StableDiffusionLatentUpscalePipeline.from_pretrained("stabilityai/sd-x2-latent-upscaler", torch_dtype=torch.float32)
+    pipeline = StableDiffusionLatentUpscalePipeline.from_pretrained("stabilityai/sd-x2-latent-upscaler", torch_dtype=torch.float16).to('cuda')
     upscaled_image = pipeline(prompt=prompt, image=image).images[0]
     return upscaled_image