Update app.py

app.py

@@ -17,6 +17,8 @@ from PIL import Image
 from cobra import load
 import time
 
+subprocess.check_call([sys.executable, '-m', 'pip', 'install', 'packaging'])
+subprocess.check_call([sys.executable, '-m', 'pip', 'install', 'ninja'])
 subprocess.check_call([sys.executable, '-m', 'pip', 'install', 'mamba-ssm'])
 subprocess.check_call([sys.executable, '-m', 'pip', 'install', 'causal-conv1d'])
 
@@ -31,10 +33,11 @@ else:
 vlm.to(DEVICE, dtype=DTYPE)
 
 prompt_builder = vlm.get_prompt_builder()
-system_prompt = prompt_builder.system_prompt
 
 @spaces.GPU
-def bot_streaming(message, history):
+def bot_streaming(message, history, temperature, top_k, max_new_tokens):
+    if len(history) == 0:
+        prompt_builder.prompt, prompt_builder.turn_count = "", 0
     print(message)
     if message["files"]:
         image = message["files"][-1]["path"]
@@ -47,48 +50,32 @@ def bot_streaming(message, history):
   
     image = Image.open(image).convert("RGB")
     
-    prompt_builder.add_turn(role="human", message=message)
+    prompt_builder.add_turn(role="human", message=message['text'])
     prompt_text = prompt_builder.get_prompt()
-
+    
     # Generate from the VLM
-    generated_text = vlm.generate(
-        image,
-        prompt_text,
-        cg=True,
-        do_sample=False,
-        temperature=1.0,
-        max_new_tokens=2048,
-        # do_sample=cfg.do_sample,
-        # temperature=cfg.temperature,
-        # max_new_tokens=cfg.max_new_tokens,
-    )
+    with torch.no_grad():
+        generated_text = vlm.generate(
+            image,
+            prompt_text,
+            cg=True,
+            do_sample=True,
+            temperature=temperature,
+            top_k=top_k,
+            max_new_tokens=max_new_tokens,
+        )
     prompt_builder.add_turn(role="gpt", message=generated_text)
 
-    # streamer = TextIteratorStreamer(processor, **{"skip_special_tokens": True})
-    # generation_kwargs = dict(inputs, streamer=streamer, max_new_tokens=100)
-    # generation_kwargs = dict(image, prompt_text, cg=True, do_sample=cfg.do_sample, temperature=cfg.temperature, max_new_tokens=cfg.max_new_tokens)
-    generation_kwargs = dict(image, prompt_text, cg=True, do_sample=True, temperature=1.0, max_new_tokens=2048)
-
-    thread = Thread(target=vlm.generate, kwargs=generation_kwargs)
-    thread.start()
-
-    text_prompt =f"[INST]  \n{message['text']} [/INST]"
-    print(generated_text)
-    
-
-    buffer = ""
+    time.sleep(0.04)
     yield generated_text
-    # for new_text in streamer:
-        
-    #     buffer += new_text
-        
-    #     generated_text_without_prompt = buffer[len(text_prompt):]
-    #     time.sleep(0.04)
-    #     yield generated_text_without_prompt
-
-
-demo = gr.ChatInterface(fn=bot_streaming, title="LLaVA Next", examples=[{"text": "What is on the flower?", "files":["./bee.jpg"]},
-                                                                      {"text": "How to make this pastry?", "files":["./baklava.png"]}], 
-                        description="Try [LLaVA Next](https://huggingface.co/papers/2310.03744) in this demo. Upload an image and start chatting about it, or simply try one of the examples below.",
-                        stop_btn="Stop Generation", multimodal=True)
+    
+    
+demo = gr.ChatInterface(fn=bot_streaming, 
+                        additional_inputs=[gr.Slider(0, 1, value=0.2, label="Temperature"),
+                                           gr.Slider(1, 3, value=1, step=1, label="Top k"),
+                                           gr.Slider(1, 2048, value=256, step=1, label="Max New Tokens")],
+                        title="Cobra", 
+                        description="Try [Cobra](https://huggingface.co/papers/2403.14520) in this demo. Upload an image and start chatting about it.",
+                        stop_btn="Stop Generation", multimodal=True,
+                        examples=[{"text": "Describe this image", "files":["./cobra.png"]}])
 demo.launch(debug=True)