Update app.py

app.py

@@ -1,5 +1,57 @@
 import gradio as gr
 from PIL import Image 
+import argparse
+import os
+import random
+
+import numpy as np
+import torch
+import torch.backends.cudnn as cudnn
+
+from omegaconf import OmegaConf
+from cheetah.common.config import Config
+from cheetah.common.registry import registry
+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
+
+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)
+model_cls = registry.get_model_class(cfg.model.arch)
+model = model_cls.from_config(cfg.model).to('cuda:{}'.format(0))
+
+vis_processor_cfg = cfg.preprocess.vis_processor.eval
+vis_processor = registry.get_processor_class(vis_processor_cfg.name).from_config(vis_processor_cfg)
+chat = Chat(model, vis_processor, device='cuda:{}'.format(0))
+print('Initialization Finished')
+
 
 def chat(img, question):
     img.save("img.jpg")