Add predict function for image processing and refactor http_bot for cleaner code

app.py

@@ -181,6 +181,23 @@ model = AutoModel.from_pretrained(
 tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True, use_fast=False)
 
 @spaces.GPU
+def predict(message, image_path):
+        pixel_values = load_image(image_path, max_num=6).to(torch.bfloat16).cuda()
+        print(f"pixel_values: {pixel_values}")
+        generation_config = dict(max_new_tokens= 700, do_sample=False, num_beams = 3, repetition_penalty=2.5)
+        print(f"######################")
+        print(f"message: {message}")
+        if pixel_values is not None:
+            question = '<image>\n'+message
+        else:
+            question = message
+        print("Model: ", model)
+        print("Tokenizer: ", tokenizer)
+        print("Question: ", question)
+        response, conv_history = model.chat(tokenizer, pixel_values, question, generation_config, history=None, return_history=True)
+        print(f"AI response: {response}")
+        return response, conv_history
+
 def http_bot(
     state,
     temperature,
@@ -231,23 +248,9 @@ def http_bot(
         # Stream output
         # response = requests.post(worker_addr, json=pload, headers=headers, stream=True, timeout=300)
         print(f"all_image_paths: {all_image_paths}")
-
-        pixel_values = load_image(all_image_paths[0], max_num=6).to(torch.bfloat16).cuda()
-        print(f"pixel_values: {pixel_values}")
-        generation_config = dict(max_new_tokens= 700, do_sample=False, num_beams = 3, repetition_penalty=2.5)
         message = state.get_user_message(source=state.USER)
-        print(f"######################")
         print(f"message: {message}")
-        if pixel_values is not None:
-            question = '<image>\n'+message
-        else:
-            question = message
-        print("Model: ", model)
-        print("Tokenizer: ", tokenizer)
-        print("Question: ", question)
-        response, conv_history = model.chat(tokenizer, pixel_values, question, generation_config, history=None, return_history=True)
-        print(f"AI response: {response}")
-
+        response, conv_history = predict(message, all_image_paths[0])
 
         # streamer = TextIteratorStreamer(
         #     tokenizer, skip_prompt=True, skip_special_tokens=True