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):
    """
    从多个文件中加载候选项。
    """
    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()])
    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():
    """
    从 Hugging Face 数据集中加载所有包含 "DTR" 的文件内容。
    """
    try:
        response = requests.get(DTR_DATASET_PATTERN)
        response.raise_for_status()
        return response.text.splitlines()
    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."

    tag_descriptions = "\n".join([f"{key}: {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
):
    """
    生成随机提示词和描述。
    """
    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)

    number_of_characters = ", ".join(selected_categories) if selected_categories else []

    tags = {
        "number_of_characters": number_of_characters,
        "character_name": characters,
        "artist_prompt": f"(artist:{', '.join(artists)})",
        "style": styles,
        "scene": get_random_items(SCENES, scene_count),
        "camera_angle": get_random_items(CAMERA_ANGLES, angle_count),
        "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
    }

    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)
    combined_output = f"{final_tags}\n\n{description}"
    return final_tags, description, combined_output

# ========== Gradio 界面 ==========
def gradio_interface():
    """
    定义 Gradio 应用界面。
    """
    with gr.Blocks() as demo:
        gr.Markdown("## Random Prompt Generator with Adjustable Tag Counts")

        api_key_input = gr.Textbox(
            label="Enter your OpenAI API Key (Optional)",
            placeholder="sk-...",
            type="password"
        )

        deepseek_key_input = gr.Textbox(
            label="Enter your DeepSeek API Key (Optional)",
            placeholder="sk-...",
            type="password"
        )

        use_deepseek = gr.Checkbox(label="Use DeepSeek API - 用DeepSeek别忘了勾选这个")

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

        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 - 当前不可用2025-01-12")

        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=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)

        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.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],
        )

    return demo

# 启动 Gradio 应用
if __name__ == "__main__":
    gradio_interface().launch(share=True)