Use gradio state instead of class variable

app.py

@@ -23,7 +23,6 @@ class ChatBotCheese:
         model_ckpt_path = hf_hub_download("alvanlii/fromage", "pretrained_ckpt.pth.tar")
         self.model = models.load_fromage(MODEL_DIR, model_ckpt_path)
         self.curr_image = None
-        self.chat_history = ''
 
     def add_image(self, state, image_in):
         state = state + [(f"![](/file={image_in.name})", "Ok, now type your message")]
@@ -35,14 +34,14 @@ class ChatBotCheese:
         image_pil.save(file_name)
         return file_name
 
-    def chat(self, input_text, state, ret_scale_factor, num_ims, num_words, temp):
-        # model_outputs = ["heyo", []]
-        self.chat_history += f'Q: {input_text} \nA:'
+    def chat(self, input_text, state, ret_scale_factor, num_ims, num_words, temp, chat_state):
+        chat_state.append(f'Q: {input_text} \nA:')
+        chat_history = " ".join(chat_state)
         if self.curr_image is not None:
-            model_outputs = self.model.generate_for_images_and_texts([self.curr_image, self.chat_history], num_words=num_words, max_num_rets=num_ims, ret_scale_factor=ret_scale_factor, temperature=temp)
+            model_outputs = self.model.generate_for_images_and_texts([self.curr_image, chat_history], num_words=num_words, max_num_rets=num_ims, ret_scale_factor=ret_scale_factor, temperature=temp)
         else:
-            model_outputs = self.model.generate_for_images_and_texts([self.chat_history], max_num_rets=num_ims, num_words=num_words, ret_scale_factor=ret_scale_factor, temperature=temp)
-        self.chat_history += ' '.join([s for s in model_outputs if type(s) == str]) + '\n'
+            model_outputs = self.model.generate_for_images_and_texts([chat_history], max_num_rets=num_ims, num_words=num_words, ret_scale_factor=ret_scale_factor, temperature=temp)
+        chat_state.append(' '.join([s for s in model_outputs if type(s) == str]) + '\n')
 
         im_names = []
         if len(model_outputs) > 1:
@@ -52,11 +51,10 @@ class ChatBotCheese:
         for im_name in im_names:
             response += f'<img src="/file={im_name}">'
         state.append((input_text, response.replace("[RET]", "")))
-        self.curr_image = None
-        return state, state    
+        # self.curr_image = None
+        return state, state, chat_state    
 
     def reset(self):
-        self.chat_history = ""
         self.curr_image = None
         return [], []
 
@@ -66,7 +64,7 @@ class ChatBotCheese:
                 """
                 ### FROMAGe: Grounding Language Models to Images for Multimodal Generation
                 Jing Yu Koh, Ruslan Salakhutdinov, Daniel Fried <br/>
-                [Paper](https://arxiv.org/abs/2301.13823) [Github](https://github.com/kohjingyu/fromage) <br/>
+                [Paper](https://arxiv.org/abs/2301.13823) [Github](https://github.com/kohjingyu/fromage) [Official Demo](https://huggingface.co/spaces/jykoh/fromage) <br/>
                 This is an unofficial Gradio demo for the paper FROMAGe <br/>
                 - Instructions (in order):
                   - [Optional] Upload an image (the button with a photo emoji)
@@ -83,6 +81,7 @@ class ChatBotCheese:
 
             chatbot = gr.Chatbot(elem_id="chatbot")
             gr_state = gr.State([])
+            gr_chat_state = gr.State([])
 
             with gr.Row():
                 with gr.Column(scale=0.85):
@@ -104,7 +103,7 @@ class ChatBotCheese:
                 gr.Image("example_3.png", label="Example 3")
                 
 
-            txt.submit(self.chat, [txt, gr_state, gr_ret_scale_factor, gr_num_ims, gr_num_words, gr_temp], [gr_state, chatbot])
+            txt.submit(self.chat, [txt, gr_state, gr_ret_scale_factor, gr_num_ims, gr_num_words, gr_temp, gr_chat_state], [gr_state, chatbot, gr_chat_state])
             txt.submit(lambda :"", None, txt)
             btn.upload(self.add_image, [gr_state, btn], [gr_state, chatbot])
             reset_btn.click(self.reset, [], [gr_state, chatbot])
@@ -119,5 +118,4 @@ def main():
     cheddar.main()
 
 if __name__ == "__main__":
-    cheddar = ChatBotCheese()
-    cheddar.main()
+    main()