1

app/litserve_serve.py

@@ -1,10 +1,10 @@
+import logging
 import uuid
 
 import torch
 from minio import Minio
 
 import litserve as ls
-from PIL import Image
 import io
 from diffusers import Flux2KleinPipeline
 
@@ -13,21 +13,6 @@ from app.utils.new_oss_client import oss_get_image, oss_upload_image, MINIO_URL,
 minio_client = Minio(MINIO_URL, access_key=MINIO_ACCESS, secret_key=MINIO_SECRET, secure=MINIO_SECURE)
 
 
-# 保持原有的辅助函数
-def aspect_to_wh(aspect_ratio: str, base_long_edge: int) -> tuple[int, int]:
-    w_str, h_str = aspect_ratio.split(":")
-    w, h = float(w_str), float(h_str)
-    if w >= h:
-        width = base_long_edge
-        height = int(round(base_long_edge * (h / w)))
-    else:
-        height = base_long_edge
-        width = int(round(base_long_edge * (w / h)))
-    width = max(64, (width // 8) * 8)
-    height = max(64, (height // 8) * 8)
-    return width, height
-
-
 class FluxKleinAPI(ls.LitAPI):
     def setup(self, device):
         # 1. 模型初始化
@@ -42,7 +27,7 @@ class FluxKleinAPI(ls.LitAPI):
         )
         self.pipe.to(device)
 
-    def decode_request(self, request):
+    async def decode_request(self, request):
         """
         解析请求参数并加载OSS图片的接口函数
 
@@ -67,8 +52,6 @@ class FluxKleinAPI(ls.LitAPI):
               推理步数，控制模型生成过程的迭代次数
             - guidance : float (可选，默认值4.0)
               引导系数，调节提示词对生成结果的影响程度
-            - seed : int (可选，默认值42)
-              随机种子，保证生成结果的可复现性
 
         返回值说明
         -------
@@ -80,7 +63,6 @@ class FluxKleinAPI(ls.LitAPI):
             - prompt: 文本提示词（默认空字符串）
             - steps: 推理步数（默认28）
             - guidance: 引导系数（默认4.0）
-            - seed: 随机种子（默认42）
             - height: 图片高度（默认512）
             - width: 图片宽度（默认512）
 
@@ -106,19 +88,15 @@ class FluxKleinAPI(ls.LitAPI):
             "prompt": request.get("prompt", ""),
             "steps": request.get("steps", 4),
             "guidance": request.get("guidance", 4.0),
-            "seed": request.get("seed", 42),
             "height": H,
             "width": W
         }
 
-    @torch.inference_mode()
-    def predict(self, payload):
+    async def predict(self, payload):
         # 3. 执行推理逻辑
         images = payload.get("images", [])
         prompt = payload.get("prompt", "")
         gen = torch.Generator(device=self.device)
-        seed = gen.seed()
-        print(f"本次使用的随机种子是: {seed}")
         if images:
             output = self.pipe(
                 image=images,
@@ -144,14 +122,15 @@ class FluxKleinAPI(ls.LitAPI):
         image_bytes = image_data.read()
         req = oss_upload_image(oss_client=minio_client, bucket=payload.get("bucket_name", "test"), object_name=payload.get("object_name", f"fida_generate_image/{uuid.uuid4().hex}.png"), image_bytes=image_bytes)
         output_path = req.bucket_name + "/" + req.object_name
+        logging.info(f"output_path :{output_path}")
         return output_path
 
-    def encode_response(self, output_path):
+    async def encode_response(self, output_path):
         return {"output_path": output_path}
 
 
 if __name__ == "__main__":
     # 启动服务器
-    api = FluxKleinAPI()
+    api = FluxKleinAPI(enable_async=True)
     server = ls.LitServer(api, accelerator="cuda", devices=1)
     server.run(port=8451)