Update app.py

app.py

@@ -17,33 +17,8 @@ from cheetah.conversation.conversation_llama2 import Chat, CONV_VISION
 from cheetah.models import *
 from cheetah.processors import *
 
-'''
-def parse_args():
-    parser = argparse.ArgumentParser(description="Demo")
-    parser.add_argument("--cfg-path", required=True, help="path to configuration file.")
-    parser.add_argument("--gpu-id", type=int, default=0, help="specify the gpu to load the model.")
-    parser.add_argument(
-        "--options",
-        nargs="+",
-        help="override some settings in the used config, the key-value pair "
-        "in xxx=yyy format will be merged into config file (deprecate), "
-        "change to --cfg-options instead.",
-    )
-    args = parser.parse_args()
-    return args
-
-def setup_seeds(seed = 50):
-    random.seed(seed)
-    np.random.seed(seed)
-    torch.manual_seed(seed)
-    torch.cuda.manual_seed_all(seed)
-    cudnn.benchmark = False
-    cudnn.deterministic = True
-'''
-
+# Initialize Cheetah
 print('Initializing Chat')
-#args = parse_args()
-#args = Namespace(cfg_path='eval_configs/cheetah_eval_vicuna.yaml', gpu_id=0, options=None)
 
 config = OmegaConf.load('eval_configs/cheetah_eval_llama2.yaml')
 cfg = Config.build_model_config(config)
@@ -56,28 +31,42 @@ chat = Chat(model, vis_processor, device='cuda:{}'.format(0))
 print('Initialization Finished')
 
 
+# Respond to user
 def respond(imgs, question):
 
+    # List of image paths
     raw_img_list = []
-    
+
+    # Go through each image
     for i, img in enumerate(imgs):
+        # Open image
         temp_img = Image.open(io.BytesIO(img))
+
+        # Convert to RGB
         temp_img = temp_img.convert(mode='RGB')
+
+        # Save image
         save_path = "./img" + str(i) + ".jpg"
         temp_img.save(save_path)
+
+        # Add to image path list
         raw_img_list.append(save_path)
     
     context = question
 
+    # Format: <Img><HereForImage></Img> <Img><HereForImage></Img> Question
     for i in raw_img_list:
         context = "<Img><HereForImage></Img> " + context
 
     print("Question: ", context)
+
+    # Get response from Cheetah
     llm_message = chat.answer(raw_img_list, context)
+    
     print("Answer: ", llm_message)
 
     return raw_img_list,llm_message
 
-#iface = gr.Interface(fn=respond, title="Cheetah", inputs=[gr.Image(type="pil"), "text"], outputs="text")
+
 iface = gr.Interface(fn=respond, title="Cheetah", inputs=[gr.File(file_count='multiple',file_types=['image'],type='binary'), "text"], outputs=[gr.Gallery(),"text"])
 iface.launch()