up

app.py

@@ -12,6 +12,9 @@ import random
 import io
 import base64
 import json
+from diffusers import DiffusionPipeline
+from diffusers import StableDiffusionLatentUpscalePipeline, StableDiffusionPipeline
+
 
 def removeFurniture(input_img1,
             input_img2,
@@ -65,8 +68,11 @@ def segmentation(img):
     return json.dumps(results)
     
 
-def upscale(image):
-    return image
+def upscale(image,prompt):
+    # pipeline = DiffusionPipeline.from_pretrained("stabilityai/stable-diffusion-x4-upscaler")
+    pipeline = StableDiffusionLatentUpscalePipeline.from_pretrained("stabilityai/sd-x2-latent-upscaler", torch_dtype=torch.float32)
+    upscaled_image = pipeline(prompt=prompt, image=image).images[0]
+    return upscaled_image
 
 with gr.Blocks() as app:    
     with gr.Row():
@@ -93,8 +99,9 @@ with gr.Blocks() as app:
                                                 gr.Image()])
         with gr.Column():  
             gr.Button("Segmentation").click(segmentation, inputs=gr.Image(type="pil"), outputs=gr.JSON())
+
         with gr.Column():
-            gr.Button("Upscale").click(upscale, inputs=gr.Image(type="pil"), outputs=gr.Image())
+            gr.Button("Upscale").click(upscale, inputs=[gr.Image(type="pil"),gr.Textbox(label="prompt",value="empty room")], outputs=gr.Image())
 
 app.launch(debug=True,share=True)