Upload app.py

app.py

@@ -1,212 +1,310 @@
-import httpx
-import re
-import base64
-import os
-import asyncio
-import time
-import gc # 引入垃圾回收
-from datetime import datetime
-from PIL import Image
-from io import BytesIO
-from openai import AsyncOpenAI
-from fastapi import FastAPI, HTTPException, Request
-from fastapi.staticfiles import StaticFiles
-from fastapi.responses import FileResponse, JSONResponse
-from pydantic import BaseModel
-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),
-    "16:9 (1080p Full HD)": (1920, 1080),
-    "9:16 (手机竖屏)": (768, 1344),
-    "4:3 (标准)":  (1152, 896),
-    "21:9 (宽屏)": (1536, 640)
-}
-
-SCALE_MULTIPLIERS = {
-    "1k": 1.0,
-    "2k": 1.5,
-    "4k": 2.0
-}
-
-class GenerateRequest(BaseModel):
-    api_key: str
-    model: str
-    prompt: str
-    ratio: str
-    scale: str
-    custom_size: str
-    format: str
-
-def process_image_in_memory(image_bytes: bytes, target_format: str) -> str:
-    try:
-        # 使用上下文管理器打开，确保及时关闭
-        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")
-            
-            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=60) as client:
-        response = await client.get(url)
-        if response.status_code == 200:
-            return response.content
-        else:
-            raise Exception(f"下载失败 HTTP {response.status_code}")
-
-# ==========================================
-# 2. 核心生成逻辑 (异步)
-# ==========================================
-
-async def core_generate(req: GenerateRequest):
-    # 1. 计算分辨率
-    final_size = "1024x1024"
-    ratio_str = "1:1"
-
-    if req.ratio == "custom":
-        final_size = req.custom_size.strip()
-    else:
-        base_w, base_h = AR_BASES.get(req.ratio, (1024, 1024))
-        multiplier = SCALE_MULTIPLIERS.get(req.scale, 1.0)
-        final_size = f"{int(base_w * multiplier)}x{int(base_h * multiplier)}"
-        ratio_str = req.ratio.split(" ")[0]
-
-    # 2. 处理 Prompt
-    final_prompt = req.prompt
-    if "gemini" in req.model.lower():
-        if req.ratio == "custom":
-             final_prompt = f"{req.prompt} --resolution {final_size}"
-        else:
-             final_prompt = f"{req.prompt} --ar {ratio_str} --resolution {final_size}"
-
-    # 3. 初始化客户端
-    client = AsyncOpenAI(
-        api_key=req.api_key, 
-        base_url="https://api.cloudapp.ink/v1",
-        max_retries=0
-    )
-
-    image_bytes = None
-
-    # 4. 执行请求
-    if "gemini" in req.model.lower():
-        # --- Gemini Logic ---
-        print(f"🟣 [Chat模式] 请求 Gemini: {req.model}")
-        response = await client.chat.completions.create(
-            model=req.model,
-            messages=[{"role": "user", "content": final_prompt}],
-            extra_body={"modalities": ["image", "text"]}
-        )
-        content = response.choices[0].message.content
-        
-        # 提取链接或 Base64
-        url_match = re.search(r"!\[.*?\]\((https?://[^\)]+)\)", content)
-        if not url_match:
-            url_match = re.search(r"(https?://\S+\.(?:png|jpg|jpeg|webp))", content)
-        
-        if url_match:
-            image_bytes = await download_image_as_bytes(url_match.group(1))
-        elif "base64," in content:
-            print(f"🟣 [Chat模式] 检测到 Base64 流")
-            b64_str = content.split("base64,")[1].split(")")[0].strip()
-            image_bytes = base64.b64decode(b64_str)
-        else:
-            raise Exception("未能识别图片内容")
-            
-    else:
-        # --- Nano Banana (Standard) Logic ---
-        print(f"🔵 [绘图模式] 请求 Standard: {req.model} | 分辨率: {final_size}")
-        response = await client.images.generate(
-            model=req.model,
-            prompt=final_prompt,
-            n=1,
-            size=final_size, 
-            response_format="b64_json"
-        )
-        print(f"🔵 [绘图模式] 请求成功")
-        
-        data = response.data[0]
-        if getattr(data, 'b64_json', None):
-            image_bytes = base64.b64decode(data.b64_json)
-        elif hasattr(data, 'url'):
-            image_bytes = await download_image_as_bytes(data.url)
-        else:
-            raise Exception("API 未返回有效图片数据")
-
-    # 5. 格式转换
-    return process_image_in_memory(image_bytes, req.format)
-
-
-# ==========================================
-# 3. FastAPI App
-# ==========================================
-app = FastAPI()
-
-app.mount("/static", StaticFiles(directory="static"), name="static")
-
-@app.get("/")
-async def read_index():
-    return FileResponse('static/index.html')
-
-@app.get("/health")
-def health():
-    return {"status": "ok"}
-
-@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"📌 模型: {req.model} | 规格: {req.ratio} ({req.scale})")
-    print("-" * 60)
-
-    # --- 关键修改：使用 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()
+import httpx
+import re
+import base64
+import os
+import asyncio
+import time
+import gc
+from datetime import datetime
+from PIL import Image
+from io import BytesIO
+from openai import AsyncOpenAI
+from fastapi import FastAPI, HTTPException, Request
+from fastapi.staticfiles import StaticFiles
+from fastapi.responses import FileResponse, JSONResponse
+from pydantic import BaseModel
+import uvicorn
+
+# ==========================================
+# 1. 配置区域
+# ==========================================
+
+# --- 百度内容审核配置 (从环境变量获取) ---
+# 如果环境变量没配置，默认为空字符串，相关功能会自动关闭
+BAIDU_API_KEY = os.getenv("BAIDU_API_KEY", "")
+BAIDU_SECRET_KEY = os.getenv("BAIDU_SECRET_KEY", "")
+
+# --- 并发限制 ---
+CONCURRENCY_LIMITER = asyncio.Semaphore(3)
+
+# --- 基础配置 ---
+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)
+}
+
+SCALE_MULTIPLIERS = {
+    "1k": 1.0, "2k": 1.5, "4k": 2.0
+}
+
+# --- 本地敏感词 (第一道防线) ---
+SENSITIVE_WORDS = [
+    "nsfw", "nude", "naked", "sex", "porn", "hentai", 
+    "裸", "色情", "全裸", "blood", "kill", "murder", 
+    "血腥", "尸体", "毒品"
+]
+
+# ==========================================
+# 2. 工具函数 (百度审核 & 图像处理)
+# ==========================================
+
+# 全局缓存 Access Token
+_BAIDU_TOKEN_CACHE = {"token": None, "expires_at": 0}
+
+async def get_baidu_access_token():
+    """获取百度 Access Token"""
+    # 如果环境变量没填，直接返回 None，跳过百度审核
+    if not BAIDU_API_KEY or not BAIDU_SECRET_KEY:
+        return None
+
+    now = time.time()
+    if _BAIDU_TOKEN_CACHE["token"] and now < _BAIDU_TOKEN_CACHE["expires_at"]:
+        return _BAIDU_TOKEN_CACHE["token"]
+
+    url = "https://aip.baidubce.com/oauth/2.0/token"
+    params = {
+        "grant_type": "client_credentials",
+        "client_id": BAIDU_API_KEY,
+        "client_secret": BAIDU_SECRET_KEY
+    }
+    
+    try:
+        async with httpx.AsyncClient(timeout=10) as client:
+            resp = await client.post(url, params=params)
+            data = resp.json()
+            token = data.get("access_token")
+            if token:
+                _BAIDU_TOKEN_CACHE["token"] = token
+                _BAIDU_TOKEN_CACHE["expires_at"] = now + (20 * 24 * 3600)
+                return token
+            else:
+                print(f"⚠️ 百度 Token 获取失败: {data}")
+                return None
+    except Exception as e:
+        print(f"⚠️ 百度 Token 网络错误: {e}")
+        return None
+
+async def check_baidu_text_censor(text):
+    """调用百度 API 检查文本"""
+    token = await get_baidu_access_token()
+    if not token:
+        # 如果获取不到 token (未配置环境变量或接口挂了)，默认放行
+        return False, None
+
+    url = f"https://aip.baidubce.com/rest/2.0/solution/v1/text_censor/v2/user_defined?access_token={token}"
+    data = {'text': text}
+    
+    try:
+        async with httpx.AsyncClient(timeout=10) as client:
+            resp = await client.post(url, data=data)
+            result = resp.json()
+            
+            # conclusionType: 1.合规 2.不合规 3.疑似 4.审核失败
+            conclusion_type = result.get('conclusionType')
+            
+            if conclusion_type == 2:
+                msg_list = []
+                if 'data' in result:
+                    for item in result['data']:
+                        msg_list.append(item.get('msg', '违规内容'))
+                reason = ",".join(msg_list) if msg_list else "内容不合规"
+                return True, reason
+            elif conclusion_type == 3:
+                return True, "内容疑似违规"
+            
+            return False, None
+
+    except Exception as e:
+        print(f"⚠️ 百度审核接口异常: {e}")
+        return False, None
+
+def contains_local_sensitive_words(text: str):
+    text_lower = text.lower()
+    for word in SENSITIVE_WORDS:
+        if word.lower() in text_lower:
+            return True, word
+    return False, None
+
+class GenerateRequest(BaseModel):
+    api_key: str
+    model: str
+    prompt: str
+    ratio: str
+    scale: str
+    custom_size: str
+    format: str
+
+def process_image_in_memory(image_bytes: bytes, target_format: str) -> str:
+    try:
+        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')
+    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=60) as client:
+        response = await client.get(url)
+        if response.status_code == 200:
+            return response.content
+        else:
+            raise Exception(f"下载失败 HTTP {response.status_code}")
+
+# ==========================================
+# 3. 核心生成逻辑
+# ==========================================
+
+async def core_generate(req: GenerateRequest):
+    final_size = "1024x1024"
+    ratio_str = "1:1"
+    if req.ratio == "custom":
+        final_size = req.custom_size.strip()
+    else:
+        base_w, base_h = AR_BASES.get(req.ratio, (1024, 1024))
+        multiplier = SCALE_MULTIPLIERS.get(req.scale, 1.0)
+        final_size = f"{int(base_w * multiplier)}x{int(base_h * multiplier)}"
+        ratio_str = req.ratio.split(" ")[0]
+
+    final_prompt = req.prompt
+    if "gemini" in req.model.lower():
+        if req.ratio == "custom":
+             final_prompt = f"{req.prompt} --resolution {final_size}"
+        else:
+             final_prompt = f"{req.prompt} --ar {ratio_str} --resolution {final_size}"
+
+    client = AsyncOpenAI(
+        api_key=req.api_key, 
+        base_url="https://api.cloudapp.ink/v1",
+        max_retries=0, 
+        timeout=60.0
+    )
+
+    image_bytes = None
+
+    if "gemini" in req.model.lower():
+        print(f"🟣 [Chat模式] 请求 Gemini: {req.model}")
+        response = await client.chat.completions.create(
+            model=req.model,
+            messages=[{"role": "user", "content": final_prompt}],
+            extra_body={"modalities": ["image", "text"]}
+        )
+        content = response.choices[0].message.content
+        
+        url_match = re.search(r"!\[.*?\]\((https?://[^\)]+)\)", content)
+        if not url_match:
+            url_match = re.search(r"(https?://\S+\.(?:png|jpg|jpeg|webp))", content)
+        
+        if url_match:
+            image_bytes = await download_image_as_bytes(url_match.group(1))
+        elif "base64," in content:
+            print(f"🟣 [Chat模式] 检测到 Base64 流")
+            b64_str = content.split("base64,")[1].split(")")[0].strip()
+            image_bytes = base64.b64decode(b64_str)
+        else:
+            raise Exception("未能识别图片内容")
+            
+    else:
+        print(f"🔵 [绘图模式] 请求 Standard: {req.model} | 分辨率: {final_size}")
+        response = await client.images.generate(
+            model=req.model,
+            prompt=final_prompt,
+            n=1,
+            size=final_size, 
+            response_format="b64_json"
+        )
+        print(f"🔵 [绘图模式] 请求成功")
+        
+        data = response.data[0]
+        if getattr(data, 'b64_json', None):
+            image_bytes = base64.b64decode(data.b64_json)
+        elif hasattr(data, 'url'):
+            image_bytes = await download_image_as_bytes(data.url)
+        else:
+            raise Exception("API 未返回有效图片数据")
+
+    return process_image_in_memory(image_bytes, req.format)
+
+# ==========================================
+# 4. FastAPI App
+# ==========================================
+app = FastAPI()
+
+app.mount("/static", StaticFiles(directory="static"), name="static")
+
+@app.get("/")
+async def read_index():
+    return FileResponse('static/index.html')
+
+@app.get("/health")
+def health():
+    return {"status": "ok"}
+
+@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"🚀 [Request] {timestamp} | IP: {client_host}")
+    print(f"📌 Model: {req.model}")
+    print(f"💡 Prompt: {req.prompt}")
+
+    # 1. 本地敏感词检查
+    is_local_sensitive, word = contains_local_sensitive_words(req.prompt)
+    if is_local_sensitive:
+        print(f"🚫 [Blocked-Local] 命中本地词库: {word}")
+        return JSONResponse(
+            content={
+                "status": "error", 
+                "message": f"提示词包含违禁词汇 ({word})，请求被拒绝。"
+            },
+            status_code=200 
+        )
+
+    # 2. 百度 API 检查 (动态读取环境变量)
+    is_baidu_sensitive, reason = await check_baidu_text_censor(req.prompt)
+    if is_baidu_sensitive:
+        print(f"🚫 [Blocked-Baidu] 百度审核拦截: {reason}")
+        return JSONResponse(
+            content={
+                "status": "error", 
+                "message": f"提示词未通过智能安全审核: {reason}"
+            },
+            status_code=200 
+        )
+
+    print("-" * 60)
+
+    # 3. 正常生成
+    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()
+
+if __name__ == "__main__":
+    uvicorn.run(app, host="0.0.0.0", port=7860)