添加 qwen edits 暂未支持

Author: shell-nlp

gpt_server/database/models/process_manager.py

@@ -0,0 +1,104 @@
+"""暂时没有使用此代码"""
+
+from typing import List, Dict, Optional, Any
+from multiprocessing import Process
+from sqlmodel import SQLModel, Field, create_engine, Session, select
+from datetime import datetime
+import json
+from uuid import uuid4
+
+
+# 数据库模型
+class ProcessRecord(SQLModel, table=True):
+    id: int | None = Field(default=None, primary_key=True, description="主键ID")
+    pid: int | None = Field(default=None, description="进程ID")
+    args: str = Field(default="", description="进程参数")
+    status: str = Field(
+        default="created", description="进程状态"
+    )  # created, started, stopped
+    created_at: datetime = Field(default_factory=datetime.now, description="创建时间")
+    started_at: Optional[datetime] = Field(default=None, description="启动时间")
+    stopped_at: Optional[datetime] = Field(default=None, description="停止时间")
+
+
+class ProcessManager:
+    def __init__(self, write_db: bool = False, db_url: str = "sqlite:///processes.db"):
+        """进程管理类
+
+        Parameters
+        ----------
+        write_db : bool, optional
+            是否将进程信息写入到数据库, by default False
+        db_url : str, optional
+            数据库的连接 url, by default "sqlite:///processes.db"
+        """
+        self.processes: List[Dict[Process, dict]] | None = []
+        self.write_db = write_db
+        if self.write_db:
+            self.engine = create_engine(db_url)
+            # 创建表
+            SQLModel.metadata.create_all(self.engine)
+
+    def add_process(
+        self,
+        target,
+        args=(),
+    ):
+        p = Process(target=target, args=args)
+        process_id = uuid4().int & ((1 << 64) - 1)
+        self.processes.append({p: {"args": args, "process_id": process_id}})
+        if self.write_db:
+            # 记录到数据库
+            with Session(self.engine) as session:
+
+                process_record = ProcessRecord(
+                    id=process_id,
+                    pid=None,
+                    args=json.dumps(args, ensure_ascii=False),
+                    status="created",
+                )
+                session.add(process_record)
+                session.commit()
+                session.refresh(process_record)
+
+    def start_all(self):
+        for process in self.processes:
+            for _process, process_info in process.items():
+                _process.start()
+                process_info["pid"] = _process.pid
+                if self.write_db:
+                    process_id = process_info["process_id"]
+                    # 更新数据库记录
+                    with Session(self.engine) as session:
+                        # 根据PID查找记录（这里简化处理，实际可能需要更好的标识）
+                        statement = select(ProcessRecord).where(
+                            ProcessRecord.id == process_id
+                        )
+                        result = session.exec(statement)
+                        process_record = result.first()
+                        if process_record:
+                            process_record.pid = _process.pid
+                            process_record.status = "started"
+                            process_record.started_at = datetime.now()
+                            session.add(process_record)
+                            session.commit()
+                            session.refresh(process_record)
+
+    def join_all(self):
+        for process in self.processes:
+            for _process, process_info in process.items():
+                _process.join()
+                if self.write_db:
+                    process_id = process_info["process_id"]
+                    # 更新数据库记录为完成状态
+                    with Session(self.engine) as session:
+                        statement = select(ProcessRecord).where(
+                            ProcessRecord.id == process_id
+                        )
+                        results = session.exec(statement)
+                        record = results.first()
+                        if record:
+                            record.status = "finished"
+                            record.finished_at = datetime.now()
+                            session.add(record)
+                            session.commit()

gpt_server/model_worker/qwen_image_edit.py

@@ -0,0 +1,106 @@
+import asyncio
+
+import io
+import os
+from typing import List
+import uuid
+from loguru import logger
+import shortuuid
+from gpt_server.model_worker.base.model_worker_base import ModelWorkerBase
+from gpt_server.model_worker.utils import (
+    pil_to_base64,
+    load_base64_or_url,
+    bytesio2image,
+)
+from gpt_server.utils import STATIC_DIR
+import torch
+from diffusers import QwenImageEditPipeline
+
+root_dir = os.path.dirname(os.path.dirname(os.path.dirname(__file__)))
+
+
+class QwenImageEditWorker(ModelWorkerBase):
+    def __init__(
+        self,
+        controller_addr: str,
+        worker_addr: str,
+        worker_id: str,
+        model_path: str,
+        model_names: List[str],
+        limit_worker_concurrency: int,
+        conv_template: str = None,  # type: ignore
+    ):
+        super().__init__(
+            controller_addr,
+            worker_addr,
+            worker_id,
+            model_path,
+            model_names,
+            limit_worker_concurrency,
+            conv_template,
+            model_type="image",
+        )
+        backend = os.environ["backend"]
+        self.device = "cuda" if torch.cuda.is_available() else "cpu"
+        self.pipe = QwenImageEditPipeline.from_pretrained(model_path)
+        self.pipe.to(torch.bfloat16)
+        self.pipe.to(self.device)
+        self.pipe.set_progress_bar_config(disable=None)
+        logger.warning(f"模型：{model_names[0]}")
+
+    async def get_image_output(self, params):
+        prompt = params["prompt"]
+        response_format = params.get("response_format", "b64_json")
+        bytes_io = await load_base64_or_url(params["image"])
+        image = bytesio2image(bytes_io)
+        inputs = {
+            "image": image,
+            "prompt": prompt,
+            "negative_prompt": None,
+            "generator": torch.manual_seed(0),
+            "true_cfg_scale": 4.0,
+            "negative_prompt": " ",
+            "num_inference_steps": 50,
+        }
+        with torch.inference_mode():
+            output = self.pipe(**inputs)
+            image = output.images[0]
+
+        result = {}
+        if response_format == "b64_json":
+            # Convert PIL image to base64
+            base64 = pil_to_base64(pil_img=image)
+            result = {
+                "created": shortuuid.random(),
+                "data": [{"b64_json": base64}],
+                "usage": {
+                    "total_tokens": 0,
+                    "input_tokens": 0,
+                    "output_tokens": 0,
+                    "input_tokens_details": {"text_tokens": 0, "image_tokens": 0},
+                },
+            }
+            return result
+        elif response_format == "url":
+            # 生成唯一文件名（避免冲突）
+            file_name = str(uuid.uuid4()) + ".png"
+            save_path = STATIC_DIR / file_name
+            image.save(save_path, format="PNG")
+            WORKER_PORT = os.environ["WORKER_PORT"]
+            WORKER_HOST = os.environ["WORKER_HOST"]
+            url = f"http://{WORKER_HOST}:{WORKER_PORT}/static/{file_name}"
+            result = {
+                "created": shortuuid.random(),
+                "data": [{"url": url}],
+                "usage": {
+                    "total_tokens": 0,
+                    "input_tokens": 0,
+                    "output_tokens": 0,
+                    "input_tokens_details": {"text_tokens": 0, "image_tokens": 0},
+                },
+            }
+        return result
+
+
+if __name__ == "__main__":
+    QwenImageEditWorker.run()

