up test

app.py

@@ -14,6 +14,7 @@ import base64
 import json
 from diffusers import DiffusionPipeline
 from diffusers import StableDiffusionLatentUpscalePipeline, StableDiffusionPipeline
+from diffusers import StableDiffusionUpscalePipeline
 
 
 def removeFurniture(input_img1,
@@ -68,13 +69,19 @@ def segmentation(img):
     return json.dumps(results)
     
 
-def upscale(image,prompt):
+def upscale(image, prompt):
     print("upscale")
 
     image = image.thumbnail((512, 512))
 
-    pipeline = DiffusionPipeline.from_pretrained("stabilityai/stable-diffusion-x4-upscaler")
+    model_id = "stabilityai/stable-diffusion-x4-upscaler"
+    pipeline = StableDiffusionUpscalePipeline.from_pretrained(model_id, torch_dtype=torch.float16)
+    pipeline = pipeline.to("cuda")
+
+
+    #pipeline = DiffusionPipeline.from_pretrained("stabilityai/stable-diffusion-x4-upscaler")
     # 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