Update app.py

app.py

@@ -41,7 +41,7 @@ priority = {
 
 
 @spaces.GPU
-def run_llava(prompt, pil_image):
+def run_llava(prompt, pil_image, temperature, top_p, max_new_tokens):
     image_size = pil_image.size
     image_tensor = image_processor.preprocess(pil_image, return_tensors='pt')['pixel_values'].half().cuda()
     # images_tensor = load_images(images, image_processor)
@@ -54,11 +54,11 @@ def run_llava(prompt, pil_image):
             images=image_tensor,
             image_sizes=[image_size],
             do_sample=True,
-            temperature=0.2,
-            top_p=0.95,
+            temperature=temperature,
+            top_p=top_p,
             top_k=50,
             num_beams=2,
-            max_new_tokens=1024,
+            max_new_tokens=max_new_tokens,
             use_cache=True,
             stopping_criteria=[KeywordsStoppingCriteria(['}'], tokenizer, input_ids)]
         )
@@ -84,11 +84,12 @@ def get_conv_log_filename():
     name = os.path.join(LOGDIR, f"{t.year}-{t.month:02d}-{t.day:02d}-conv.json")
     return name
 
+
 def get_model_list():
     models = [
-                 'LukasHug/LlavaGuard-7B-hf',
-                 'LukasHug/LlavaGuard-13B-hf',
-                 'LukasHug/LlavaGuard-34B-hf', ]
+        'LukasHug/LlavaGuard-7B-hf',
+        'LukasHug/LlavaGuard-13B-hf',
+        'LukasHug/LlavaGuard-34B-hf', ][:2]
     return models
 
 
@@ -249,7 +250,6 @@ def llava_bot(state, model_selector, temperature, top_p, max_new_tokens, request
         new_state.append_message(new_state.roles[1], None)
         state = new_state
 
-
     # Construct prompt
     prompt = state.get_prompt()
 
@@ -262,13 +262,12 @@ def llava_bot(state, model_selector, temperature, top_p, max_new_tokens, request
             os.makedirs(os.path.dirname(filename), exist_ok=True)
             image.save(filename)
 
-    output = run_llava(prompt, all_images[0])
+    output = run_llava(prompt, all_images[0], temperature, top_p, max_new_tokens)
 
     state.messages[-1][-1] = output
 
     yield (state, state.to_gradio_chatbot()) + (disable_btn,) * 5
 
-
     finish_tstamp = time.time()
     logger.info(f"{output}")
 
@@ -406,7 +405,10 @@ def build_demo(embed_mode, cur_dir=None, concurrency_count=10):
             [textbox, upvote_btn, downvote_btn, flag_btn]
         )
 
-        model_selector.change(load_selected_model)
+        model_selector.change(
+            load_selected_model,
+            [model_selector],
+        )
 
         regenerate_btn.click(
             regenerate,
@@ -517,7 +519,6 @@ Set the environment variable `model` to change the model:
     tokenizer, model, image_processor, context_len = load_pretrained_model(model_path, None, model_name, token=api_key)
     model.config.tokenizer_model_max_length = 2048 * 2
 
-
     exit_status = 0
     try:
         demo = build_demo(embed_mode=False, cur_dir='./', concurrency_count=concurrency_count)
@@ -534,4 +535,4 @@ Set the environment variable `model` to change the model:
         print(e)
         exit_status = 1
     finally:
-        sys.exit(exit_status)
+        sys.exit(exit_status)