gpt_server/model_worker/utils.py

@@ -4,7 +4,7 @@
 import base64
 import io
 import os
-from PIL.Image import Image
+from PIL import Image
 import re
 
 
@@ -14,38 +14,47 @@ def is_base64_image(data_string):
 
 
 # 转换为Base64
-def pil_to_base64(pil_img: Image, format: str = "PNG"):
+def pil_to_base64(pil_img: Image.Image, format: str = "PNG"):
     buffered = io.BytesIO()
     pil_img.save(buffered, format=format)  # 明确指定PNG格式
     return base64.b64encode(buffered.getvalue()).decode("utf-8")
 
 
-def extract_base64(data_url: str):
+def _extract_base64(data_url: str):
     """从Data URL中提取纯Base64数据"""
     return data_url.split(",", 1)[-1]  # 从第一个逗号后分割
 
 
-async def get_bytes_from_url(url: str) -> bytes:
+async def _get_bytes_from_url(url: str) -> bytes:
     async with httpx.AsyncClient() as client:
         response = await client.get(url)
         if response.status_code != 200:
             raise HTTPException(status_code=400, detail="无法从指定 URL 下载数据")
         return response.content
 
 
-async def load_base64_or_url(base64_or_url):
+def bytesio2image(bytes_io: io.BytesIO) -> Image.Image:
+    return Image.open(bytes_io)
+
+
+def bytes2image(bytes_: bytes) -> Image.Image:
+    bytes_io = io.BytesIO(bytes_)
+    return Image.open(bytes_io)
+
+
+async def load_base64_or_url(base64_or_url) -> io.BytesIO:
     # 根据 reference_audio 内容判断读取方式
     if base64_or_url.startswith("http://") or base64_or_url.startswith("https://"):
-        audio_bytes = await get_bytes_from_url(base64_or_url)
+        audio_bytes = await _get_bytes_from_url(base64_or_url)
     else:
         try:
             if "data:" in base64_or_url:
-                base64_or_url = extract_base64(data_url=base64_or_url)
+                base64_or_url = _extract_base64(data_url=base64_or_url)
             audio_bytes = base64.b64decode(base64_or_url)
         except Exception as e:
             logger.warning("无效的 base64 数据: " + str(e))
             raise HTTPException(status_code=400, detail="无效的 base64 数据: " + str(e))
-    # 利用 BytesIO 包装字节数据，然后使用 soundfile 读取为 numpy 数组
+    # 利用 BytesIO 包装字节数据
     try:
         bytes_io = io.BytesIO(audio_bytes)
     except Exception as e:

gpt_server/serving/openai_api_server.py

@@ -714,9 +714,37 @@ async def generate_completion(payload: Dict[str, Any], worker_addr: str):
     SpeechRequest,
     OpenAISpeechRequest,
     ImagesGenRequest,
+    ImagesEditsRequest,
 )
 
 
+async def get_images_edits(payload: Dict[str, Any]):
+    model_name = payload["model"]
+    worker_addr = get_worker_address(model_name)
+
+    transcription = await fetch_remote(
+        worker_addr + "/worker_get_image_output", payload
+    )
+    return json.loads(transcription)
+
+
+@app.post("/v1/images/edits", dependencies=[Depends(check_api_key)])
+async def images_edits(request: ImagesEditsRequest):
+    """图片编辑"""
+    error_check_ret = check_model(request)
+    if error_check_ret is not None:
+        return error_check_ret
+    payload = {
+        "image": request.image,
+        "model": request.model,
+        "prompt": request.prompt,
+        "output_format": request.output_format,
+        "response_format": request.response_format,
+    }
+    result = await get_images_edits(payload=payload)
+    return result
+
+
 async def get_images_gen(payload: Dict[str, Any]):
     model_name = payload["model"]
     worker_addr = get_worker_address(model_name)
@@ -728,7 +756,8 @@ async def get_images_gen(payload: Dict[str, Any]):
 
 
 @app.post("/v1/images/generations", dependencies=[Depends(check_api_key)])
-async def speech(request: ImagesGenRequest):
+async def images_generations(request: ImagesGenRequest):
+    """文生图"""
     error_check_ret = check_model(request)
     if error_check_ret is not None:
         return error_check_ret