优化架构实现

Author: shell-nlp

gpt_server/model_worker/base/base_model_worker.py

@@ -6,7 +6,6 @@
 from fastapi.responses import StreamingResponse, JSONResponse
 import requests
 
-from fastchat.constants import WORKER_HEART_BEAT_INTERVAL
 from fastchat.conversation import Conversation
 from fastchat.utils import pretty_print_semaphore
 
@@ -19,11 +18,11 @@ def build_logger():
 
 worker = None
 logger = None
-
+WORKER_HEART_BEAT_INTERVAL = 6
 app = FastAPI()
 
 
-def heart_beat_worker(obj):
+def heart_beat_worker(obj: "BaseModelWorker"):
     while True:
         time.sleep(WORKER_HEART_BEAT_INTERVAL)
         obj.send_heart_beat()
@@ -96,7 +95,7 @@ def register_to_controller(self):
 
         url = self.controller_addr + "/register_worker"
         data = {
-            "worker_name": self.worker_addr,
+            "worker_addr": self.worker_addr,
             "check_heart_beat": True,
             "worker_status": self.get_status(),
             "multimodal": self.multimodal,
@@ -105,12 +104,6 @@ def register_to_controller(self):
         assert r.status_code == 200
 
     def send_heart_beat(self):
-        # logger.info(
-        #     f"Send heart beat. Models: {self.model_names}. "
-        #     f"Semaphore: {pretty_print_semaphore(self.semaphore)}. "
-        #     f"call_ct: {self.call_ct}. "
-        #     f"worker_id: {self.worker_id}. "
-        # )
 
         url = self.controller_addr + "/receive_heart_beat"
 
@@ -119,7 +112,7 @@ def send_heart_beat(self):
                 ret = requests.post(
                     url,
                     json={
-                        "worker_name": self.worker_addr,
+                        "worker_addr": self.worker_addr,
                         "queue_length": self.get_queue_length(),
                     },
                     timeout=5,

gpt_server/model_worker/qwen_image.py

@@ -1,3 +1,4 @@
+import asyncio
 import os
 from typing import List
 import uuid
@@ -65,21 +66,24 @@ def __init__(
         logger.warning(f"模型：{model_names[0]}")
 
     async def get_image_output(self, params):
+        self.call_ct += 1
         prompt = params["prompt"]
         if contains_chinese(prompt):
             prompt += positive_magic["zh"]
         else:
             prompt += positive_magic["en"]
         response_format = params.get("response_format", "b64_json")
-        image = self.pipe(
-            prompt,
-            negative_prompt=" ",
-            height=height,
-            width=width,
-            num_inference_steps=50,
-            true_cfg_scale=4.0,
-            generator=torch.Generator(self.device).manual_seed(0),
-        ).images[0]
+        inputs = {
+            "prompt": prompt,
+            "negative_prompt": " ",
+            "height": height,
+            "width": width,
+            "num_inference_steps": 50,
+            "true_cfg_scale": 4.0,
+            "generator": torch.Generator(self.device).manual_seed(0),
+        }
+        output = await asyncio.to_thread(self.pipe, **inputs)
+        image = output.images[0]
         result = {}
         if response_format == "b64_json":
             # Convert PIL image to base64

gpt_server/model_worker/qwen_image_edit.py

@@ -63,7 +63,7 @@ async def get_image_output(self, params):
             "num_inference_steps": 50,
         }
         with torch.inference_mode():
-            output = self.pipe(**inputs)
+            output = await asyncio.to_thread(self.pipe, **inputs)
             image = output.images[0]
 
         result = {}

gpt_server/openai_api_protocol/custom_api_protocol.py

@@ -380,6 +380,11 @@ class CustomModelCard(BaseModel):
     owned_by: str = "gpt_server"
 
 
+class ModelList(BaseModel):
+    object: str = "list"
+    data: List[CustomModelCard] = []
+
+
 class CustomEmbeddingsRequest(BaseModel):
     model: Optional[str] = None
     engine: Optional[str] = None