app.py

import gradio as gr
import random
import glob
import os
import requests
from openai import OpenAI
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"]

# Hugging Face DTR 数据集路径（示例，若不可用请忽略）
DTR_DATASET_PATTERN = "https://huggingface.co/datasets/X779/Danbooruwildcards/resolve/main/*DTR*.txt"


# ========== 工具函数 ==========
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:
                    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):
    """
    从候选项中随机选取指定数量的选项。
    """
    return random.sample(candidates, min(num_items, len(candidates))) if candidates else []

def load_dtr_from_huggingface(excluded_tags=None):
    """
    从 Hugging Face 数据集中加载所有包含 "DTR" 的文件内容，同时排除不需要的tag。
    """
    if excluded_tags is None:
        excluded_tags = set()
    try:
        response = requests.get(DTR_DATASET_PATTERN)
        response.raise_for_status()
        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 []

def generate_natural_language_description(tags, api_key=None, base_url=None, model="gpt-4o"):
    """
    使用 OpenAI GPT 或 DeepSeek API 生成自然语言描述。
    """
    if not api_key:
        api_key = os.getenv("OPENAI_API_KEY")
    if not api_key:
        return "Error: No API Key provided and none found in environment variables."

    # 将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=[
                {
                    "role": "system",
                    "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."
                    ),
                },
                {
                    "role": "user",
                    "content": f"Here are the tags and details:\n{tag_descriptions}\nPlease generate a vivid, imaginative scene description.",
                },
            ],
            model=model,
        )
        return response.choices[0].message.content.strip()
    except Exception as e:
        return f"GPT generation failed. Error: {e}"

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, user_custom_tags, excluded_tags
):
    """
    生成随机提示词和描述。
    """
    # 处理排除 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] if number_of_characters else [],
        "character_name": characters,
        "artist_prompt": [f"(artist:{', '.join(artists)})"] if artists else [],
        "style": styles,
        "scene": random_scenes,
        "camera_angle": random_angles,
        "action": actions,
        "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（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 (Enhanced)")

        # 用于存储收藏内容的状态（最多缓存3条）
        favorites_state = gr.State([])

        with gr.Row():
            # 左侧：文件上传、参数选择、排除/自定义输入
            with gr.Column(scale=1):
                api_key_input = gr.Textbox(
                    label="OpenAI API Key (可选)",
                    placeholder="sk-...",
                    type="password"
                )

                deepseek_key_input = gr.Textbox(
                    label="DeepSeek API Key (可选)",
                    placeholder="sk-...",
                    type="password"
                )

                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"])
                    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"
                )

                # 输入排除和自定义
                excluded_tags = gr.Textbox(
                    label="排除 Tags (逗号分隔)",
                    placeholder="如：angry, sword"
                )
                user_custom_tags = gr.Textbox(
                    label="自定义附加 Tags (逗号分隔)",
                    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)
                    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.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 应用（在 Hugging Face Spaces 上只需要保留最后的 launch 即可）
if __name__ == "__main__":
    gradio_interface().launch()