1

app.py

@@ -19,6 +19,52 @@ pipe = pipe.to(device)
 
 MAX_SEED = np.iinfo(np.int32).max
 MAX_IMAGE_SIZE = 768  # 减小最大尺寸以提高生成速度
+# 提示词优化：调用 Gemini Pro API
+import requests
+
+def optimize_prompt(en_text):
+    api_key = "AIzaSyB8Qu7XLzR6vnmnBN19z2cAXVRrJYjr2KY"
+    url = "https://generativelanguage.googleapis.com/v1beta/models/gemini-2.0-flash:generateContent"
+    headers = {
+        "Content-Type": "application/json",
+        "X-goog-api-key": api_key
+    }
+    data = {
+        "contents": [{
+            "parts": [{
+                "text": f"请将以下英文提示词优化为适合AI文生图生成的英文提示词，要求简洁、描述清晰、突出画面细节：{en_text}"
+            }]
+        }]
+    }
+    print(f"[LOG] 优化提示词请求: {en_text}")
+    try:
+        resp = requests.post(url, headers=headers, json=data, timeout=10)
+        print(f"[LOG] Gemini API响应状态: {resp.status_code}")
+        resp.raise_for_status()
+        result = resp.json()
+        print(f"[LOG] Gemini API响应内容: {result}")
+        opt_text = result["candidates"][0]["content"]["parts"][0]["text"]
+        print(f"[LOG] 优化后提示词: {opt_text}")
+        return opt_text.strip()
+    except Exception as e:
+        print(f"[ERROR] Gemini优化失败: {e}")
+        return en_text
+# 添加自动翻译功能
+from transformers import MarianMTModel, MarianTokenizer
+
+def translate_prompt(text):
+    # 检查是否包含中文字符
+    if any('\u4e00' <= ch <= '\u9fff' for ch in text):
+        print(f"[LOG] 检测到中文提示词: {text}")
+        model_name = 'Helsinki-NLP/opus-mt-zh-en'
+        tokenizer = MarianTokenizer.from_pretrained(model_name)
+        model = MarianMTModel.from_pretrained(model_name)
+        translated = model.generate(**tokenizer(text, return_tensors="pt", padding=True))
+        en_text = tokenizer.decode(translated[0], skip_special_tokens=True)
+        print(f"[LOG] 翻译后的英文提示词: {en_text}")
+        return en_text
+    print(f"[LOG] 非中文提示词，直接使用: {text}")
+    return text
 
 # 熊猫烧香相关的提示词模板
 PANDA_INCENSE_PROMPTS = [
@@ -45,13 +91,25 @@ def infer(
     num_inference_steps,
     progress=gr.Progress(track_tqdm=True),
 ):
+    # 自动翻译中文提示词为英文
+    print(f"[LOG] 用户输入提示词: {prompt}")
+    prompt_en = translate_prompt(prompt)
+    print(f"[LOG] 英文提示词: {prompt_en}")
+    prompt_opt = optimize_prompt(prompt_en)
+    print(f"[LOG] 最终用于生成的提示词: {prompt_opt}")
+
     if randomize_seed:
         seed = random.randint(0, MAX_SEED)
+        print(f"[LOG] 随机生成种子: {seed}")
+    else:
+        print(f"[LOG] 使用用户指定种子: {seed}")
 
     generator = torch.Generator().manual_seed(seed)
+    print(f"[LOG] 生成器初始化完成")
 
+    print(f"[LOG] 开始生成图片，参数: guidance_scale={guidance_scale}, steps={num_inference_steps}, width={width}, height={height}")
     image = pipe(
-        prompt=prompt,
+        prompt=prompt_opt,
         negative_prompt=negative_prompt,
         guidance_scale=guidance_scale,
         num_inference_steps=num_inference_steps,
@@ -59,7 +117,7 @@ def infer(
         height=height,
         generator=generator,
     ).images[0]
-
+    print(f"[LOG] 图片生成完成")
     return image, seed
 
 

requirements.txt

@@ -4,4 +4,5 @@ invisible_watermark
 torch
 transformers
 xformers
-gradio
+gradio
+requests