Upload app.py

app.py

@@ -4,6 +4,7 @@ import base64
 import os
 import asyncio
 import time
+import gc # 引入垃圾回收
 from datetime import datetime
 from PIL import Image
 from io import BytesIO
@@ -17,6 +18,13 @@ import uvicorn
 # ==========================================
 # 1. 配置与工具
 # ==========================================
+
+# --- 关键修改：全局并发限制 ---
+# 限制同时进行处理的请求数量为 3。
+# 这样即使有 100 个人同时点，同时也只有 3 个人在占用内存处理图片，其他人排队。
+# 既保护了内存，又不会报错。
+CONCURRENCY_LIMITER = asyncio.Semaphore(5)
+
 AR_BASES = {
     "1:1 (SDXL标准 1024)": (1024, 1024),
     "16:9 (SDXL标准 768p)": (1344, 768),
@@ -43,22 +51,31 @@ class GenerateRequest(BaseModel):
 
 def process_image_in_memory(image_bytes: bytes, target_format: str) -> str:
     try:
-        img = Image.open(BytesIO(image_bytes))
-        output_buffer = BytesIO()
-        
-        if target_format == "jpg":
-            img = img.convert("RGB")
-            img.save(output_buffer, format="JPEG", quality=95)
-        else:
-            img.save(output_buffer, format="PNG")
+        # 使用上下文管理器打开，确保及时关闭
+        with Image.open(BytesIO(image_bytes)) as img:
+            output_buffer = BytesIO()
+            
+            if target_format == "jpg":
+                img = img.convert("RGB")
+                img.save(output_buffer, format="JPEG", quality=95)
+            else:
+                img.save(output_buffer, format="PNG")
             
-        return base64.b64encode(output_buffer.getvalue()).decode('utf-8')
+            result = base64.b64encode(output_buffer.getvalue()).decode('utf-8')
+            
+            # 显式关闭 buffer
+            output_buffer.close()
+            return result
     except Exception as e:
         raise Exception(f"格式转换失败: {str(e)}")
+    finally:
+        # 强制垃圾回收，防止大图片占用内存不释放
+        gc.collect()
 
 async def download_image_as_bytes(url):
     print(f"📥 [System] 正在下载图片: {url}")
-    async with httpx.AsyncClient(timeout=45) as client:
+    # 增加超时时间，防止大图下载断开
+    async with httpx.AsyncClient(timeout=60) as client:
         response = await client.get(url)
         if response.status_code == 200:
             return response.content
@@ -82,7 +99,7 @@ async def core_generate(req: GenerateRequest):
         final_size = f"{int(base_w * multiplier)}x{int(base_h * multiplier)}"
         ratio_str = req.ratio.split(" ")[0]
 
-    # 2. 处理 Prompt (加入分辨率指令)
+    # 2. 处理 Prompt
     final_prompt = req.prompt
     if "gemini" in req.model.lower():
         if req.ratio == "custom":
@@ -108,7 +125,6 @@ async def core_generate(req: GenerateRequest):
             extra_body={"modalities": ["image", "text"]}
         )
         content = response.choices[0].message.content
-        print(f"🟣 [Chat模式] 收到回复 (长度: {len(content)})")
         
         # 提取链接或 Base64
         url_match = re.search(r"!\[.*?\]\((https?://[^\)]+)\)", content)
@@ -122,7 +138,6 @@ async def core_generate(req: GenerateRequest):
             b64_str = content.split("base64,")[1].split(")")[0].strip()
             image_bytes = base64.b64decode(b64_str)
         else:
-            print(f"❌ [Error] 无法识别: {content[:100]}")
             raise Exception("未能识别图片内容")
             
     else:
@@ -166,31 +181,31 @@ def health():
 
 @app.post("/api/generate")
 async def generate_api(req: GenerateRequest, request: Request):
-    # --- 日志记录核心部分 ---
     client_host = request.client.host
     timestamp = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
     
     print("\n" + "="*60)
     print(f"🚀 [New Request] {timestamp} from {client_host}")
-    print(f"📌 模型 (Model)  : {req.model}")
-    print(f"📐 规格 (Specs)  : Ratio={req.ratio} | Scale={req.scale}")
-    print(f"🎨 格式 (Format) : {req.format}")
+    print(f"📌 模型: {req.model} | 规格: {req.ratio} ({req.scale})")
     print("-" * 60)
-    print(f"💡 提示词 (Prompt):")
-    print(f"{req.prompt}")
-    print("="*60 + "\n")
-    # -----------------------
 
-    try:
-        start_time = time.time()
-        base64_img = await core_generate(req)
-        duration = time.time() - start_time
-        print(f"✅ [Success] 生成完成，耗时: {duration:.2f}s")
-        
-        return {"status": "success", "image_base64": base64_img}
-    except Exception as e:
-        print(f"❌ [Failed] 生成失败: {str(e)}")
-        return JSONResponse(
-            content={"status": "error", "message": str(e)},
-            status_code=200 
-        )
+    # --- 关键修改：使用 Semaphore 进行限流 ---
+    # 只有拿到“通行证”的请求才能进入下面的代码块执行
+    # 没拿到的会在这一行等待，直到别人执行完释放内存
+    async with CONCURRENCY_LIMITER:
+        try:
+            start_time = time.time()
+            base64_img = await core_generate(req)
+            duration = time.time() - start_time
+            print(f"✅ [Success] 耗时: {duration:.2f}s")
+            
+            return {"status": "success", "image_base64": base64_img}
+        except Exception as e:
+            print(f"❌ [Failed] 失败: {str(e)}")
+            return JSONResponse(
+                content={"status": "error", "message": str(e)},
+                status_code=200 
+            )
+        finally:
+            # 每次请求结束后，强制清理一次内存
+            gc.collect()