Upload app.py

app.py

@@ -6,7 +6,7 @@ import asyncio
 import time
 import gc
 from datetime import datetime
-from collections import deque  # 引入双端队列
+from collections import deque
 from PIL import Image
 from io import BytesIO
 from openai import AsyncOpenAI
@@ -21,28 +21,25 @@ import uvicorn
 # ==========================================
 CONCURRENCY_LIMITER = asyncio.Semaphore(3)
 
-# --- 🚀 动态违禁词缓存 (核心修改) ---
-# maxlen=200: 自动保持最新的200个，旧的自动丢弃
 DYNAMIC_BLOCK_CACHE = deque(maxlen=200)
 
-# 硬编码的基础高危词 (保留最基本的底线，防止API挂了时裸奔)
 SENSITIVE_WORDS = [
     "nsfw", "nude", "naked", "sex", "porn", "hentai", 
     "裸", "色情", "全裸", "blood", "kill", "murder", 
     "血腥", "尸体", "毒品", "杀"
 ]
 
-AR_BASES = {
-    "1:1 (SDXL标准 1024)": (1024, 1024),
-    "16:9 (SDXL标准 768p)": (1344, 768),
-    "16:9 (1080p Full HD)": (1920, 1080),
-    "9:16 (手机竖屏)": (768, 1344),
-    "4:3 (标准)":  (1152, 896),
-    "21:9 (宽屏)": (1536, 640)
+# 精确的枚举表
+RESOLUTION_MAP = {
+    "1:1":  {"1k": (1024, 1024), "2k": (2048, 2048), "4k": (4096, 4096)},
+    "16:9": {"1k": (1280, 720),  "2k": (2560, 1440), "4k": (3840, 2160)},
+    "9:16": {"1k": (720, 1280),  "2k": (1440, 2560), "4k": (2160, 3840)},
+    "4:3":  {"1k": (1152, 864),  "2k": (2048, 1536), "4k": (2880, 2160)},
+    "3:4":  {"1k": (864, 1152),  "2k": (1536, 2048), "4k": (2160, 2880)},
+    "21:9": {"1k": (1680, 720),  "2k": (3360, 1440), "4k": (5040, 2160)},
+    "2:3":  {"1k": (852, 1280),  "2k": (1704, 2560), "4k": (2440, 3660)}
 }
-SCALE_MULTIPLIERS = {"1k": 1.0, "2k": 1.5, "4k": 2.0}
 
-# IP 限流记录
 IP_REQUEST_HISTORY = {}
 
 # ==========================================
@@ -68,25 +65,14 @@ def check_ip_rate_limit(ip: str, limit: int = 5, window: int = 60):
     return False
 
 def contains_local_sensitive_words(text: str):
-    """
-    检查本地缓存
-    1. 检查最近被百度封杀的200个Prompt (精确匹配)
-    2. 检查硬编码的关键词 (模糊匹配)
-    """
-    # 1. 检查动态缓存 (O(n) 遍历，但n最大200，极快)
-    if text in DYNAMIC_BLOCK_CACHE:
-        return True, "最近违规记录 (已缓存)"
-
-    # 2. 检查静态关键词
+    if text in DYNAMIC_BLOCK_CACHE: return True, text
     text_lower = text.lower()
     for word in SENSITIVE_WORDS:
-        if word and word.lower() in text_lower:
-            return True, word
-            
+        if word and word.lower() in text_lower: return True, word
     return False, None
 
 # ==========================================
-# 3. 百度审核 & 生成
+# 3. 百度审核 & 生成逻辑
 # ==========================================
 
 _BAIDU_TOKEN_CACHE = {"token": None, "expires_at": 0}
@@ -104,8 +90,7 @@ async def get_baidu_access_token():
     try:
         async with httpx.AsyncClient(timeout=10) as client:
             resp = await client.post(url, params=params)
-            data = resp.json()
-            token = data.get("access_token")
+            token = resp.json().get("access_token")
             if token:
                 _BAIDU_TOKEN_CACHE["token"] = token
                 _BAIDU_TOKEN_CACHE["expires_at"] = now + 1728000
@@ -124,16 +109,12 @@ async def check_baidu_text_censor(text):
         async with httpx.AsyncClient(timeout=10) as client:
             resp = await client.post(url, data=data, headers=headers)
             result = resp.json()
-            
             if result.get('conclusionType') == 2:
                 msg = [i.get('msg') for i in result.get('data', [])]
                 reason = ",".join(msg) if msg else "违规"
-                
-                # 🚀 发现新词：存入内存队列
                 if text not in DYNAMIC_BLOCK_CACHE:
                     DYNAMIC_BLOCK_CACHE.append(text)
-                    print(f"🔒 [Cache] 缓存违规Prompt ({len(DYNAMIC_BLOCK_CACHE)}/200): {text[:15]}...")
-                
+                    print(f"🔒 [Cache] 新增违规: {text[:10]}...")
                 return True, reason
             elif result.get('conclusionType') == 3: return True, "疑似违规"
             return False, None
@@ -145,60 +126,96 @@ class GenerateRequest(BaseModel):
     prompt: str
     ratio: str
     scale: str
-    custom_size: str
-    format: str
+    format: str 
+    # ⚠️ custom_size 已移除
 
-def process_image_in_memory(image_bytes: bytes, target_format: str) -> str:
+def process_image_in_memory(image_bytes: bytes, target_format: str, target_size: tuple = None) -> str:
     try:
         with Image.open(BytesIO(image_bytes)) as img:
             out = BytesIO()
+            final_img = img
+            status_msg = f"[原图直出] 尺寸: {img.size}"
+
+            if target_size:
+                w_diff = abs(img.size[0] - target_size[0])
+                h_diff = abs(img.size[1] - target_size[1])
+                if w_diff > 50 or h_diff > 50:
+                    print(f"⚠️ [Resize] {img.size} -> {target_size}")
+                    final_img = img.resize(target_size, Image.Resampling.LANCZOS)
+                    status_msg = f"[触发放大] -> {target_size}"
+                else:
+                    status_msg = f"[原图直出] 符合预期"
+            print(f"✅ {status_msg}")
+
             fmt = "JPEG" if target_format == "jpg" else "PNG"
-            img = img.convert("RGB") if target_format == "jpg" else img
-            img.save(out, format=fmt, quality=95)
+            final_img = final_img.convert("RGB") if target_format == "jpg" else final_img
+            final_img.save(out, format=fmt, quality=95)
             return base64.b64encode(out.getvalue()).decode('utf-8')
-    except Exception as e: raise Exception(f"IMG Error: {e}")
-    finally: gc.collect()
+    except Exception as e: 
+        raise Exception(f"Image Error: {e}")
+    finally: 
+        gc.collect()
 
 async def download_image_as_bytes(url):
     print(f"📥 下载: {url}")
-    async with httpx.AsyncClient(timeout=60) as client:
+    async with httpx.AsyncClient(timeout=180) as client:
         r = await client.get(url)
         if r.status_code == 200: return r.content
         raise Exception(f"HTTP {r.status_code}")
 
 async def core_generate(req: GenerateRequest):
-    fs = "1024x1024"
-    if req.ratio == "custom": fs = req.custom_size.strip()
-    else:
-        w, h = AR_BASES.get(req.ratio, (1024, 1024))
-        m = SCALE_MULTIPLIERS.get(req.scale, 1.0)
-        fs = f"{int(w*m)}x{int(h*m)}"
+    # 1. 纯查表逻辑 (不再处理 custom)
+    ratio_key = req.ratio.split(' ')[0] # 提取 "16:9"
+    
+    # 默认兜底
+    target_w, target_h = 1024, 1024
     
-    fp = req.prompt
+    if ratio_key in RESOLUTION_MAP:
+        if req.scale in RESOLUTION_MAP[ratio_key]:
+            target_w, target_h = RESOLUTION_MAP[ratio_key][req.scale]
+        else:
+            # 如果 scale 不在表里，默认 1k
+            target_w, target_h = RESOLUTION_MAP[ratio_key]["1k"]
+            
+    api_size_str = f"{target_w}x{target_h}"
+    
+    # 2. Prompt 处理
+    final_prompt = req.prompt
     if "gemini" in req.model.lower():
-        suffix = f"--resolution {fs}" if req.ratio=="custom" else f"--ar {req.ratio.split()[0]} --resolution {fs}"
-        fp = f"{req.prompt} {suffix}"
-
+        quality_prompt = "standard quality"
+        if req.scale == "4k": quality_prompt = "Extreme High Quality, 4K Resolution"
+        elif req.scale == "2k": quality_prompt = "High Quality, 2K Resolution"
+        
+        suffix = f". (Settings: Aspect Ratio {ratio_key}, Quality {quality_prompt}, Target Size {api_size_str})"
+        final_prompt = f"{req.prompt} {suffix}"
+    
+    # 3. API 请求
     client = AsyncOpenAI(api_key=req.api_key, base_url="https://api.cloudapp.ink/v1", max_retries=0, timeout=60.0)
-
     img_b = None
+
     if "gemini" in req.model.lower():
         print(f"🟣 Gemini: {req.model}")
-        res = await client.chat.completions.create(model=req.model, messages=[{"role":"user","content":fp}], extra_body={"modalities":["image","text"]})
+        res = await client.chat.completions.create(
+            model=req.model, 
+            messages=[{"role":"user","content":final_prompt}], 
+            extra_body={"modalities":["image","text"]}
+        )
         c = res.choices[0].message.content
         u = re.search(r"!\[.*?\]\((https?://[^\)]+)\)", c) or re.search(r"(https?://\S+\.(?:png|jpg|jpeg|webp))", c)
         if u: img_b = await download_image_as_bytes(u.group(1))
         elif "base64," in c: img_b = base64.b64decode(c.split("base64,")[1].split(")")[0].strip())
-        else: raise Exception("No Image")
+        else: raise Exception("No Image Found")
     else:
-        print(f"🔵 Standard: {req.model}")
-        res = await client.images.generate(model=req.model, prompt=fp, n=1, size=fs, response_format="b64_json")
+        print(f"🔵 Standard: {req.model} | 请求 {api_size_str}")
+        res = await client.images.generate(
+            model=req.model, prompt=final_prompt, n=1, size=api_size_str, response_format="b64_json"
+        )
         d = res.data[0]
         if getattr(d,'b64_json',None): img_b = base64.b64decode(d.b64_json)
         elif hasattr(d,'url'): img_b = await download_image_as_bytes(d.url)
         else: raise Exception("No Data")
         
-    return process_image_in_memory(img_b, req.format)
+    return process_image_in_memory(img_b, req.format, target_size=(target_w, target_h))
 
 # ==========================================
 # 4. FastAPI App
@@ -215,19 +232,17 @@ def health(): return {"status": "ok"}
 @app.post("/api/generate")
 async def generate_api(req: GenerateRequest, request: Request):
     ip = get_real_ip(request)
-    print(f"\n🚀 [{datetime.now().strftime('%H:%M:%S')}] IP: {ip} | Prompt: {req.prompt[:20]}...")
+    print(f"\n🚀 [{datetime.now().strftime('%H:%M:%S')}] IP: {ip} | Target: {req.scale} ({req.ratio})")
 
-    if check_ip_rate_limit(ip, 5, 60):
+    if check_ip_rate_limit(ip, 3, 60):
         return JSONResponse({"status": "error", "message": "请求过快"}, 429)
 
     is_loc, word = contains_local_sensitive_words(req.prompt)
     if is_loc:
-        print(f"🚫 [Local] 拦截: {word}")
         return JSONResponse({"status": "error", "message": f"违规内容: {word}"}, 200)
 
     is_bd, reason = await check_baidu_text_censor(req.prompt)
     if is_bd:
-        print(f"🚫 [Baidu] 拦截: {reason}")
         return JSONResponse({"status": "error", "message": f"审核未通过: {reason}"}, 200)
 
     async with CONCURRENCY_LIMITER: