UI changes

app.py

@@ -1,5 +1,5 @@
 import tempfile
-from share_btn import community_icon_html, loading_icon_html, share_js
+from share_btn import community_icon_html, loading_icon_html, share_js, save_js
 import huggingface_hub
 import gradio as gr
 from fromage import utils
@@ -15,7 +15,7 @@ os.environ["HF_HUB_ENABLE_HF_TRANSFER"] = "False"
 css = """
     #share-btn-container {
         display: flex; padding-left: 0.5rem !important; padding-right: 0.5rem !important; background-color: #000000; justify-content: center; align-items: center; border-radius: 9999px !important; width: 13rem;
-        margin-top: 10px;
+        margin-top: 3px;
         margin-left: auto;
         flex: unset;
     }
@@ -33,6 +33,22 @@ css = """
         display: none !important;
     }
     #chatbot { min-height: 300px; }
+    #save-btn {
+        background-image: linear-gradient(to right bottom, rgba(130,217,244, 0.9), rgba(158,231,214, 1.0));
+    }
+    #save-btn:hover {
+        background-image: linear-gradient(to right bottom, rgba(130,217,244, 0.9), rgba(158,231,214, 1.0));
+    }
+    #share-btn-2 {
+        background-image: linear-gradient(to right bottom, rgba(130,217,244, 0.9), rgba(158,231,214, 1.0));
+    }
+    #share-btn-2:hover {
+        background-image: linear-gradient(to right bottom, rgba(130,217,244, 0.9), rgba(158,231,214, 1.0));
+    }
+    .message .user {
+    }
+    .message .bot {
+    }
 """
 
 examples = [
@@ -133,7 +149,7 @@ def generate_for_prompt(input_text, state, ret_scale_factor, max_num_rets, num_w
     # Set input image to None.
     print('state', state, flush=True)
     print('updated state', [conversation, chat_history, None], flush=True)
-    return [conversation, chat_history, None], conversation, gr.update(visible=True)
+    return [conversation, chat_history, None], conversation, gr.update(visible=True), gr.update(visible=True)
 
 
 with gr.Blocks(css=css) as demo:
@@ -161,11 +177,11 @@ with gr.Blocks(css=css) as demo:
                         "Submit", interactive=True, variant="primary")
                     clear_last_btn = gr.Button("Undo")
                     clear_btn = gr.Button("Reset All")
+                    with gr.Row(visible=False) as save_group:
+                        save_button = gr.Button("💾 Save Conversation as .png", elem_id="save-btn")
+
                     with gr.Row(visible=False) as share_group:
-                        with gr.Group(elem_id="share-btn-container"):
-                            community_icon = gr.HTML(community_icon_html)
-                            loading_icon = gr.HTML(loading_icon_html)
-                            share_button = gr.Button("Share to community", elem_id="share-btn")
+                        share_button = gr.Button("🤗 Share to Community", elem_id="share-btn-2")
 
 
         with gr.Column(scale=0.3, min_width=200):
@@ -183,17 +199,18 @@ with gr.Blocks(css=css) as demo:
             # ).style(grid=[2], height="auto")
 
     text_input.submit(generate_for_prompt, [text_input, gr_state, ret_scale_factor,
-                      max_ret_images, gr_max_len, gr_temperature], [gr_state, chatbot, share_group])
+                      max_ret_images, gr_max_len, gr_temperature], [gr_state, chatbot, share_group, save_group])
     text_input.submit(lambda: "", None, text_input)  # Reset chatbox.
     submit_btn.click(generate_for_prompt, [text_input, gr_state, ret_scale_factor,
-                     max_ret_images, gr_max_len, gr_temperature], [gr_state, chatbot, share_group])
+                     max_ret_images, gr_max_len, gr_temperature], [gr_state, chatbot, share_group, save_group])
     submit_btn.click(lambda: "", None, text_input)  # Reset chatbox.
 
     image_btn.upload(upload_image, [gr_state, image_btn], [gr_state, chatbot])
     clear_last_btn.click(reset_last, [gr_state], [gr_state, chatbot])
     clear_btn.click(reset, [], [gr_state, chatbot])
     share_button.click(None, [], [], _js=share_js)
+    save_button.click(None, [], [], _js=save_js)
 
 
-# demo.queue(concurrency_count=1, api_open=False, max_size=16)
-demo.launch(debug=True, server_name="127.0.0.1")
+demo.queue(concurrency_count=1, api_open=False, max_size=16)
+demo.launch(debug=True, server_name="0.0.0.0")

fromage/models.py

@@ -635,7 +635,7 @@ def load_fromage(embeddings_dir: str, model_args_path: str, model_ckpt_path: str
   assert len(ret_token_idx) == 1, ret_token_idx
   model_kwargs['retrieval_token_idx'] = ret_token_idx[0]
 
-  debug = True
+  debug = False
   if debug:
     model_kwargs['opt_version'] = 'facebook/opt-125m'
     model_kwargs['visual_encoder'] = 'openai/clip-vit-base-patch32'

share_btn.py

@@ -68,4 +68,45 @@ async () => {
     shareIconEl.style.removeProperty('display');
     loadingIconEl.style.display = 'none';
 }
+"""
+
+save_js = """
+async () => {
+    const html2canvas = (await import('https://cdnjs.cloudflare.com/ajax/libs/html2canvas/1.4.1/html2canvas.esm.js')).default;
+
+    function saveAs(uri, filename) {
+      var link = document.createElement('a');
+      if (typeof link.download === 'string') {
+        link.href = uri;
+        link.download = filename;
+  
+        //Firefox requires the link to be in the body
+        document.body.appendChild(link);
+  
+        //simulate click
+        link.click();
+  
+        //remove the link when done
+        document.body.removeChild(link);
+      } else {
+        window.open(uri);
+      }
+    }
+
+    async function getImageFile(div) {
+        return new Promise((resolve, reject) =>
+            html2canvas(div)
+                .then((canvas) => {
+                    const imageId = Date.now();
+                    const fileName = "FROMAGe-" + imageId + ".jpg";
+                    saveAs(canvas.toDataURL(), fileName);
+                })
+
+        )
+    }
+    const gradioEl = document.querySelector("gradio-app").shadowRoot || document.querySelector('body > gradio-app');
+    const chatbotEl = gradioEl.querySelector('#chatbot')
+    const imageFile = await getImageFile(chatbotEl);
+    console.log(imageFile);
+}
 """