add support for all three steps

app.py

@@ -16,8 +16,8 @@ auth_token = os.environ.get("TOKEN_FROM_SECRET")
 # LLM part
 ##########################################
 import torch
-from transformers import AutoProcessor, AutoTokenizer
-from transformers import Qwen2ForCausalLM, Qwen2_5_VLForConditionalGeneration, TextIteratorStreamer
+from transformers import AutoProcessor, AutoTokenizer, AutoModelForCausalLM
+from transformers import Qwen2_5_VLForConditionalGeneration, TextIteratorStreamer
 from qwen_vl_utils import process_vision_info
 from threading import Thread
 
@@ -36,7 +36,7 @@ processor = AutoProcessor.from_pretrained(MLLM_MODEL_PATH)
 tokenizer = AutoTokenizer.from_pretrained(LLM_MODEL_PATH)
 
 mllm = Qwen2_5_VLForConditionalGeneration.from_pretrained(MLLM_MODEL_PATH, torch_dtype=torch.bfloat16, device_map="auto")
-llm = Qwen2ForCausalLM.from_pretrained(LLM_MODEL_PATH, torch_dtype=torch.bfloat16, device_map="auto")
+llm = AutoModelForCausalLM.from_pretrained(LLM_MODEL_PATH, torch_dtype=torch.bfloat16, device_map="auto")
 
 mllm_sampling = dict(do_sample=False, temperature=0, max_new_tokens=8192)
 llm_sampling = dict(temperature=0.6, top_p=0.95, max_new_tokens=8192)
@@ -52,25 +52,6 @@ def build_messages(image_path, question):
     ]
     return cap_msgs, qa_msgs
 
-# === Run Captioning and QA ===
-def run_mllm_tentative(image_tensor, cap_prompt, qa_prompt):
-    qa_output = mllm.generate([{"multi_modal_data": {"image": image_tensor}, "prompt": qa_prompt[0]}], sampling_params=mllm_sampling)
-    return qa_output[0].outputs[0].text
-
-def run_mllm_caption(image_tensor, cap_prompt, qa_prompt):
-    cap_output = mllm.generate([{"multi_modal_data": {"image": image_tensor}, "prompt": cap_prompt[0]}], sampling_params=mllm_sampling)
-    return cap_output[0].outputs[0].text
-
-# === Final Reasoning Step ===
-def run_llm_reasoning(caption, question, answer):
-    messages = [
-        {"role": "system", "content": SYSTEM_PROMPT_LLM},
-        {"role": "user", "content": LLM_PROMPT.format(caption, question, answer)}
-    ]
-    prompt = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
-    output = llm.generate([{"prompt": prompt}], sampling_params=llm_sampling)
-    return output[0].outputs[0].text
-
 ##########################################
 # Streaming
 ##########################################
@@ -204,18 +185,31 @@ def http_bot(state):
     logging.info(f"Query-conditioned Caption: {caption_text}")
     state.messages[-1][-1] = state.messages[-1][-1][:-1]
     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
+    
+    # Step 3: Text-only Reasoning
+    reason_msgs = [
+        {"role": "system", "content": SYSTEM_PROMPT_LLM},
+        {"role": "user", "content": LLM_PROMPT.format(caption_text, prompt, tentative_answer)}
+    ]
+    reason_prompt = tokenizer.apply_chat_template(reason_msgs, tokenize=False, add_generation_prompt=True)
+    reason_inputs = tokenizer(reason_prompt, return_tensors="pt").to(llm.device)
+    
+    state.append_message(state.roles[1], "# Text-only Reasoning\n\n▌")
+    try:
+        for generated_text in stream_response(llm, reason_inputs, llm_streamer, reason_prompt, llm_sampling):
+            output = generated_text[len(reason_prompt):].strip()
+            state.messages[-1][-1] = "# Text-only Reasoning\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
+    final_response = output
+    logging.info(f"Text-only Reasoning: {final_response}")
+    state.messages[-1][-1] = state.messages[-1][-1][:-1]
+    yield (state, state.to_gradio_chatbot_public()) + (disable_btn,) * 2
 
 ############
 # Layout Markdown