Add application file

app.py

@@ -1,8 +1,9 @@
 import os
-
+import pickle
 import gradio as gr
 import torch
 from PIL import Image
+import matplotlib.pyplot as plt
 
 from mmgpt.models.builder import create_model_and_transforms
 
@@ -13,7 +14,9 @@ response_split = "### Response:"
 class Inferencer:
 
     def __init__(self, finetune_path, llama_path, open_flamingo_path):
+        print("inferencer initialization begun")
         ckpt = torch.load(finetune_path, map_location="cpu")
+        print("ckpt: ", ckpt)
         if "model_state_dict" in ckpt:
             state_dict = ckpt["model_state_dict"]
             # remove the "module." prefix
@@ -23,6 +26,7 @@ class Inferencer:
             }
         else:
             state_dict = ckpt
+        print("state_dict has been set")
         tuning_config = ckpt.get("tuning_config")
         if tuning_config is None:
             print("tuning_config not found in checkpoint")
@@ -46,15 +50,19 @@ class Inferencer:
         self.model = model
         self.image_processor = image_processor
         self.tokenizer = tokenizer
+        print("finished inferencer initialization")
 
     def __call__(self, prompt, imgpaths, max_new_token, num_beams, temperature,
                  top_k, top_p, do_sample):
+        print("inferecer called")
         if len(imgpaths) > 1:
             raise gr.Error(
                 "Current only support one image, please clear gallery and upload one image"
             )
         lang_x = self.tokenizer([prompt], return_tensors="pt")
+        print("tokenized")
         if len(imgpaths) == 0 or imgpaths is None:
+            print("imgpath len is 0 or None")
             for layer in self.model.lang_encoder._get_decoder_layers():
                 layer.condition_only_lang_x(True)
             output_ids = self.model.lang_encoder.generate(
@@ -70,10 +78,16 @@ class Inferencer:
             for layer in self.model.lang_encoder._get_decoder_layers():
                 layer.condition_only_lang_x(False)
         else:
+            print("imgpath is valid")
             images = (Image.open(fp) for fp in imgpaths)
+            print("images retrieved")
             vision_x = [self.image_processor(im).unsqueeze(0) for im in images]
             vision_x = torch.cat(vision_x, dim=0)
             vision_x = vision_x.unsqueeze(1).unsqueeze(0).half()
+            print("vision_x retrieved")
+            torch.cuda.empty_cache()
+            print(f"Allocated GPU memory: {torch.cuda.memory_allocated() / 1024**3:.2f} GB")
+            print(f"Available GPU memory: {torch.cuda.memory_reserved() / 1024**3:.2f} GB")
 
             output_ids = self.model.generate(
                 vision_x=vision_x.cuda(),
@@ -86,12 +100,24 @@ class Inferencer:
                 top_p=top_p,
                 do_sample=do_sample,
             )[0]
+            print("output_ids retrieved")
         generated_text = self.tokenizer.decode(
             output_ids, skip_special_tokens=True)
-        # print(generated_text)
+        print("text generated:", generated_text)
         result = generated_text.split(response_split)[-1].strip()
+        print("result: ", result)
         return result
 
+    def save(self, file_path):
+        print("Saving model components...")
+        data = {
+            "model_state_dict": self.model.state_dict(),
+            "tokenizer": self.tokenizer,
+            "image_processor": self.image_processor,
+        }
+        with open(file_path, "wb") as f:
+            pickle.dump(data, f)
+        print(f"Model components saved to {file_path}")
 
 class PromptGenerator:
 
@@ -103,7 +129,7 @@ class PromptGenerator:
         sep: str = "\n\n### ",
         buffer_size=0,
     ):
-        self.all_history = list()
+        self.all_history = [("user", "Welcome to the chatbot!")]
         self.ai_prefix = ai_prefix
         self.user_prefix = user_prefix
         self.buffer_size = buffer_size
@@ -217,16 +243,23 @@ def bot(
     state.sep = seperator
     state.buffer_size = history_buffer
     if image:
+        print(image)
+        print(text)
         state.add_message(user_prefix, (text, image))
+        print("added message")
     else:
         state.add_message(user_prefix, text)
     state.add_message(ai_prefix, None)
+    print("added ai_prefix message")
     inputs = state.get_prompt()
+    print("retrived inputs")
     image_paths = state.get_images()[-1:]
+    print("retrieved image_paths")
 
     inference_results = inferencer(inputs, image_paths, max_new_token,
                                    num_beams, temperature, top_k, top_p,
                                    do_sample)
+    print(inference_results)
     state.all_history[-1][-1] = inference_results
     memory_allocated = str(round(torch.cuda.memory_allocated() / 1024**3,
                                  2)) + 'GB'
@@ -284,14 +317,13 @@ def build_conversation_demo():
             with gr.Column(scale=6):
                 with gr.Row():
                     with gr.Column():
-                        chatbot = gr.Chatbot(elem_id="chatbot").style(
-                            height=750)
+                        chatbot = gr.Chatbot(elem_id="chatbot", height=750)
                 with gr.Row():
                     with gr.Column(scale=8):
                         textbox = gr.Textbox(
                             show_label=False,
                             placeholder="Enter text and press ENTER",
-                        ).style(container=False)
+                            container=False)
                         submit_btn = gr.Button(value="Submit")
                         clear_btn = gr.Button(value="🗑️  Clear history")
         cur_dir = os.path.dirname(os.path.abspath(__file__))
@@ -354,7 +386,6 @@ def build_conversation_demo():
                         [state, chatbot, textbox, imagebox, model_inputs])
     return demo
 
-
 if __name__ == "__main__":
     llama_path = "checkpoints/llama-7b_hf"
     open_flamingo_path = "checkpoints/OpenFlamingo-9B/checkpoint.pt"
@@ -365,8 +396,11 @@ if __name__ == "__main__":
         open_flamingo_path=open_flamingo_path,
         finetune_path=finetune_path)
     init_memory = str(round(torch.cuda.memory_allocated() / 1024**3, 2)) + 'GB'
+
+    inferencer.save("inferencer.pkl")
+
     demo = build_conversation_demo()
-    demo.queue(concurrency_count=3)
+    demo.queue()
     IP = "0.0.0.0"
     PORT = 8997
     demo.launch(server_name=IP, server_port=PORT, share=True)