add mask

app.py

@@ -104,16 +104,15 @@ def gr_caption(x):
     caption_from_model = caption(20, latent, model, tokenizer, )
     
     masked, masked_recon, recon = map(unnorm_pix, (masked, masked_recon, recon))
+    return_img = np.concatenate([masked, masked_recon, recon], axis=1)
     
-    return masked, masked_recon, recon, caption_from_model
+    return return_img, caption_from_model
 
 import gradio as gr
 
 demo = gr.Interface(gr_caption, 
                     inputs=[gr.Image(shape=(224, 224))], 
-                    outputs=[gr.Image(shape=(224, 224)), 
-                             gr.Image(shape=(224, 224)), 
-                             gr.Image(shape=(224, 224)), 
+                    outputs=[gr.Image(shape=(224, 224 * 3)), 
                              'text'])
 demo.launch()