model selection added, tiling forced to 2048

app.py

@@ -4,36 +4,54 @@ from gradio_imageslider import ImageSlider
 from image_gen_aux import UpscaleWithModel
 from image_gen_aux.utils import load_image
 
+MODELS = {"UltraSharp": "OzzyGT/UltraSharp", "DAT X4": "OzzyGT/DAT_X4"}
+
 
 @spaces.GPU
-def upscale_image(image):
+def upscale_image(image, model_selection):
     original = load_image(image)
 
-    upscaler = UpscaleWithModel.from_pretrained("OzzyGT/UltraSharp").to("cuda")
-    image = upscaler(original)
+    upscaler = UpscaleWithModel.from_pretrained(MODELS[model_selection]).to("cuda")
+    image = upscaler(original, tiling=True, tile_width=2048, tile_height=2048)
 
     return original, image
 
 
-title = """<h1 align="center">Image Upscaler</h1>"""
+def clear_result():
+    return gr.update(value=None)
+
+
+title = """<h1 align="center">Image Upscaler</h1>
+<div align="center">This space is a showcase of the different super resolution models you can use to upscale with the 
+<a href="https://github.com/asomoza/image_gen_aux">Image Generation Auxiliary Tools</a> library.</div>
+"""
 
 with gr.Blocks() as demo:
     gr.HTML(title)
     with gr.Row():
         with gr.Column():
             input_image = gr.Image(type="pil", label="Input Image")
+
+            model_selection = gr.Dropdown(
+                choices=list(MODELS.keys()),
+                value="UltraSharp",
+                label="Model",
+            )
+
             run_button = gr.Button("Upscale")
         with gr.Column():
             result = ImageSlider(
                 interactive=False,
                 label="Generated Image",
-                elem_id="result-image",
-                position=0.1,
             )
 
     run_button.click(
+        fn=clear_result,
+        inputs=None,
+        outputs=result,
+    ).then(
         fn=upscale_image,
-        inputs=[input_image],
+        inputs=[input_image, model_selection],
         outputs=result,
     )