Update app.py

app.py

@@ -10,12 +10,15 @@ from dotenv import load_dotenv
 load_dotenv()
 
 # ========== 默认选项和数据 ==========
-EXPRESSIONS = ["smiling", "determined", "surprised", "serene"]
-ITEMS = ["magic wand", "sword", "flower", "book of spells", "ancient scroll", "music instrument"]
-OTHER_DETAILS = ["sparkles", "magical aura", "lens flare", "fireworks in the background"]
-SCENES = ["sunset beach", "rainy city street at night", "fantasy forest with glowing mushrooms", "futuristic skyline at dawn"]
-CAMERA_ANGLES = ["low-angle shot", "close-up shot", "bird's-eye view", "wide-angle shot"]
-QUALITY_PROMPTS = ["8k", "ultra-realistic", "high detail", "cinematic lighting", "award-winning", "masterpiece"]
+EXPRESSIONS = ["smiling", "determined", "surprised", "serene", "laughing", "angry", "pensive", "confident"]
+ITEMS = ["magic wand", "sword", "flower", "book of spells", "ancient scroll", "music instrument", "shield", "dagger", "staff", "potion"]
+OTHER_DETAILS = ["sparkles", "magical aura", "lens flare", "fireworks in the background", "smoke effects", "light trails", "falling leaves", "glowing embers"]
+SCENES = ["sunset beach", "rainy city street at night", "fantasy forest with glowing mushrooms", "futuristic skyline at dawn", "abandoned castle", "snowy mountain peak", "desert ruins", "underwater city"]
+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"]
+QUALITY_PROMPTS = ["cinematic lighting", "award-winning", "masterpiece", "vivid colors", "high dynamic range", "immersive", "studio quality", "fine art", "dreamlike", "8K", "HD", "high quality", "best quality"]
+
+# Hugging Face DTR 数据集路径
+DTR_DATASET_PATTERN = "https://huggingface.co/datasets/X779/Danbooruwildcards/resolve/main/*DTR*.txt"
 
 # ========== 工具函数 ==========
 def load_candidates_from_files(files):
@@ -36,12 +39,12 @@ 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(dataset_url):
+def load_dtr_from_huggingface():
     """
-    从 Hugging Face 数据集中加载 DTR 内容。
+    从 Hugging Face 数据集中加载所有包含 "DTR" 的文件内容。
     """
     try:
-        response = requests.get(dataset_url)
+        response = requests.get(DTR_DATASET_PATTERN)
         response.raise_for_status()
         return response.text.splitlines()
     except Exception as e:
@@ -84,33 +87,33 @@ def generate_natural_language_description(tags, api_key=None):
         return f"GPT generation failed. Error: {e}"
 
 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
+    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
 ):
     """
     生成随机提示词和描述。
     """
-    actions = load_candidates_from_files([action_file]) if action_file else []
-    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_huggingface(dtr_url) if dtr_enabled else []
+    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 [], 1)
+    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)
 
     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_items(characters, 1),
-        "artist_prompt": f"(artist:{get_random_items(artists, 1)})",
-        "style": get_random_items(styles, 1),
+        "character_name": characters,
+        "artist_prompt": f"(artist:{artists})",
+        "style": styles,
         "scene": get_random_items(SCENES, scene_count),
         "camera_angle": get_random_items(CAMERA_ANGLES, angle_count),
-        "action": get_random_items(actions, 1),
+        "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": get_random_items(dtr_candidates, 1)
+        "dtr": dtr_candidates
     }
 
     description = generate_natural_language_description(tags, api_key)
@@ -143,12 +146,7 @@ def gradio_interface():
             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 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"
-        )
+        dtr_enabled = gr.Checkbox(label="Enable DTR")
 
         selected_categories = gr.CheckboxGroup(
             ["1boy", "1girl", "furry", "mecha", "fantasy monster", "animal", "still life"],
@@ -158,14 +156,14 @@ def gradio_interface():
         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)
+            action_count = gr.Slider(label="Number of Actions", minimum=1, maximum=5, step=1, value=1)
+            style_count = gr.Slider(label="Number of Styles", minimum=1, maximum=5, step=1, value=1)
 
         with gr.Row():
             tags_output = gr.Textbox(label="Generated Tags")
@@ -177,8 +175,8 @@ def gradio_interface():
         generate_button.click(
             generate_prompt,
             inputs=[
-                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
+                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
             ],
             outputs=[tags_output, description_output, combined_output],
         )