Update app.py

app.py

@@ -10,12 +10,40 @@ from dotenv import load_dotenv
 load_dotenv()
 
 # ========== 默认选项和数据 ==========
-EXPRESSIONS = ["smiling", "determined", "surprised", "serene", "smug", "thinking", "looking back", "laughing", "angry", "pensive", "confident", "grinning", "thoughtful", "sad tears", "bewildered"]
-ITEMS = ["magic wand", "sword", "flower", "book of spells", "earrings", "loincloth", "slippers", "ancient scroll", "music instrument", "shield", "dagger", "headband", "leg ties", "staff", "potion", "crystal ball", "anklet", "ribbon", "lantern", "amulet", "ring"]
-OTHER_DETAILS = ["sparkles", "magical aura", "lens flare", "fireworks in the background", "smoke effects", "light trails", "falling leaves", "glowing embers", "floating particles", "rays of light", "shimmering mist", "ethereal glow"]
-SCENES = ["sunset beach", "rainy city street at night", "floating ash land", "particles magic world", "high blue sky", "top of the building", "fantasy forest with glowing mushrooms", "futuristic skyline at dawn", "abandoned castle", "snowy mountain peak", "desert ruins", "underwater city", "enchanted meadow", "haunted mansion", "steampunk marketplace", "glacial cavern"]
-CAMERA_ANGLES = ["low-angle shot", "close-up shot", "bird's-eye view", "wide-angle shot", "over-the-shoulder shot", "extreme close-up", "panoramic view", "dynamic tracking shot", "fisheye view", "point-of-view shot"]
-QUALITY_PROMPTS = ["cinematic lighting", "sharp shadow", "award-winning", "masterpiece", "vivid colors", "high dynamic range", "immersive", "studio quality", "fine art", "dreamlike", "8K", "HD", "high quality", "best quality", "artistic", "vibrant"]
+EXPRESSIONS = [
+    "smiling", "determined", "surprised", "serene", "smug", "thinking", 
+    "looking back", "laughing", "angry", "pensive", "confident", 
+    "grinning", "thoughtful", "sad tears", "bewildered"
+]
+ITEMS = [
+    "magic wand", "sword", "flower", "book of spells", "earrings", "loincloth", 
+    "slippers", "ancient scroll", "music instrument", "shield", "dagger", 
+    "headband", "leg ties", "staff", "potion", "crystal ball", "anklet", 
+    "ribbon", "lantern", "amulet", "ring"
+]
+OTHER_DETAILS = [
+    "sparkles", "magical aura", "lens flare", "fireworks in the background", 
+    "smoke effects", "light trails", "falling leaves", "glowing embers", 
+    "floating particles", "rays of light", "shimmering mist", "ethereal glow"
+]
+SCENES = [
+    "sunset beach", "rainy city street at night", "floating ash land", 
+    "particles magic world", "high blue sky", "top of the building", 
+    "fantasy forest with glowing mushrooms", "futuristic skyline at dawn", 
+    "abandoned castle", "snowy mountain peak", "desert ruins", "underwater city", 
+    "enchanted meadow", "haunted mansion", "steampunk marketplace", "glacial cavern"
+]
+CAMERA_ANGLES = [
+    "low-angle shot", "close-up shot", "bird's-eye view", "wide-angle shot", 
+    "over-the-shoulder shot", "extreme close-up", "panoramic view", 
+    "dynamic tracking shot", "fisheye view", "point-of-view shot"
+]
+QUALITY_PROMPTS = [
+    "cinematic lighting", "sharp shadow", "award-winning", "masterpiece", 
+    "vivid colors", "high dynamic range", "immersive", "studio quality", 
+    "fine art", "dreamlike", "8K", "HD", "high quality", "best quality", 
+    "artistic", "vibrant"
+]
 
 # Hugging Face DTR 数据集路径（示例，若不可用请忽略）
 DTR_DATASET_PATTERN = "https://huggingface.co/datasets/X779/Danbooruwildcards/resolve/main/*DTR*.txt"
@@ -68,7 +96,7 @@ def load_dtr_from_huggingface(excluded_tags=None):
         print(f"Error loading DTR dataset: {e}")
         return []
 
-def generate_natural_language_description(tags, api_key=None, base_url=None, model="gpt-4o"):
+def generate_natural_language_description(tags, api_key=None, base_url=None, model="gpt-4"):
     """
     使用 OpenAI GPT 或 DeepSeek API 生成自然语言描述。
     """
