Update app.py

app.py

@@ -17,20 +17,31 @@ SCENES = ["sunset beach", "rainy city street at night", "floating ash land", "pa
 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 数据集路径
+# Hugging Face DTR 数据集路径（示例，若不可用请忽略）
 DTR_DATASET_PATTERN = "https://huggingface.co/datasets/X779/Danbooruwildcards/resolve/main/*DTR*.txt"
 
 # ========== 工具函数 ==========
-def load_candidates_from_files(files):
+def load_candidates_from_files(files, excluded_tags=None):
     """
-    从多个文件中加载候选项。
+    从多个文件中加载候选项，同时排除用户不想要的标签（精确匹配）。
     """
+    if excluded_tags is None:
+        excluded_tags = set()
     all_lines = []
     if files:
         for file in files:
             if isinstance(file, str):
+                # 说明是路径字符串
                 with open(file, "r", encoding="utf-8") as f:
-                    all_lines.extend([line.strip() for line in f if line.strip()])
+                    lines = [line.strip() for line in f if line.strip()]
+                    filtered = [l for l in lines if l not in excluded_tags]
+                    all_lines.extend(filtered)
+            else:
+                # 说明是一个上传的 file-like 对象
+                file_data = file.read().decode("utf-8", errors="ignore")
+                lines = [line.strip() for line in file_data.splitlines() if line.strip()]
+                filtered = [l for l in lines if l not in excluded_tags]
+                all_lines.extend(filtered)
     return all_lines
 
 def get_random_items(candidates, num_items=1):
@@ -39,14 +50,19 @@ def get_random_items(candidates, num_items=1):
     """
     return random.sample(candidates, min(num_items, len(candidates))) if candidates else []
 
-def load_dtr_from_huggingface():
+def load_dtr_from_huggingface(excluded_tags=None):
     """
-    从 Hugging Face 数据集中加载所有包含 "DTR" 的文件内容。
+    从 Hugging Face 数据集中加载所有包含 "DTR" 的文件内容，同时排除不需要的tag。
     """
+    if excluded_tags is None:
+        excluded_tags = set()
     try:
         response = requests.get(DTR_DATASET_PATTERN)
         response.raise_for_status()
-        return response.text.splitlines()
+        lines = response.text.splitlines()
+        # 只过滤精确匹配
+        lines = [l for l in lines if l not in excluded_tags]
+        return lines
     except Exception as e:
         print(f"Error loading DTR dataset: {e}")
         return []
@@ -60,11 +76,14 @@ 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."
 
-    tag_descriptions = "\n".join([f"{key}: {value}" for key, value in tags.items() if value])
+    # 将dict转成可读字符串
+    tag_descriptions = "\n".join([
+        f"{key}: {', '.join(value) if isinstance(value, list) else value}"
+        for key, value in tags.items() if value
+    ])
 
     try:
         client = OpenAI(api_key=api_key, base_url=base_url) if base_url else OpenAI(api_key=api_key)
-
         response = client.chat.completions.create(
             messages=[
                 {
@@ -72,7 +91,7 @@ def generate_natural_language_description(tags, api_key=None, base_url=None, mod
                     "content": (
                         "You are a creative assistant that generates detailed and imaginative scene descriptions for AI generation prompts. "
                         "Focus on the details provided and incorporate them into a cohesive narrative. "
-                        "Use at least three sentences but no more than five sentences"
+                        "Use at least three sentences but no more than five sentences."
                     ),
                 },
                 {
@@ -89,114 +108,227 @@ def generate_natural_language_description(tags, api_key=None, base_url=None, mod
 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,
-    artist_count, use_deepseek, deepseek_key
+    artist_count, use_deepseek, deepseek_key, user_custom_tags, excluded_tags
 ):
     """
     生成随机提示词和描述。
     """
-    actions = get_random_items(load_candidates_from_files([action_file]) if action_file else [], action_count)
-    styles = get_random_items(load_candidates_from_files([style_file]) if style_file else [], style_count)
-    artists = get_random_items(load_candidates_from_files(artist_files) if artist_files else [], artist_count)
-    characters = get_random_items(load_candidates_from_files(character_files) if character_files else [], 1)
-    dtr_candidates = get_random_items(load_dtr_from_huggingface() if dtr_enabled else [], 1)
+    # 处理排除 Tags（逗号分隔 -> 去重 set）
+    excluded_set = set(
+        [tag.strip() for tag in excluded_tags.split(",") if tag.strip()]
+    ) if excluded_tags else set()
+
+    # 从文件中加载可选 action、style、artist、character
+    actions = get_random_items(load_candidates_from_files([action_file], excluded_set) if action_file else [], action_count)
+    styles = get_random_items(load_candidates_from_files([style_file], excluded_set) if style_file else [], style_count)
+    artists = get_random_items(load_candidates_from_files(artist_files, excluded_set) if artist_files else [], artist_count)
+    characters = get_random_items(load_candidates_from_files(character_files, excluded_set) if character_files else [], 1)
+
+    # 处理 DTR
+    dtr_candidates = get_random_items(load_dtr_from_huggingface(excluded_set) if dtr_enabled else [], 1)
+
+    # 处理预设列表中的随机筛选
+    filtered_expressions = [e for e in EXPRESSIONS if e not in excluded_set]
+    filtered_items = [i for i in ITEMS if i not in excluded_set]
+    filtered_details = [d for d in OTHER_DETAILS if d not in excluded_set]
+    filtered_scenes = [s for s in SCENES if s not in excluded_set]
+    filtered_angles = [c for c in CAMERA_ANGLES if c not in excluded_set]
+    filtered_quality = [q for q in QUALITY_PROMPTS if q not in excluded_set]
+
+    # 随机抽取
+    random_expression = get_random_items(filtered_expressions, expression_count)
+    random_items = get_random_items(filtered_items, item_count)
+    random_details = get_random_items(filtered_details, detail_count)
+    random_scenes = get_random_items(filtered_scenes, scene_count)
+    random_angles = get_random_items(filtered_angles, angle_count)
+    random_quality = get_random_items(filtered_quality, quality_count)
 
     number_of_characters = ", ".join(selected_categories) if selected_categories else []
 
+    # 整理为字典
     tags = {
-        "number_of_characters": number_of_characters,
+        "number_of_characters": [number_of_characters] if number_of_characters else [],
         "character_name": characters,
-        "artist_prompt": f"(artist:{', '.join(artists)})",
+        "artist_prompt": [f"(artist:{', '.join(artists)})"] if artists else [],
         "style": styles,
-        "scene": get_random_items(SCENES, scene_count),
-        "camera_angle": get_random_items(CAMERA_ANGLES, angle_count),
+        "scene": random_scenes,
+        "camera_angle": random_angles,
         "action": actions,
-        "expression": get_random_items(EXPRESSIONS, expression_count),
-        "items": get_random_items(ITEMS, item_count),
-        "other_details": get_random_items(OTHER_DETAILS, detail_count),
-        "quality_prompts": get_random_items(QUALITY_PROMPTS, quality_count),
-        "dtr": dtr_candidates
+        "expression": random_expression,
+        "items": random_items,
+        "other_details": random_details,
+        "quality_prompts": random_quality,
+        "dtr": dtr_candidates,
     }
 
+    # 如果用户有自定义输入
+    if user_custom_tags.strip():
+        tags["custom_tags"] = [t.strip() for t in user_custom_tags.split(",") if t.strip()]
+
+    # 生成自然语言描述
     if use_deepseek:
         description = generate_natural_language_description(tags, api_key=deepseek_key, base_url="https://api.deepseek.com", model="deepseek-chat")
     else:
         description = generate_natural_language_description(tags, api_key=api_key)
 
-    tags_list = [item for sublist in tags.values() for item in (sublist if isinstance(sublist, list) else [sublist])]  # Flatten
-    unique_tags = list(dict.fromkeys(tags_list))
-    final_tags = ", ".join(unique_tags)
+    # 整理最终 Tags（flatten 并去重）
+    tags_list = []
+    for v in tags.values():
+        if isinstance(v, list):
+            tags_list.extend(v)
+        else:
+            tags_list.append(v)
+
+    # 去重保持顺序
+    seen = set()
+    final_tags_list = []
+    for t in tags_list:
+        if t not in seen and t:
+            seen.add(t)
+            final_tags_list.append(t)
+
+    final_tags = ", ".join(final_tags_list)
+
+    # 默认 Combined = Tags + Description
     combined_output = f"{final_tags}\n\n{description}"
     return final_tags, description, combined_output
 
+# ========== Favorite 相关函数 ==========
+def add_to_favorites(combined_output, current_favorites):
+    """
+    将当前生成的 combined_output 添加到收藏列表中（最多存 3 条）。
+    """
+    # current_favorites 是一个列表
+    current_favorites.append(combined_output)
+    # 如果超过3条，移除最早的一条
+    if len(current_favorites) > 3:
+        current_favorites.pop(0)
+    # 格式化输出
+    favorites_text = "\n\n".join(
+        [f"[Favorite {i+1}]\n{fav}" for i, fav in enumerate(current_favorites)]
+    )
+    return favorites_text, current_favorites
+
 # ========== Gradio 界面 ==========
 def gradio_interface():
     """
     定义 Gradio 应用界面。
     """
     with gr.Blocks() as demo:
-        gr.Markdown("## Random Prompt Generator with Adjustable Tag Counts")
+        gr.Markdown("## Random Prompt Generator with Adjustable Tag Counts (Enhanced)")
 
-        api_key_input = gr.Textbox(
-            label="Enter your OpenAI API Key (Optional)",
-            placeholder="sk-...",
-            type="password"
-        )
+        # 用于存储收藏内容的状态（最多缓存3条）
+        favorites_state = gr.State([])
 
-        deepseek_key_input = gr.Textbox(
-            label="Enter your DeepSeek API Key (Optional)",
-            placeholder="sk-...",
-            type="password"
-        )
+        with gr.Row():
+            # 左侧：文件上传、参数选择、排除/自定义输入
+            with gr.Column(scale=1):
+                api_key_input = gr.Textbox(
+                    label="OpenAI API Key (可选)",
+                    placeholder="sk-...",
+                    type="password"
+                )
 
-        use_deepseek = gr.Checkbox(label="Use DeepSeek API - 用DeepSeek别忘了勾选这个")
+                deepseek_key_input = gr.Textbox(
+                    label="DeepSeek API Key (可选)",
+                    placeholder="sk-...",
+                    type="password"
+                )
 
-        with gr.Row():
-            action_file = gr.File(label="Upload Action File (Optional)", file_types=[".txt"])
-            style_file = gr.File(label="Upload Style File (Optional)", file_types=[".txt"])
+                use_deepseek = gr.Checkbox(label="Use DeepSeek API")
 
-        with gr.Row():
-            artist_files = gr.Files(label="Upload Artist Files (Multiple Allowed)", file_types=[".txt"])
-            character_files = gr.Files(label="Upload Character Files (Multiple Allowed)", file_types=[".txt"])
+                dtr_enabled = gr.Checkbox(label="Enable DTR (如不可用请勿勾选)")
 
-        dtr_enabled = gr.Checkbox(label="Enable DTR - 当前不可用2025-01-12")
+                # 上传文件部分
+                with gr.Box():
+                    gr.Markdown("**上传文件 (可选)：**")
+                    action_file = gr.File(label="Action File", file_types=[".txt"])
+                    style_file = gr.File(label="Style File", file_types=[".txt"])
+                    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 (Optional)"
-        )
+                # 选择角色类型
+                selected_categories = gr.CheckboxGroup(
+                    ["1boy", "1girl", "furry", "mecha", "fantasy monster", "animal", "still life"],
+                    label="Choose Character Categories"
+                )
 
-        with gr.Row():
-            expression_count = gr.Slider(label="Number of Expressions", minimum=0, maximum=10, step=1, value=1)
-            item_count = gr.Slider(label="Number of Items", minimum=0, maximum=10, step=1, value=1)
-            detail_count = gr.Slider(label="Number of Other Details", minimum=0, maximum=10, step=1, value=1)
-            scene_count = gr.Slider(label="Number of Scenes", minimum=0, maximum=10, step=1, value=1)
+                # 输入排除和自定义
+                excluded_tags = gr.Textbox(
+                    label="排除 Tags (逗号分隔)",
+                    placeholder="如：angry, sword"
+                )
+                user_custom_tags = gr.Textbox(
+                    label="自定义附加 Tags (逗号分隔)",
+                    placeholder="如：glowing eyes, giant wings"
+                )
 
-        with gr.Row():
-            angle_count = gr.Slider(label="Number of Camera Angles", minimum=0, maximum=10, step=1, value=1)
-            quality_count = gr.Slider(label="Number of Quality Prompts", minimum=0, maximum=10, step=1, value=1)
-            action_count = gr.Slider(label="Number of Actions", minimum=1, maximum=10, step=1, value=1)
-            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.Box():
+                    gr.Markdown("**随机数量设置：**")
+                    expression_count = gr.Slider(label="Number of Expressions", minimum=0, maximum=10, step=1, value=1)
+                    item_count = gr.Slider(label="Number of Items", minimum=0, maximum=10, step=1, value=1)
+                    detail_count = gr.Slider(label="Number of Other Details", minimum=0, maximum=10, step=1, value=1)
+                    scene_count = gr.Slider(label="Number of Scenes", minimum=0, maximum=10, step=1, value=1)
+                    angle_count = gr.Slider(label="Number of Camera Angles", minimum=0, maximum=10, step=1, value=1)
+                    quality_count = gr.Slider(label="Number of Quality Prompts", minimum=0, maximum=10, step=1, value=1)
+                    action_count = gr.Slider(label="Number of Actions", minimum=1, maximum=10, step=1, value=1)
+                    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.Row():
-            tags_output = gr.Textbox(label="Generated Tags")
-            description_output = gr.Textbox(label="Generated Description")
-            combined_output = gr.Textbox(label="Combined Output: Tags + Description")
-
-        generate_button = gr.Button("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
-            ],
-            outputs=[tags_output, description_output, combined_output],
-        )
+            # 右侧：生成按钮 + 生成结果 + 收藏
+            with gr.Column(scale=2):
+                generate_button = gr.Button("Generate Prompt", variant="primary")
+
+                # 生成的结果（可编辑）
+                tags_output = gr.Textbox(
+                    label="Generated Tags",
+                    placeholder="等待生成...",
+                    lines=4,
+                    interactive=True
+                )
+                description_output = gr.Textbox(
+                    label="Generated Description",
+                    placeholder="等待生成...",
+                    lines=4,
+                    interactive=True
+                )
+                combined_output = gr.Textbox(
+                    label="Combined Output: Tags + Description",
+                    placeholder="等待生成...",
+                    lines=6
+                )
+
+                with gr.Row():
+                    # 收藏操作
+                    favorite_button = gr.Button("收藏本次结果")
+                    favorites_box = gr.Textbox(
+                        label="收藏夹 (最多 3 条)",
+                        placeholder="暂无收藏",
+                        lines=6
+                    )
+
+                # 点击生成按钮
+                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
+                    ],
+                    outputs=[tags_output, description_output, combined_output],
+                )
+
+                # 收藏按钮点击事件
+                favorite_button.click(
+                    fn=add_to_favorites,
+                    inputs=[combined_output, favorites_state],
+                    outputs=[favorites_box, favorites_state],
+                )
 
+        # 整个界面启动
     return demo
 
-# 启动 Gradio 应用
+# 启动 Gradio 应用（在 Hugging Face Spaces 上只需要保留最后的 launch 即可）
 if __name__ == "__main__":
-    gradio_interface().launch(share=True)
+    gradio_interface().launch()