debug

app.py

@@ -71,6 +71,26 @@ def run_llm_reasoning(caption, question, answer):
     output = llm.generate([{"prompt": prompt}], sampling_params=llm_sampling)
     return output[0].outputs[0].text
 
+##########################################
+# Streaming
+##########################################
+mllm_streamer = TextIteratorStreamer(processor.tokenizer, skip_prompt=True, skip_special_tokens=True, timeout=15)
+llm_streamer = TextIteratorStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True, timeout=15)
+
+def stream_response(model, inputs, streamer, prompt, gen_kwargs):
+    thread = Thread(target=model.generate, kwargs=dict(
+        streamer=streamer,
+        **inputs,
+        **gen_kwargs
+        )
+    )
+    thread.start()
+
+    generated_text = prompt
+    for new_text in streamer:
+        generated_text += new_text
+        yield generated_text
+
 ##########################################
 # Gradio part
 ##########################################
@@ -136,43 +156,42 @@ def http_bot(state):
         return
 
     # Retrive prompt
-    logging.info(state.messages)
-    logging.info("=================")
-    logging.info(str(state.messages))
-    logging.info("=================")
-    logging.info(state.messages[-1])
-    logging.info("=================")
-    logging.info(state.messages[-1][-1])
-    logging.info("=================")
-    logging.info(state.messages[-1][-1][0])
-    logging.info("=================")
     prompt = state.messages[-1][-1][0]
     all_images = state.get_images(return_pil=True)[0]
     pload = {"prompt": prompt, "images": f'List of {len(state.get_images())} images: {all_images}'}
     logging.info(f"==== request ====\n{pload}")
     
-    return
-    
     # Construct prompt
     cap_msgs, qa_msgs = build_messages(all_images, prompt)
     cap_prompt = processor.apply_chat_template([cap_msgs], tokenize=False, add_generation_prompt=True)
     qa_prompt = processor.apply_chat_template([qa_msgs], tokenize=False, add_generation_prompt=True)
-
     image_tensor, _ = process_vision_info(cap_msgs)
-    tentative_answer = run_mllm_tentative(image_tensor, cap_prompt, qa_prompt)
-    state.append_message(state.roles[1], "# Tentative Response\n\n" + tentative_answer)
-    logging.info("# Tentative Response\n\n" + tentative_answer)
-    yield (state, state.to_gradio_chatbot_public()) + (disable_btn,) * 2
+    cap_inputs = processor(text=[cap_prompt], images=image_tensor, return_tensors="pt").to(mllm.device)
+    qa_inputs = processor(text=[qa_prompt], images=image_tensor, return_tensors="pt").to(mllm.device)    
+    
+    # Step 1: Tentative Response
+    state.append_message(state.roles[1], "# Tentative Response\n\n▌")
+    try:
+        for generated_text in stream_response(mllm, qa_inputs, mllm_streamer, qa_prompt, mllm_sampling):
+            output = generated_text[len(prompt):].strip()
+            state.messages[-1][-1] = "# Tentative Response\n\n" + output + "▌"
+            yield (state, state.to_gradio_chatbot_public()) + (disable_btn,) * 2
+    except Exception as e:
+        os.system("nvidia-smi")
+        logging.info(traceback.print_exc())
+        state.messages[-1][-1] = server_error_msg
+        yield (state, state.to_gradio_chatbot_public()) + (enable_btn,) * 2
+        return
     
-    caption_text = run_mllm_caption(image_tensor, cap_prompt, qa_prompt)
-    state.append_message(state.roles[1], "# Caption\n\n" + caption_text)
-    logging.info("# Caption\n\n" + caption_text)
-    yield (state, state.to_gradio_chatbot_public()) + (disable_btn,) * 2
+    # caption_text = run_mllm_caption(image_tensor, cap_prompt, qa_prompt)
+    # state.append_message(state.roles[1], "# Caption\n\n" + caption_text)
+    # logging.info("# Caption\n\n" + caption_text)
+    # yield (state, state.to_gradio_chatbot_public()) + (disable_btn,) * 2
         
-    final_answer = run_llm_reasoning(caption_text, QUESTION, tentative_answer)
-    state.append_message(state.roles[1], "# Final Response\n\n" + final_answer)
-    logging.info("# Final Response\n\n" + final_answer)
-    yield (state, state.to_gradio_chatbot_public()) + (enable_btn,) * 2
+    # final_answer = run_llm_reasoning(caption_text, QUESTION, tentative_answer)
+    # state.append_message(state.roles[1], "# Final Response\n\n" + final_answer)
+    # logging.info("# Final Response\n\n" + final_answer)
+    # yield (state, state.to_gradio_chatbot_public()) + (enable_btn,) * 2
 
 ############
 # Layout Markdown