Update app.py

app.py

@@ -73,8 +73,7 @@ def infer(img_path, seed=0, randomize_seed=False, target_modalities=None, top_p=
         seed = None
     img = img_from_path(img_path).to(device)
     preds = sampler({'rgb@224': img}, seed=seed, target_modalities=target_modalities, top_p=top_p, top_k=top_k) 
-    sampler.plot_modalities(preds, ncols_max=4, use_fixed_plotting_order=True, save_path='./output.png')
-    return './output.png'
+    return sampler.modalities_to_pil(preds, use_fixed_plotting_order=True, resize=512)
 
 
 examples = [
@@ -140,7 +139,10 @@ with gr.Blocks(css=css, theme=gr.themes.Base()) as demo:
                     top_p = gr.Slider(label="Top-p", minimum=0.0, maximum=1.0, step=0.01, value=0.8)
                     top_k = gr.Slider(label="Top-k", minimum=0.0, maximum=1.0, step=0.01, value=0.0)
         
-            result = gr.Image(label="Predictions", show_label=False)
+            result = gr.Gallery(
+                label="Predictions", show_label=True, elem_id="gallery", type='pil',
+                columns=[4], rows=None, object_fit="contain", height="auto"
+            )
 
         gr.Examples(
             examples = examples,