@@ -77,7 +105,7 @@ def generate_natural_language_description(tags, api_key=None, base_url=None, mod
     if not api_key:
         return "Error: No API Key provided and none found in environment variables."
 
-    # 将dict转成可读字符串
+    # 将 dict 转成可读字符串
     tag_descriptions = "\n".join([
         f"{key}: {', '.join(value) if isinstance(value, list) else value}"
         for key, value in tags.items() if value
@@ -106,6 +134,8 @@ def generate_natural_language_description(tags, api_key=None, base_url=None, mod
     except Exception as e:
         return f"GPT generation failed. Error: {e}"
 
+
+# ========== 核心函数：随机生成 prompt ==========
 def generate_prompt(
     action_file, style_file, artist_files, character_files, dtr_enabled, api_key, selected_categories,
     expression_count, item_count, detail_count, scene_count, angle_count, quality_count, action_count, style_count,
@@ -194,12 +224,98 @@ def generate_prompt(
     combined_output = f"{final_tags}\n\n{description}"
     return final_tags, description, combined_output
 
-# ========== Favorite 相关函数 ==========
+
+# ========== 部分更新：只根据用户修改后的 tags_text 生成新的描述和合并输出 ==========
+def update_description(tags_text, api_key, use_deepseek, deepseek_key):
+    """
+    只根据用户提供的 tags_text 生成描述和合并输出。
+    不再重新随机抽取，以免破坏用户手动修改过的 Tags。
+    """
+    if not api_key and not deepseek_key:
+        # 没有提供任意可用 API Key
+        return "(No API Key provided)", f"{tags_text}\n\n(No API Key provided)"
+
+    # 构造给 GPT 的 prompt
+    user_prompt = (
+        "You are a creative assistant that generates detailed, imaginative scene descriptions for AI generation.\n"
+        "Below is the user's current tags (prompt elements). "
+        "Generate a new descriptive text (3-5 sentences) that incorporates these tags.\n\n"
+        f"User Tags: {tags_text}\n"
+        "Please generate a vivid, imaginative scene description."
+    )
+
+    try:
+        if use_deepseek:
+            # 调用 DeepSeek
+            client = OpenAI(api_key=deepseek_key, base_url="https://api.deepseek.com")
+            model = "deepseek-chat"
+        else:
+            # 调用 OpenAI
+            client = OpenAI(api_key=api_key)
+            model = "gpt-4"  # 或其他可用模型，比如 "gpt-3.5-turbo"
+        response = client.chat.completions.create(
+            messages=[
+                {
+                    "role": "system",
+                    "content": "You are a creative assistant that generates imaginative scene descriptions..."
+                },
+                {
+                    "role": "user",
+                    "content": user_prompt,
+                },
+            ],
+            model=model,
+        )
+        new_description = response.choices[0].message.content.strip()
+    except Exception as e:
+        new_description = f"(GPT generation failed: {e})"
+
+    new_combined_output = f"{tags_text}\n\n{new_description}"
+    return new_description, new_combined_output
+
+
+# ========== 翻译功能：将 combined_output 翻译成用户选定语言 ==========
+def translate_combined_output(combined_text, target_language, api_key, use_deepseek, deepseek_key):
+    """
+    使用 GPT 或 DeepSeek API，将 combined_text 翻译成 target_language。
+    """
+    if not api_key and not deepseek_key:
+        return "(No API Key provided)"
+
+    # 简单用 GPT 做翻译，也可改成其他翻译 API
+    translation_prompt = (
+        f"You are a professional translator. Please translate the following text into {target_language}.\n\n"
+        f"{combined_text}"
+    )
+
+    try:
+        if use_deepseek:
+            # 调用 DeepSeek
+            client = OpenAI(api_key=deepseek_key, base_url="https://api.deepseek.com")
+            model = "deepseek-chat"
+        else:
+            # 调用 OpenAI
+            client = OpenAI(api_key=api_key)
+            model = "gpt-3.5-turbo"  # 或者别的模型
+        response = client.chat.completions.create(
+            messages=[
+                {"role": "system", "content": "You are a professional translator."},
+                {"role": "user", "content": translation_prompt},
+            ],
+            model=model,
+        )
+        translated_text = response.choices[0].message.content.strip()
+    except Exception as e:
+        translated_text = f"(Translation failed: {e})"
+
+    return translated_text
+
+
+# ========== 收藏功能：最多存 3 条 ==========
 def add_to_favorites(combined_output, current_favorites):
     """
     将当前生成的 combined_output 添加到收藏列表中（最多存 3 条）。
     """
-    # current_favorites 是一个列表
     current_favorites.append(combined_output)
     # 如果超过3条，移除最早的一条
     if len(current_favorites) > 3:
@@ -210,6 +326,7 @@ def add_to_favorites(combined_output, current_favorites):
     )
     return favorites_text, current_favorites
 
+
 # ========== Gradio 界面 ==========
 def gradio_interface():
     """
@@ -237,10 +354,8 @@ def gradio_interface():
                 )
 
                 use_deepseek = gr.Checkbox(label="Use DeepSeek API")
-
                 dtr_enabled = gr.Checkbox(label="Enable DTR (如不可用请勿勾选)")
 
-                # 上传文件部分
                 with gr.Group():
                     gr.Markdown("**上传文件 (可选)：**")
                     action_file = gr.File(label="Action File", file_types=[".txt"])
@@ -248,13 +363,11 @@ def gradio_interface():
                     artist_files = gr.Files(label="Artist Files", file_types=[".txt"])
                     character_files = gr.Files(label="Character Files", file_types=[".txt"])
 
-                # 选择角色类型
                 selected_categories = gr.CheckboxGroup(
                     ["1boy", "1girl", "furry", "mecha", "fantasy monster", "animal", "still life"],
                     label="Choose Character Categories"
                 )
 
-                # 输入排除和自定义
                 excluded_tags = gr.Textbox(
                     label="排除 Tags (逗号分隔)",
                     placeholder="如：angry, sword"
@@ -264,7 +377,6 @@ def gradio_interface():
                     placeholder="如：glowing eyes, giant wings"
                 )
 
-                # 各种数量
                 with gr.Group():
                     gr.Markdown("**随机数量设置：**")
                     expression_count = gr.Slider(label="Number of Expressions", minimum=0, maximum=10, step=1, value=1)
@@ -277,11 +389,10 @@ def gradio_interface():
                     style_count = gr.Slider(label="Number of Styles", minimum=1, maximum=10, step=1, value=1)
                     artist_count = gr.Slider(label="Number of Artists", minimum=1, maximum=10, step=1, value=1)
 
-            # 右侧：生成按钮 + 生成结果 + 收藏
+            # 右侧：生成按钮 + 生成结果 + 收藏 + 翻译
             with gr.Column(scale=2):
                 generate_button = gr.Button("Generate Prompt", variant="primary")
 
-                # 生成的结果（可编辑）
                 tags_output = gr.Textbox(
                     label="Generated Tags",
                     placeholder="等待生成...",
@@ -300,8 +411,25 @@ def gradio_interface():
                     lines=6
                 )
 
+                # 新增一个按钮，只更新 description 和 combined
+                update_desc_button = gr.Button("Update Description Only")
+
+                # 翻译相关
+                with gr.Row():
+                    target_language = gr.Dropdown(
+                        choices=["English", "Chinese", "Japanese"],
+                        value="English",
+                        label="Target Language"
+                    )
+                    translate_button = gr.Button("Translate to selected language")
+                translated_output = gr.Textbox(
+                    label="Translated Output",
+                    placeholder="等待翻译...",
+                    lines=6
+                )
+
+                # 收藏
                 with gr.Row():
-                    # 收藏操作
                     favorite_button = gr.Button("收藏本次结果")
                     favorites_box = gr.Textbox(
                         label="收藏夹 (最多 3 条)",
@@ -309,17 +437,45 @@ def gradio_interface():
                         lines=6
                     )
 
-                # 点击生成按钮
+                # 点击“Generate Prompt”按钮
                 generate_button.click(
                     generate_prompt,
                     inputs=[
-                        action_file, style_file, artist_files, character_files, dtr_enabled, api_key_input, selected_categories,
-                        expression_count, item_count, detail_count, scene_count, angle_count, quality_count, action_count, style_count,
-                        artist_count, use_deepseek, deepseek_key_input, user_custom_tags, excluded_tags
+                        action_file, style_file, artist_files, character_files, 
+                        dtr_enabled, api_key_input, selected_categories,
+                        expression_count, item_count, detail_count, scene_count, 
+                        angle_count, quality_count, action_count, style_count,
+                        artist_count, use_deepseek, deepseek_key_input, 
+                        user_custom_tags, excluded_tags
                     ],
                     outputs=[tags_output, description_output, combined_output],
                 )
 
+                # 点击“Update Description Only”按钮
+                update_desc_button.click(
+                    update_description,
+                    inputs=[
+                        tags_output,     # 用户在文本框里编辑后的 Tags
+                        api_key_input, 
+                        use_deepseek,
+                        deepseek_key_input,
+                    ],
+                    outputs=[description_output, combined_output],
+                )
+
+                # 点击“Translate to selected language”按钮
+                translate_button.click(
+                    fn=translate_combined_output,
+                    inputs=[
+                        combined_output,   # 要翻译的源文本
+                        target_language,
+                        api_key_input,
+                        use_deepseek,
+                        deepseek_key_input
+                    ],
+                    outputs=[translated_output],
+                )
+
                 # 收藏按钮点击事件
                 favorite_button.click(
                     fn=add_to_favorites,
@@ -329,6 +485,7 @@ def gradio_interface():
 
     return demo
 
-# 启动 Gradio 应用（在 Hugging Face Spaces 上只需要保留最后的 launch 即可）
+
+# 启动 Gradio 应用
 if __name__ == "__main__":
     gradio_interface().launch()