update backend api

app.py

@@ -47,40 +47,6 @@ def write2file(path, content):
             fout.write(content)
 
 
-def sort_models(models):
-    def custom_sort_key(model_name):
-        # InternVL-Chat-V1-5 should be the first item
-        if model_name == "InternVL2-Pro":
-            return (2, model_name)  # 2 indicates highest precedence
-        elif model_name.startswith("InternVL2-8B"):
-            return (1, model_name)  # 0 indicates highest precedence
-        else:
-            return (0, model_name)  # 0 indicates normal order
-
-    models.sort(key=custom_sort_key, reverse=True)
-    # try:  # We have five InternVL-Chat-V1-5 models, randomly choose one to be the first
-    #     first_three = models[:4]
-    #     random.shuffle(first_three)
-    #     models[:4] = first_three
-    # except:
-    #     pass
-    return models
-
-
-def get_model_list():
-    logger.info(f"Call `get_model_list`")
-    ret = requests.post(args.controller_url + "/refresh_all_workers")
-    logger.info(f"status_code from `get_model_list`: {ret.status_code}")
-    assert ret.status_code == 200
-    ret = requests.post(args.controller_url + "/list_models")
-    logger.info(f"status_code from `list_models`: {ret.status_code}")
-    models = ret.json()["models"]
-    models = sort_models(models)
-
-    logger.info(f"Models (from {args.controller_url}): {models}")
-    return models
-
-
 get_window_url_params = """
 function() {
     const params = new URLSearchParams(window.location.search);
@@ -154,48 +120,6 @@ def find_bounding_boxes(state, response):
     return returned_image if len(matches) > 0 else None
 
 
-def query_image_generation(response, sd_worker_url, timeout=15):
-    if not sd_worker_url:
-        return None
-    sd_worker_url = f"{sd_worker_url}/generate_image/"
-    pattern = r"```drawing-instruction\n(.*?)\n```"
-    match = re.search(pattern, response, re.DOTALL)
-    if match:
-        payload = {"caption": match.group(1)}
-        print("drawing-instruction:", payload)
-        response = requests.post(sd_worker_url, json=payload, timeout=timeout)
-        response.raise_for_status()  # 检查HTTP请求是否成功
-        image = Image.open(BytesIO(response.content))
-        return image
-    else:
-        return None
-
-
-def load_demo(url_params, request: gr.Request = None):
-    if not request:
-        logger.info(f"load_demo. ip: {request.client.host}. params: {url_params}")
-
-    dropdown_update = gr.Dropdown(visible=True)
-    if "model" in url_params:
-        model = url_params["model"]
-        if model in models:
-            dropdown_update = gr.Dropdown(value=model, visible=True)
-
-    state = init_state()
-    return state, dropdown_update
-
-
-def load_demo_refresh_model_list(request: gr.Request = None):
-    if not request:
-        logger.info(f"load_demo. ip: {request.client.host}")
-    models = get_model_list()
-    state = init_state()
-    dropdown_update = gr.Dropdown(
-        choices=models, value=models[0] if len(models) > 0 else ""
-    )
-    return state, dropdown_update
-
-
 def vote_last_response(state, liked, model_selector, request: gr.Request):
     conv_data = {
         "tstamp": round(time.time(), 4),
@@ -249,7 +173,7 @@ def flag_last_response(state, model_selector, request: gr.Request):
 def regenerate(state, image_process_mode, request: gr.Request):
     logger.info(f"regenerate. ip: {request.client.host}")
     # state.messages[-1][-1] = None
-    state.update_message(Conversation.ASSISTANT, None, -1)
+    state.update_message(Conversation.ASSISTANT, content='', image=None, idx=-1)
     prev_human_msg = state.messages[-2]
     if type(prev_human_msg[1]) in (tuple, list):
         prev_human_msg[1] = (*prev_human_msg[1][:2], image_process_mode)
@@ -327,16 +251,11 @@ def http_bot(
         ) + (no_change_btn,) * 5
         return
 
-    # Query worker address
-    controller_url = args.controller_url
-    ret = requests.post(
-        controller_url + "/get_worker_address", json={"model": model_name}
-    )
-    worker_addr = ret.json()["address"]
-    if worker_addr.startswith("http://0.0.0.0") and args.worker_ip:
-        worker_addr = worker_addr.replace("0.0.0.0", args.worker_ip)
-    logger.info(f"model_name: {model_name}, worker_addr: {worker_addr}")
+    worker_addr = os.environ.get("WORKER_ADDR", "")
+    api_token = os.environ.get("API_TOKEN", "")
+    headers = {"Authorization": f"{api_token}", "Content-Type": "application/json"}
 
+    state.append_message(Conversation.ASSISTANT, state.streaming_placeholder)
     # No available worker
     if worker_addr == "":
         # state.messages[-1][-1] = server_error_msg
@@ -359,20 +278,14 @@ def http_bot(
     # Make requests
     pload = {
         "model": model_name,
-        "prompt": state.get_prompt(),
+        "messages": state.get_prompt_v2(inlude_image=True, max_dynamic_patch=max_input_tiles),
         "temperature": float(temperature),
         "top_p": float(top_p),
-        "max_new_tokens": max_new_tokens,
-        "max_input_tiles": max_input_tiles,
-        # "bbox_threshold": bbox_threshold,
-        # "mask_threshold": mask_threshold,
+        "max_tokens": max_new_tokens,
         "repetition_penalty": repetition_penalty,
-        "images": f"List of {len(all_images)} images: {all_image_paths}",
+        "stream": True
     }
     logger.info(f"==== request ====\n{pload}")
-    pload.pop("images")
-    pload["prompt"] = state.get_prompt(inlude_image=True)
-    state.append_message(Conversation.ASSISTANT, state.streaming_placeholder)
     yield (
         state,
         state.to_gradio_chatbot(),
@@ -381,50 +294,25 @@ def http_bot(
 
     try:
         # Stream output
-        response = requests.post(
-            worker_addr + "/worker_generate_stream",
-            headers=headers,
-            json=pload,
-            stream=True,
-            timeout=20,
-        )
-        for chunk in response.iter_lines(decode_unicode=False, delimiter=b"\0"):
+        response = requests.post(worker_addr, json=pload, headers=headers, stream=True)
+        finnal_output = ''
+        for chunk in response.iter_lines(decode_unicode=False, delimiter=b"\n"):
             if chunk:
-                data = json.loads(chunk.decode())
-                if data["error_code"] == 0:
-                    if "text" in data:
-                        output = data["text"].strip()
-                        output += state.streaming_placeholder
-
-                    image = None
-                    if "image" in data:
-                        image = load_image_from_base64(data["image"])
-                        _ = state.save_image(image)
-
-                    state.update_message(Conversation.ASSISTANT, output, image)
-                    yield (
-                        state,
-                        state.to_gradio_chatbot(),
-                        gr.MultimodalTextbox(interactive=False),
-                    ) + (disable_btn,) * 5
-                else:
-                    output = (
-                        f"**{data['text']}**" + f" (error_code: {data['error_code']})"
-                    )
-
-                    state.update_message(Conversation.ASSISTANT, output, None)
-                    yield (
-                        state,
-                        state.to_gradio_chatbot(),
-                        gr.MultimodalTextbox(interactive=True),
-                    ) + (
-                        disable_btn,
-                        disable_btn,
-                        disable_btn,
-                        enable_btn,
-                        enable_btn,
-                    )
-                    return
+                chunk = chunk.decode()
+                if chunk == 'data: [DONE]':
+                    break
+                if chunk.startswith("data:"):
+                    chunk = chunk[5:]
+                    chunk = json.loads(chunk)
+                    output = chunk['choices'][0]['delta']['content']
+                    finnal_output += output
+                
+                state.update_message(Conversation.ASSISTANT, finnal_output + state.streaming_placeholder, None)
+                yield (
+                    state,
+                    state.to_gradio_chatbot(),
+                    gr.MultimodalTextbox(interactive=False),
+                ) + (disable_btn,) * 5
     except requests.exceptions.RequestException as e:
         state.update_message(Conversation.ASSISTANT, server_error_msg, None)
         yield (
@@ -445,12 +333,6 @@ def http_bot(
         returned_image = find_bounding_boxes(state, ai_response)
         returned_image = [returned_image] if returned_image else []
         state.update_message(Conversation.ASSISTANT, ai_response, returned_image)
-    if "```drawing-instruction" in ai_response:
-        returned_image = query_image_generation(
-            ai_response, sd_worker_url=sd_worker_url
-        )
-        returned_image = [returned_image] if returned_image else []
-        state.update_message(Conversation.ASSISTANT, ai_response, returned_image)
 
     state.end_of_current_turn()
 
@@ -577,7 +459,7 @@ def build_demo(embed_mode):
         theme=gr.themes.Default(),
         css=block_css,
     ) as demo:
-        models = get_model_list()
+        models = ['InternVL2-Pro']
         state = gr.State()
 
         if not embed_mode:
@@ -797,27 +679,6 @@ def build_demo(embed_mode):
             [state, chatbot, textbox] + btn_list,
         )
 
-        # NOTE: The following code will be not triggered when deployed on HF space.
-        # It's very strange. I don't know why.
-        """
-        if args.model_list_mode == "once":
-            demo.load(
-                load_demo,
-                [url_params],
-                [state, model_selector],
-                js=js,
-            )
-        elif args.model_list_mode == "reload":
-            demo.load(
-                load_demo_refresh_model_list,
-                None,
-                [state, model_selector],
-                js=js,
-            )
-        else:
-            raise ValueError(f"Unknown model list mode: {args.model_list_mode}")
-        """
-
     return demo
 
 
@@ -825,31 +686,13 @@ if __name__ == "__main__":
     parser = argparse.ArgumentParser()
     parser.add_argument("--host", type=str, default="0.0.0.0")
     parser.add_argument("--port", type=int, default=7860)
-    parser.add_argument("--controller-url", type=str, default=None)
-    parser.add_argument("--worker-ip", type=str, default=None)
     parser.add_argument("--concurrency-count", type=int, default=10)
-    parser.add_argument(
-        "--model-list-mode", type=str, default="reload", choices=["once", "reload"]
-    )
-    parser.add_argument("--sd-worker-url", type=str, default=None)
     parser.add_argument("--share", action="store_true")
     parser.add_argument("--moderate", action="store_true")
     parser.add_argument("--embed", action="store_true")
     args = parser.parse_args()
     logger.info(f"args: {args}")
-    if not args.controller_url:
-        args.controller_url = os.environ.get("CONTROLLER_URL", None)
-
-    if not args.controller_url:
-        raise ValueError("controller-url is required.")
-
-    if not args.worker_ip:
-        args.worker_ip = os.environ.get("WORKER_IP", None)
 
-    model_lists = ["OpenGVLab/InternVL-Chat-V1-5", "OpenGVLab/InternVL2-1B", "OpenGVLab/InternVL2-2B",
-                   "OpenGVLab/InternVL2-4B", "OpenGVLab/InternVL2-8B", "OpenGVLab/InternVL2-26B", 
-                   "OpenGVLab/InternVL2-40B", "OpenGVLab/InternVL2-Llama3-76B"]
-    sd_worker_url = args.sd_worker_url
     logger.info(args)
     demo = build_demo(args.embed)
     demo.queue(api_open=False).launch(