Refactor API and controller files; remove unused code and update imports

api.py

@@ -30,51 +30,4 @@ def get_selected_worker_ip(controller_url, selected_model):
 def pil_image_to_base64(image):
     buffered = BytesIO()
     image.save(buffered, format='PNG')
-    return base64.b64encode(buffered.getvalue()).decode('utf-8')
-
-
-controller_url = 'http://10.140.60.209:10075'
-model_list = get_model_list(controller_url)
-print(f'Model list: {model_list}')
-
-selected_model = 'InternVL2-1B'
-worker_addr = get_selected_worker_ip(controller_url, selected_model)
-print(f'model_name: {selected_model}, worker_addr: {worker_addr}')
-
-
-# 多轮/多图对话请把数据组织成以下格式：
-# send_messages = [{'role': 'system', 'content': system_message}]
-# send_messages.append({'role': 'user', 'content': 'question1 to image1', 'image': [pil_image_to_base64(image)]})
-# send_messages.append({'role': 'assistant', 'content': 'answer1'})
-# send_messages.append({'role': 'user', 'content': 'question2 to image2', 'image': [pil_image_to_base64(image)]})
-# send_messages.append({'role': 'assistant', 'content': 'answer2'})
-# send_messages.append({'role': 'user', 'content': 'question3 to image1 & 2', 'image': []})
-
-image = Image.open('image1.jpg')
-print(f'Loading image, size: {image.size}')
-system_message = """我是书生·万象，英文名是InternVL，是由上海人工智能实验室及多家合作单位联合开发的多模态大语言模型。人工智能实验室致力于原始技术创新，开源开放，共享共创，推动科技进步和产业发展。
-请尽可能详细地回答用户的问题。"""
-send_messages = [{'role': 'system', 'content': system_message}]
-send_messages.append({'role': 'user', 'content': 'describe this image in detail', 'image': [pil_image_to_base64(image)]})
-
-pload = {
-    'model': selected_model,
-    'prompt': send_messages,
-    'temperature': 0.8,
-    'top_p': 0.7,
-    'max_new_tokens': 2048,
-    'max_input_tiles': 12,
-    'repetition_penalty': 1.0,
-}
-headers = {'User-Agent': 'InternVL-Chat Client'}
-response = requests.post(worker_addr + '/worker_generate_stream',
-                         headers=headers, json=pload, stream=True, timeout=10)
-for chunk in response.iter_lines(decode_unicode=False, delimiter=b'\0'):
-    if chunk:
-        data = json.loads(chunk.decode())
-        if data['error_code'] == 0:
-            output = data['text'] # 这里是流式输出
-        else:
-            output = data['text'] + f" (error_code: {data['error_code']})"
-# 完整的输出
-print(output)
+    return base64.b64encode(buffered.getvalue()).decode('utf-8')

controller.py

@@ -15,7 +15,7 @@ import numpy as np
 import requests
 import uvicorn
 from fastapi import FastAPI, Request
-from fastapi.responses import StreamingResponse
+from starlette.responses import StreamingResponse
 from utils import build_logger, server_error_msg
 
 CONTROLLER_HEART_BEAT_EXPIRATION = 30

gradio_web_server.py

@@ -90,81 +90,6 @@ def init_state(state=None):
         del state
     return Conversation()
 
-
-def find_bounding_boxes(state, response):
-    pattern = re.compile(r"<ref>\s*(.*?)\s*</ref>\s*<box>\s*(\[\[.*?\]\])\s*</box>")
-    matches = pattern.findall(response)
-    results = []
-    for match in matches:
-        results.append((match[0], eval(match[1])))
-    returned_image = None
-    latest_image = state.get_images(source=state.USER)[-1]
-    returned_image = latest_image.copy()
-    width, height = returned_image.size
-    draw = ImageDraw.Draw(returned_image)
-    for result in results:
-        line_width = max(1, int(min(width, height) / 200))
-        random_color = (
-            random.randint(0, 128),
-            random.randint(0, 128),
-            random.randint(0, 128),
-        )
-        category_name, coordinates = result
-        coordinates = [
-            (
-                float(x[0]) / 1000,
-                float(x[1]) / 1000,
-                float(x[2]) / 1000,
-                float(x[3]) / 1000,
-            )
-            for x in coordinates
-        ]
-        coordinates = [
-            (
-                int(x[0] * width),
-                int(x[1] * height),
-                int(x[2] * width),
-                int(x[3] * height),
-            )
-            for x in coordinates
-        ]
-        for box in coordinates:
-            draw.rectangle(box, outline=random_color, width=line_width)
-            font = ImageFont.truetype("assets/SimHei.ttf", int(20 * line_width / 2))
-            text_size = font.getbbox(category_name)
-            text_width, text_height = (
-                text_size[2] - text_size[0],
-                text_size[3] - text_size[1],
-            )
-            text_position = (box[0], max(0, box[1] - text_height))
-            draw.rectangle(
-                [
-                    text_position,
-                    (text_position[0] + text_width, text_position[1] + text_height),
-                ],
-                fill=random_color,
-            )
-            draw.text(text_position, category_name, fill="white", font=font)
-    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}")
@@ -433,16 +358,6 @@ def http_bot(
         return
 
     ai_response = state.return_last_message()
-    if "<ref>" in ai_response:
-        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()
 
@@ -823,7 +738,6 @@ if __name__ == "__main__":
     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")
@@ -837,7 +751,6 @@ if __name__ == "__main__":
 
     models = get_model_list()
 
-    sd_worker_url = args.sd_worker_url
     logger.info(args)
     demo = build_demo(args.embed)
     demo.queue(api_open=False).launch(

sd_worker.py

@@ -1,58 +0,0 @@
-# --------------------------------------------------------
-# InternVL
-# Copyright (c) 2024 OpenGVLab
-# Licensed under The MIT License [see LICENSE for details]
-# --------------------------------------------------------
-
-from io import BytesIO
-
-import torch
-from diffusers import StableDiffusion3Pipeline
-from fastapi import FastAPI
-from fastapi.responses import Response
-from pydantic import BaseModel
-
-# Initialize pipeline
-pipe = StableDiffusion3Pipeline.from_pretrained('stabilityai/stable-diffusion-3-medium-diffusers',
-                                                torch_dtype=torch.float16)
-pipe = pipe.to('cuda')
-
-# Create a FastAPI application
-app = FastAPI()
-
-
-# Define the input data model
-class CaptionRequest(BaseModel):
-    caption: str
-
-
-# Defining API endpoints
-@app.post('/generate_image/')
-async def generate_image(request: CaptionRequest):
-    caption = request.caption
-    negative_prompt = 'blurry, low resolution, artifacts, unnatural, poorly drawn, bad anatomy, out of focus'
-    image = pipe(
-        caption,
-        negative_prompt=negative_prompt,
-        num_inference_steps=20,
-        guidance_scale=7.0
-    ).images[0]
-
-    # Converts an image to a byte stream
-    img_byte_arr = BytesIO()
-    image.save(img_byte_arr, format='PNG')
-    img_byte_arr = img_byte_arr.getvalue()
-
-    return Response(content=img_byte_arr, media_type='image/png')
-
-
-# Run the Uvicorn server
-if __name__ == '__main__':
-    import argparse
-
-    import uvicorn
-    parser = argparse.ArgumentParser()
-    parser.add_argument('--port', default=11005, type=int)
-    args = parser.parse_args()
-
-    uvicorn.run(app, host='0.0.0.0', port=args.port)