Upload app.py

app.py

@@ -16,13 +16,23 @@ from pydantic import BaseModel
 import uvicorn
 
 # ==========================================
-# 1. 静态配置 (不要在这里读环境变量)
+# 1. 配置区域
 # ==========================================
 
-# 并发限制
+# --- 百度内容审核配置 ---
+BAIDU_API_KEY = os.getenv("BAIDU_API_KEY", "")
+BAIDU_SECRET_KEY = os.getenv("BAIDU_SECRET_KEY", "")
+
+# --- 启动检查日志 ---
+if BAIDU_API_KEY and BAIDU_SECRET_KEY:
+    print(f"✅ [System] 检测到百度 API Key: {BAIDU_API_KEY[:4]}****")
+else:
+    print(f"⚠️ [System] 未检测到百度 API Key 或 SECRET KEY，审核功能将关闭！")
+
+# --- 并发限制 ---
 CONCURRENCY_LIMITER = asyncio.Semaphore(3)
 
-# 基础分辨率
+# --- 基础配置 ---
 AR_BASES = {
     "1:1 (SDXL标准 1024)": (1024, 1024),
     "16:9 (SDXL标准 768p)": (1344, 768),
@@ -36,92 +46,76 @@ SCALE_MULTIPLIERS = {
     "1k": 1.0, "2k": 1.5, "4k": 2.0
 }
 
-# 本地敏感词 (第一道防线)
+# --- 本地敏感词 (第一道防线) ---
 SENSITIVE_WORDS = [
     "nsfw", "nude", "naked", "sex", "porn", "hentai", 
     "裸", "色情", "全裸", "blood", "kill", "murder", 
-    "血腥", "尸体", "毒品", "杀"
+    "血腥", "尸体", "毒品"
 ]
 
 # ==========================================
-# 2. 百度审核模块 (动态读取环境变量)
+# 2. 工具函数 (百度审核 & 图像处理)
 # ==========================================
 
 # 全局缓存 Access Token
 _BAIDU_TOKEN_CACHE = {"token": None, "expires_at": 0}
 
 async def get_baidu_access_token():
-    """获取百度 Access Token (运行时从环境变量读取)"""
-    
-    # --- 关键：运行时动态读取 ---
-    ak = os.getenv("BAIDU_API_KEY", "").strip()
-    sk = os.getenv("BAIDU_SECRET_KEY", "").strip()
-
-    # --- 调试日志 (只打印一次) ---
-    if not hasattr(get_baidu_access_token, "debug_printed"):
-        print(f"\n🔍 [Debug] 正在尝试读取环境变量:")
-        print(f"   - BAIDU_API_KEY: {'✅ 已配置' if ak else '❌ 未配置'}")
-        print(f"   - BAIDU_SECRET_KEY: {'✅ 已配置' if sk else '❌ 未配置'}")
-        get_baidu_access_token.debug_printed = True
-
-    # 如果没配置，直接返回 None
-    if not ak or not sk:
+    """获取百度 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": ak,
-        "client_secret": sk
+        "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()
-            
-            if "error" in data:
-                print(f"❌ [Baidu Error] Token 获取失败: {data.get('error_description')}")
-                return None
-
             token = data.get("access_token")
             if token:
                 _BAIDU_TOKEN_CACHE["token"] = token
-                _BAIDU_TOKEN_CACHE["expires_at"] = now + (20 * 24 * 3600) # 缓存20天
+                _BAIDU_TOKEN_CACHE["expires_at"] = now + (20 * 24 * 3600)
                 return token
             else:
+                print(f"⚠️ 百度 Token 获取失败: {data}")
                 return None
     except Exception as e:
-        print(f"⚠️ [System] 百度网络错误: {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就放行，保证服务可用
+        # 如果获取不到 token (未配置环境变量或接口挂了)，默认放行
         return False, None
 
     url = f"https://aip.baidubce.com/rest/2.0/solution/v1/text_censor/v2/user_defined?access_token={token}"
-    headers = {'Content-Type': 'application/x-www-form-urlencoded'}
     data = {'text': text}
     
     try:
         async with httpx.AsyncClient(timeout=10) as client:
-            resp = await client.post(url, data=data, headers=headers)
+            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 = [item.get('msg', '违规内容') for item in result.get('data', [])]
+                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:
@@ -130,12 +124,15 @@ async def check_baidu_text_censor(text):
             return False, None
 
     except Exception as e:
-        print(f"⚠️ [System] 审核接口异常: {e}")
+        print(f"⚠️ 百度审核接口异常: {e}")
         return False, None
 
-# ==========================================
-# 3. 图像生成模块
-# ==========================================
+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
@@ -146,13 +143,6 @@ class GenerateRequest(BaseModel):
     custom_size: str
     format: str
 
-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
-
 def process_image_in_memory(image_bytes: bytes, target_format: str) -> str:
     try:
         with Image.open(BytesIO(image_bytes)) as img:
@@ -164,12 +154,12 @@ def process_image_in_memory(image_bytes: bytes, target_format: str) -> str:
                 img.save(output_buffer, format="PNG")
             return base64.b64encode(output_buffer.getvalue()).decode('utf-8')
     except Exception as e:
-        raise Exception(f"图片处理失败: {str(e)}")
+        raise Exception(f"格式转换失败: {str(e)}")
     finally:
         gc.collect()
 
 async def download_image_as_bytes(url):
-    print(f"📥 [System] 下载图片: {url}")
+    print(f"📥 [System] 正在下载图片: {url}")
     async with httpx.AsyncClient(timeout=60) as client:
         response = await client.get(url)
         if response.status_code == 200:
@@ -177,8 +167,11 @@ async def download_image_as_bytes(url):
         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":
@@ -189,7 +182,6 @@ async def core_generate(req: GenerateRequest):
         final_size = f"{int(base_w * multiplier)}x{int(base_h * multiplier)}"
         ratio_str = req.ratio.split(" ")[0]
 
-    # Prompt
     final_prompt = req.prompt
     if "gemini" in req.model.lower():
         if req.ratio == "custom":
@@ -197,8 +189,6 @@ async def core_generate(req: GenerateRequest):
         else:
              final_prompt = f"{req.prompt} --ar {ratio_str} --resolution {final_size}"
 
-    # 初始化 OpenAI 客户端
-    # 注意：max_retries=0 禁止自动重试，防止重复计费
     client = AsyncOpenAI(
         api_key=req.api_key, 
         base_url="https://api.cloudapp.ink/v1",
@@ -209,7 +199,7 @@ async def core_generate(req: GenerateRequest):
     image_bytes = None
 
     if "gemini" in req.model.lower():
-        print(f"🟣 [Chat] 请求 Gemini: {req.model}")
+        print(f"🟣 [Chat模式] 请求 Gemini: {req.model}")
         response = await client.chat.completions.create(
             model=req.model,
             messages=[{"role": "user", "content": final_prompt}],
@@ -224,14 +214,14 @@ async def core_generate(req: GenerateRequest):
         if url_match:
             image_bytes = await download_image_as_bytes(url_match.group(1))
         elif "base64," in content:
-            print(f"🟣 [Chat] Base64数据流")
+            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"🔵 [Draw] 请求 Standard: {req.model} | {final_size}")
+        print(f"🔵 [绘图模式] 请求 Standard: {req.model} | 分辨率: {final_size}")
         response = await client.images.generate(
             model=req.model,
             prompt=final_prompt,
@@ -239,7 +229,7 @@ async def core_generate(req: GenerateRequest):
             size=final_size, 
             response_format="b64_json"
         )
-        print(f"🔵 [Draw] 响应成功")
+        print(f"🔵 [绘图模式] 请求成功")
         
         data = response.data[0]
         if getattr(data, 'b64_json', None):
@@ -247,16 +237,15 @@ async def core_generate(req: GenerateRequest):
         elif hasattr(data, 'url'):
             image_bytes = await download_image_as_bytes(data.url)
         else:
-            raise Exception("API 未返回图片数据")
+            raise Exception("API 未返回有效图片数据")
 
     return process_image_in_memory(image_bytes, req.format)
 
 # ==========================================
-# 4. FastAPI 启动入口
+# 4. FastAPI App
 # ==========================================
 app = FastAPI()
 
-# 挂载静态文件
 app.mount("/static", StaticFiles(directory="static"), name="static")
 
 @app.get("/")
@@ -273,41 +262,48 @@ async def generate_api(req: GenerateRequest, request: Request):
     timestamp = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
     
     print("\n" + "="*60)
-    print(f"🚀 [Req] {timestamp} | IP: {client_host}")
+    print(f"🚀 [Request] {timestamp} | IP: {client_host}")
     print(f"📌 Model: {req.model}")
-    
-    # --- 1. 本地敏感词拦截 ---
+    print(f"💡 Prompt: {req.prompt}")
+
+    # 1. 本地敏感词检查
     is_local_sensitive, word = contains_local_sensitive_words(req.prompt)
     if is_local_sensitive:
-        print(f"🚫 [Block-Local] 命中本地词库: {word}")
+        print(f"🚫 [Blocked-Local] 命中本地词库: {word}")
         return JSONResponse(
-            content={"status": "error", "message": f"违规词汇拦截: {word}"},
+            content={
+                "status": "error", 
+                "message": f"提示词包含违禁词汇 ({word})，请求被拒绝。"
+            },
             status_code=200 
         )
 
-    # --- 2. 百度云智能拦截 ---
+    # 2. 百度 API 检查 (动态读取环境变量)
     is_baidu_sensitive, reason = await check_baidu_text_censor(req.prompt)
     if is_baidu_sensitive:
-        print(f"🚫 [Block-Baidu] 百度审核拦截: {reason}")
+        print(f"🚫 [Blocked-Baidu] 百度审核拦截: {reason}")
         return JSONResponse(
-            content={"status": "error", "message": f"安全审核未通过: {reason}"},
+            content={
+                "status": "error", 
+                "message": f"提示词未通过智能安全审核: {reason}"
+            },
             status_code=200 
         )
 
-    print(f"💡 Prompt: {req.prompt}")
     print("-" * 60)
 
-    # --- 3. 执行生成 ---
+    # 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)}")
+            print(f"❌ [Failed] 错误: {str(e)}")
             return JSONResponse(
                 content={"status": "error", "message": str(e)},
                 status_code=200