add mask

app.py

@@ -101,7 +101,7 @@ def gr_caption(x):
         return np.clip(img, a_min=0., a_max=1.)
 
     masked, masked_recon, recon, latent = visual_recon(x, model)
-    caption_from_model = caption(8, latent, model, tokenizer, )
+    caption_from_model = caption(20, latent, model, tokenizer, )
     
     masked, masked_recon, recon = map(unnorm_pix, (masked, masked_recon, recon))
     
@@ -109,6 +109,11 @@ def gr_caption(x):
 
 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)), 'text'])
+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)), 
+                             'text'])
 demo.launch()