Update app.py

app.py

@@ -2,6 +2,7 @@ import gradio as gr
 import random
 import glob
 import os
+import requests
 from openai import OpenAI
 from dotenv import load_dotenv
 
@@ -29,28 +30,23 @@ def load_candidates_from_files(files):
                     all_lines.extend([line.strip() for line in f if line.strip()])
     return all_lines
 
-def get_random_item(candidates):
+def get_random_items(candidates, num_items=1):
     """
-    随机选取候选项。
+    从候选项中随机选取指定数量的选项。
     """
-    return random.choice(candidates) if candidates else ""
+    return random.sample(candidates, min(num_items, len(candidates))) if candidates else []
 
-def load_dtr_from_directory(directory=".", pattern="*DTR*"):
+def load_dtr_from_huggingface(dataset_url):
     """
-    从指定目录中加载所有包含特定模式的文件内容。
-    :param directory: 目标目录，默认为当前目录
-    :param pattern: 匹配文件名的模式，默认是包含 "DTR" 的文件
-    :return: 文件内容的列表
+    从 Hugging Face 数据集中加载 DTR 内容。
     """
-    dtr_candidates = []
     try:
-        files = glob.glob(os.path.join(directory, pattern))
-        for file in files:
-            with open(file, "r", encoding="utf-8") as f:
-                dtr_candidates.extend([line.strip() for line in f if line.strip()])
+        response = requests.get(dataset_url)
+        response.raise_for_status()
+        return response.text.splitlines()
     except Exception as e:
-        print(f"Error loading DTR files: {e}")
-    return dtr_candidates
+        print(f"Error loading DTR dataset: {e}")
+        return []
 
 def generate_natural_language_description(tags, api_key=None):
     """
@@ -87,7 +83,10 @@ def generate_natural_language_description(tags, api_key=None):
     except Exception as e:
         return f"GPT generation failed. Error: {e}"
 
-def generate_prompt(action_file, style_file, artist_files, character_files, dtr_directory, api_key, selected_categories):
+def generate_prompt(
+    action_file, style_file, artist_files, character_files, dtr_enabled, dtr_url, api_key, selected_categories,
+    expression_count, item_count, detail_count, scene_count, angle_count, quality_count
+):
     """
     生成随机提示词和描述。
     """
@@ -95,28 +94,28 @@ def generate_prompt(action_file, style_file, artist_files, character_files, dtr_
     styles = load_candidates_from_files([style_file]) if style_file else []
     artists = load_candidates_from_files(artist_files) if artist_files else []
     characters = load_candidates_from_files(character_files) if character_files else []
-    dtr_candidates = load_dtr_from_directory(dtr_directory) if dtr_directory else []
+    dtr_candidates = load_dtr_from_huggingface(dtr_url) if dtr_enabled else []
 
     number_of_characters = ", ".join(selected_categories) if selected_categories else random.choice(["1girl", "1boy"])
 
     tags = {
         "number_of_characters": number_of_characters,
-        "character_name": get_random_item(characters),
-        "artist_prompt": f"(artist:{get_random_item(artists)})",
-        "style": get_random_item(styles),
-        "scene": get_random_item(SCENES),
-        "camera_angle": get_random_item(CAMERA_ANGLES),
-        "action": get_random_item(actions),
-        "expression": get_random_item(EXPRESSIONS),
-        "items": get_random_item(ITEMS),
-        "other_details": get_random_item(OTHER_DETAILS),
-        "quality_prompts": get_random_item(QUALITY_PROMPTS),
-        "dtr": get_random_item(dtr_candidates)
+        "character_name": get_random_items(characters, 1),
+        "artist_prompt": f"(artist:{get_random_items(artists, 1)})",
+        "style": get_random_items(styles, 1),
+        "scene": get_random_items(SCENES, scene_count),
+        "camera_angle": get_random_items(CAMERA_ANGLES, angle_count),
+        "action": get_random_items(actions, 1),
+        "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": get_random_items(dtr_candidates, 1)
     }
 
     description = generate_natural_language_description(tags, api_key)
 
-    tags_list = [value for value in tags.values() if value]
+    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)
     combined_output = f"{final_tags}\n\n{description}"
@@ -128,7 +127,7 @@ def gradio_interface():
     定义 Gradio 应用界面。
     """
     with gr.Blocks() as demo:
-        gr.Markdown("## Random Prompt Generator with User-Provided GPT API Key")
+        gr.Markdown("## Random Prompt Generator with Adjustable Tag Counts")
 
         api_key_input = gr.Textbox(
             label="Enter your OpenAI API Key (Optional)",
@@ -145,12 +144,29 @@ def gradio_interface():
             character_files = gr.Files(label="Upload Character Files (Multiple Allowed)", file_types=[".txt"])
 
         dtr_enabled = gr.Checkbox(label="Enable DTR Directory")
+        dtr_url = gr.Textbox(
+            label="Hugging Face DTR Dataset URL",
+            placeholder="https://huggingface.co/datasets/X779/Danbooruwildcards/tree/main",
+            value="https://huggingface.co/datasets/X779/Danbooruwildcards/tree/main"
+        )
 
         selected_categories = gr.CheckboxGroup(
             ["1boy", "1girl", "furry", "mecha", "fantasy monster", "animal", "still life"],
             label="Choose Character Categories (Optional)"
         )
 
+        with gr.Row():
+            expression_count = gr.Slider(label="Number of Expressions", minimum=1, maximum=5, step=1, value=1)
+            item_count = gr.Slider(label="Number of Items", minimum=1, maximum=5, step=1, value=1)
+
+        with gr.Row():
+            detail_count = gr.Slider(label="Number of Other Details", minimum=1, maximum=5, step=1, value=1)
+            scene_count = gr.Slider(label="Number of Scenes", minimum=1, maximum=5, step=1, value=1)
+
+        with gr.Row():
+            angle_count = gr.Slider(label="Number of Camera Angles", minimum=1, maximum=5, step=1, value=1)
+            quality_count = gr.Slider(label="Number of Quality Prompts", minimum=1, maximum=5, step=1, value=1)
+
         with gr.Row():
             tags_output = gr.Textbox(label="Generated Tags")
             description_output = gr.Textbox(label="Generated Description")
@@ -161,7 +177,8 @@ def gradio_interface():
         generate_button.click(
             generate_prompt,
             inputs=[
-                action_file, style_file, artist_files, character_files, dtr_enabled, api_key_input, selected_categories
+                action_file, style_file, artist_files, character_files, dtr_enabled, dtr_url, api_key_input, selected_categories,
+                expression_count, item_count, detail_count, scene_count, angle_count, quality_count
             ],
             outputs=[tags_output, description_output, combined_output],
         )