return results in gallery

app.py

@@ -21,8 +21,8 @@ def generate_relevancy(
     labels = labels.split(",")
     prompts = [prompt]
     assert img.dtype == np.uint8
+    img = Image.fromarray(img).resize((244 * 2, 244 * 2))
     h, w, c = img.shape
-    start = time()
     grads = ClipWrapper.get_clip_saliency(
         img=img,
         text_labels=np.array(labels),
@@ -32,16 +32,13 @@ def generate_relevancy(
     if subtract_mean:
         grads -= grads.mean(axis=0)
     grads = grads.cpu().numpy()
-    num_axes = int(np.ceil(np.sqrt(len(labels))))
-    fig, axes = plt.subplots(num_axes, num_axes)
-    if num_axes == 1:
-        axes = [axes]
-    else:
-        axes = axes.flatten()
     vmin = 0.002
     cmap = plt.get_cmap("jet")
     vmax = 0.008
-    for ax, label_grad, label in zip(axes, grads, labels):
+
+    returns = []
+    for label_grad, label in zip(grads, labels):
+        fig, ax = plt.subplots(1, 1)
         ax.axis("off")
         ax.imshow(img)
         ax.set_title(label, fontsize=12)
@@ -50,10 +47,10 @@ def generate_relevancy(
         grad = 1 - grad
         colored_grad[..., -1] = grad * 0.7
         ax.imshow(colored_grad)
-    plt.tight_layout(pad=0)
-    img = plot_to_png(fig)
-    plt.close(fig)
-    return img
+        plt.tight_layout(pad=0)
+        returns.append(plot_to_png(fig))
+        plt.close(fig)
+    return returns
 
 
 iface = gr.Interface(
@@ -69,7 +66,7 @@ iface = gr.Interface(
         ),
         gr.Checkbox(value=True, label="subtract mean"),
     ],
-    outputs=gr.Image(type="numpy"),
+    outputs=gr.Gallery(label="Relevancy Maps", type="numpy"),
     examples=[
         [
             "https://semantic-abstraction.cs.columbia.edu/downloads/matterport.png",