优化图片编辑

Author: shell-nlp

gpt_server/model_worker/qwen_image_edit.py

@@ -1,6 +1,5 @@
 import asyncio
 
-import io
 import os
 from typing import List
 import uuid
@@ -14,7 +13,7 @@
 )
 from gpt_server.utils import STATIC_DIR
 import torch
-from diffusers import QwenImageEditPipeline
+from diffusers import QwenImageEditPlusPipeline
 
 root_dir = os.path.dirname(os.path.dirname(os.path.dirname(__file__)))
 
@@ -40,9 +39,8 @@ def __init__(
             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 = QwenImageEditPlusPipeline.from_pretrained(model_path)
         self.pipe.to(torch.bfloat16)
         self.pipe.to(self.device)
         self.pipe.set_progress_bar_config(disable=None)
@@ -51,16 +49,18 @@ def __init__(
     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)
+        image: list = params["image"]
+        image = [bytesio2image(await load_base64_or_url(img)) for img in image]
+        # 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,
+            "num_inference_steps": 40,
         }
         with torch.inference_mode():
             output = await asyncio.to_thread(self.pipe, **inputs)

gpt_server/serving/openai_api_server.py

@@ -1229,7 +1229,9 @@ async def get_images_edits(payload: Dict[str, Any]):
 @app.post("/v1/images/edits", dependencies=[Depends(check_api_key)])
 async def images_edits(
     model: str = Form(...),
-    image: UploadFile = File(media_type="application/octet-stream"),
+    image: Union[UploadFile, List[UploadFile]] = File(
+        ..., media_type="application/octet-stream"
+    ),
     prompt: Optional[Union[str, List[str]]] = Form(None),
     # negative_prompt: Optional[Union[str, List[str]]] = Form(None),
     response_format: Optional[str] = Form("url"),
@@ -1240,10 +1242,14 @@ async def images_edits(
     error_check_ret = check_model(model)
     if error_check_ret is not None:
         return error_check_ret
+    images = None
+    if not isinstance(image, list):  # 单
+        images = [image]
+    else:
+        images = image
+    image = [base64.b64encode(await img.read()).decode("utf-8") for img in images]
     payload = {
-        "image": base64.b64encode(await image.read()).decode(
-            "utf-8"
-        ),  # bytes → Base64 字符串,
+        "image": image,  # bytes → Base64 字符串,
         "model": model,
         "prompt": prompt,
         "output_format": output_format,