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",
    "amused", "excited", "anxious", "confused", "curious",
    "delighted", "disappointed", "elated", "embarrassed",
    "envious", "fearful", "frustrated", "hopeful", "horrified",
    "inspired", "jealous", "melancholic", "nervous", "playful",
    "relieved", "remorseful", "skeptical", "terrified", "thrilled",
    "uneasy", "wistful", "zany", "neutral", "expressionless",
    "sleepy", "tired", "relaxed", "drooling", "sick", "nauseated",
    "sneezing", "hot", "cold", "woozy", "dizzy", 
    "(๑•̀ㅂ•́)و✧", "(╯°□°）╯︵ ┻━┻", "(✧ω✧)", "(・_・;)", 
    "(｡•́︿•̀｡)", "(≧◡≦)", "(╥﹏╥)", "(￣▽￣)", 
    "(ಥ_ಥ)", "(⊙_☉)", "(¬_¬)", "(°ロ°)!", "(≧︿≦)", 
    "(´•ω•｀)", "(°ー°〃)", "(^人^)", "(╯︵╰,)", "(ಥ﹏ಥ)", 
    "(´･_･`)", "(ﾉ´ｰ`)ﾉ", "(ﾉಥ益ಥ)ﾉ", "(｡♥‿♥｡)", "(づ￣ ³￣)づ",
    "(☉｡☉)!", "(>_<)", "(☆ω☆)", "(x_x)", "o(TヘTo)", "(︶︹︺)",
    "(⌒_⌒;)", "(☍﹏⁰)｡", "(´･_･)", "(⊙﹏⊙)", "(｡-_-｡)", "(¬‿¬)",
    "(°∀°)", "(◕‿◕)", "(T_T)", "(￣︿￣)", "(>人<)", "(つД`)ノ",
    "(｀・ω・´)", "(๑´ڡ`๑)", "(✿╹◡╹)", "(´ε｀ )♡", "(•̀ᴗ•́)و",
    "(∩｀-´)⊃━☆ﾟ.*･｡ﾟ", "(☞ﾟヮﾟ)☞", "(ﾉ≧ڡ≦)!", "(´･ω･`)ﾉ", "(⁎⁍̴̛ᴗ⁍̴̛⁎)",
    "(╬ಠ益ಠ)", "(ノಠ益ಠ)ノ彡┻━┻", "(ง'̀-'́)ง", "(´｡• ω •｡`) ♡"
]
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", "enchanted cloak", "mystic orb",
    "charred tome", "golden chalice", "feather quill", "obsidian blade",
    "healing herb", "celestial compass", "shimmering veil", "phoenix feather",
    "rune stone", "dragon scale", "mirror of truth", "key of destiny",
    "shadow dagger", "firestone", "glimmering necklace", "moonlit tiara",
    "elixir of vitality", "serpent bracelet", "orb of whispers", "ivy circlet",
    "ember crown", "storm gauntlets", "twilight boots", "ethereal shawl",
    "guardian pendant", "pearl of wisdom", "arcane hourglass", "binding chain",
    "gem of clarity", "silk gloves", "wings of ascension", "crystal shard",
    "sapphire brooch", "mask of shadows", "ashwood cane", "silver flute",
    "chime of echoes", "star chart", "jade pendant", "hourglass of time",
    "frosted gem", "vial of starlight", "scroll of secrets", "golden anklet",
    "radiant crown", "veil of illusions", "ring of eternity", "soulbound locket",
    "enchanted map", "mystic key", "ancient coin", "wizard's journal",
    "magic ink", "sorcerer's robe", "elven boots", "dwarven hammer",
    "griffin feather", "fairy dust", "necromancer's book", "ranger's bow",
    "bard's lute", "monk's beads", "thief's dagger", "alchemist's flask",
    "summoner's tome", "elemental gem", "timekeeper's watch", "portal key",
    "invisibility cloak", "levitation charm", "mind control orb", "fire-breathing potion",
    "earthquake drum", "storm summoning horn", "sunstone", "moonstone",
    "stardust vial", "dreamcatcher", "truth serum", "luck charm",
    "fortune dice", "destiny thread", "life potion", "eternal torch",
    "necklace", "bracelet", "brooch", "tie clip", "cufflinks", "belt",
    "scarf", "shawl", "hat", "hairpin", "hairband", "gloves", "waist chain",
    "keychain", "bag charm", "badge", "armband", "lapel pin", "sash",
    "beaded fringe", "embroidered patch", "rhinestones", "pearls", "chains",
    "suspenders", "garter belts", "stockings", "mittens", "masks", "veils",
    "handbags", "clutches", "backpacks", "satchels", "messenger bags",
    "wristwatches", "pocket watches", "sunglasses", "spectacles", "chokers",
    "lockets", "pendants", "talismans", "charms", "ankle bracelets", "toe rings"
]
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",
    "ripples on water", "dust motes in sunlight", "swirling fog", "glittering frost",
    "reflections on surfaces", "flickering shadows", "falling snowflakes",
    "twinkling stars", "crescent moonlight", "glowing fireflies", "soft cloud wisps",
    "subtle rain", "scattered feathers", "wind-blown sand", "cracks in the ground",
    "waves crashing in the distance", "glowing mushrooms", "hanging vines",
    "faint rainbows", "drifting ash", "rising bubbles", "aurora in the sky",
    "shattered glass fragments", "falling cherry blossoms", "sunbeams breaking through clouds",
    "water droplets on leaves", "wisps of smoke", "dimly glowing lanterns",
    "undulating waves of grass", "sparkling dewdrops", "glistening cobwebs",
    "golden hour light", "tiny glowing orbs", "vapor trails in the sky",
    "distant silhouettes of birds", "slowly moving shadows", "softly pulsing glows",
    "gently falling powder snow", "mirages on the horizon", "waving tree branches",
    "radiant beams through fog", "scattered golden leaves", "softly glowing fungi",
    "rippling heatwaves", "glowing lines on surfaces",
    # 叙事辅助细节
    "ancient carvings on walls", "worn-out banners fluttering", "cracked stained glass",
    "fallen statues covered in vines", "ceremonial fires flickering", "abandoned tools scattered",
    "forgotten books on dusty shelves", "hidden inscriptions glowing faintly",
    "collapsed bridges in the distance", "ancient scripts floating in midair",
    "trails of footsteps in the snow", "discarded musical instruments", "half-buried artifacts",
    "melting candles on altars", "phantom figures fading in the background",
    "reflections of unseen movements in water", "shifting sands forming patterns",
    "layered ruins creating depth", "shattered pottery", "crumbling pillars",
    "weathered doorways", "dried riverbeds", "overgrown fences", "scattered coins",
    "glimmering gemstones in cracks", "rusted chains", "torn flags",
    "broken shields lying in the grass", "ancient bridges covered in moss",
    "swords embedded in stone", "faded murals on crumbling walls",
    "mysterious glowing symbols on the ground", "abandoned helmets half-buried in sand",
    "fossilized remains in rock walls", "golden inscriptions on blackened stone",
    "vines entwined with ancient statues", "weathered chests hidden under roots",
    "bone fragments scattered across ruins", "broken clocks frozen in time",
    "lichen-covered stones marking old pathways", "cursed objects glowing faintly in shadows",
    "forgotten wells with stagnant water", "ashes scattered across altars",
    "chains hanging from ruined walls", "rusted weapons embedded in tree trunks",
    "ancient seals broken on stone tablets", "sunken ruins peeking through murky waters",
    "melted wax pooling on ritual circles", "bird nests in abandoned helmets",
    "vines overtaking crumbling archways", "bloodstains faded into stone",
    "remains of campfires surrounded by bones", "fragments of ancient maps on walls",
    "forgotten altars surrounded by offerings", "spiral patterns carved into the ground",
    "wind chimes tangled in overgrowth", "scattered gemstones sparkling in the dirt",
    # 科幻背景细节
    "hovering drones patrolling the area", "holographic billboards flickering",
    "neon signs reflecting on wet pavement", "futuristic skyscrapers piercing the clouds",
    "flying cars zipping through the sky", "robotic vendors at street corners",
    "cybernetic enhancements visible on pedestrians", "digital advertisements projected in midair",
    "energy shields protecting buildings", "automated cleaning bots sweeping the streets",
    "glowing circuit patterns on walls", "transparent bridges connecting towers",
    "artificial intelligence interfaces displayed publicly", "bioluminescent plants lining walkways",
    "virtual reality kiosks with users immersed", "mechanical limbs discarded in alleys",
    "surveillance cameras with blinking lights", "underground labs emitting strange glows",
    "spaceships landing on rooftop pads", "androids interacting with humans seamlessly",
    "laser engravings on metallic surfaces", "plasma screens displaying news feeds",
    "synthetic food stands with diverse offerings", "hoverboards parked near entrances",
    "genetic modification clinics with bright signage", "quantum computing hubs glowing faintly",
    "fusion reactors humming in the distance", "alien flora glowing in sterile environments",
    "rusted space equipment scattered on barren ground", "forcefields shimmering in the air",
    "futuristic statues commemorating unknown heroes", "orbital stations visible in the sky",
    "cracked stasis pods leaking mist", "discarded exoskeletons in shadowy corners",
    "dimly lit corridors with blinking control panels", "holographic maps rotating slowly",
    "time-worn spacecraft hulls covered in scratches", "energy conduits glowing along walls"
]
SCENES = [
    # 自然场景
    "sunset beach", "rainy city street at night", "high blue sky",
    "snowy mountain peak", "desert ruins", "enchanted meadow",
    "misty forest trail", "hidden valley surrounded by cliffs",
    "raging waterfall in a lush jungle", "golden wheat fields at sunset",
    "stormy coastline with crashing waves", "dense rainforest with towering trees",
    "wind-swept plains under an overcast sky", "frozen tundra stretching endlessly",
    "emerald lake reflecting the mountains", "serene meadow under starlit skies",
    # 奇幻场景
    "fantasy forest with glowing mushrooms", "particles magic world",
    "floating ash land", "abandoned castle", "underwater city",
    "haunted mansion", "glacial cavern", "floating islands above a magical sea",
    "ancient ruins glowing with runes", "crystal cave with shimmering walls",
    "celestial palace above the clouds", "dark enchanted forest with eerie lights",
    "golden desert oasis with sparkling waters", "volcanic landscape with rivers of lava",
    "towering spires of an elven city", "hidden fae village among massive flowers",
    "ancient temple shrouded in mist",
    # 科幻场景
    "futuristic skyline at dawn", "steampunk marketplace",
    "orbital station overlooking a planet", "neon-lit cyberpunk cityscape",
    "crumbling space station", "alien planet with multiple moons",
    "terraforming facility on barren land", "gravity-defying sci-fi structures",
    "bioluminescent forest on an alien world", "megastructure city layered with lights",
    "robot factory with moving assembly lines", "underground research lab glowing faintly",
    "time-warped city merging eras", "abandoned alien settlement in the desert",
    "holographic gardens in the void", "massive starship docks bustling with activity",
    # 废墟与遗迹场景
    "ancient battlefield with broken weapons", "collapsed bridges over a canyon",
    "abandoned cathedral overgrown with vines", "forgotten library buried in sand",
    "crumbling statues in a misty valley", "deserted mining town under a red sky",
    "swampy ruins with submerged pillars", "eroded city reclaimed by nature",
    "old lighthouse on rocky cliffs", "sunken ship graveyard in shallow seas",
    "ruined amphitheater under moonlight", "forgotten monastery at a mountain's edge",
    "overgrown garden in a dilapidated mansion", "war-torn city with smoking rubble",
    "ancient crypt deep underground",
    # 情绪化场景
    "peaceful village at dawn", "stormy mountain pass with lightning",
    "abandoned carnival under a cloudy sky", "twilight forest bathed in golden light",
    "melancholic city square after rain", "sunlit cathedral with stained glass reflections",
    "empty train station at midnight", "serene pond surrounded by blossoms",
    "vibrant marketplace during a festival", "dim tavern with flickering candles",
    "lonely pier stretching into foggy waters", "quiet cemetery under crescent moonlight",
    "hidden spring surrounded by mossy rocks", "eerie swamp with glowing will-o'-the-wisps",
    "ancient battlefield under a blood-red sky",
    # 室内场景
    "opulent throne room in a grand castle", "dimly lit library with towering bookshelves",
    "abandoned laboratory with broken equipment", "cozy cabin with a roaring fireplace",
    "ornate ballroom with glittering chandeliers", "sprawling dining hall with long tables",
    "underground bunker with industrial walls", "luxurious palace bedroom draped in silk",
    "steampunk workshop filled with gears and tools", "mystical shrine surrounded by candles",
    "underground dungeon with damp stone walls", "hidden study behind a bookshelf",
    "sunlit greenhouse filled with exotic plants", "control room glowing with monitors",
    "train car interior with vintage decor", "haunted attic with cobwebs and trunks",
    "grand theater with red velvet curtains", "rustic kitchen with copper pots",
    "modern office with glass walls", "abandoned asylum with graffiti-covered walls",
    "elegant tea room with porcelain settings", "high-tech lab with robotic arms",
    "children’s playroom with scattered toys", "ancient council chamber with stone seats",
    "dim tavern with wooden furniture", "artist's studio with half-finished paintings",
    "futuristic sleeping pod chamber", "neon-lit arcade with glowing consoles",
    "classroom with chalkboards and empty desks", "temple hall with intricate carvings"
]

CAMERA_ANGLES = [
    "low-angle shot", "high-angle shot", "normal eye-level shot", 
    "close-up shot", "medium shot", "wide-angle shot", 
    "over-the-shoulder shot", "bird's-eye view", "worm's-eye view", 
    "extreme close-up", "panoramic view", "dynamic tracking shot", 
    "fisheye view", "point-of-view (POV) shot", "from behind shot", 
    "from top shot", "from below shot", "from the side shot", 
    "over-the-head shot", "reaction shot", "through a window shot", 
    "mirror reflection shot", "split-screen shot", 
    "silhouette framing shot", "foreground framing shot", 
    "Dutch angle", "frame within a frame", "obscured view through foliage", 
    "reflected surface view", "shadow-based composition shot", 
    "rotating 360-degree shot", "slow-motion close-up", 
    "freeze-frame shot", "fast zoom-out shot", 
    "tilt-up shot", "tilt-down shot", "panning shot", 
    "crane shot", "dolly zoom", "whip pan shot", 
    "impact tracking shot", "environment dominant framing", 
    "symmetrical center focus", "off-center dynamic composition"
]

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-4"):
    """
    使用 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}"


# ========== 核心函数：随机生成 prompt ==========
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


# ========== 部分更新：只根据用户修改后的 tags_text 生成新的描述和合并输出 ==========
def update_description(tags_text, api_key, use_deepseek, deepseek_key):
    """
    只根据用户提供的 tags_text 生成描述和合并输出。
    不再重新随机抽取，以免破坏用户手动修改过的 Tags。
    """
    if not api_key and not deepseek_key:
        # 没有提供任意可用 API Key
        return "(No API Key provided)", f"{tags_text}\n\n(No API Key provided)"

    # 构造给 GPT 的 prompt
    user_prompt = (
        "You are a creative assistant that generates detailed, imaginative scene descriptions for AI generation.\n"
        "Below is the user's current tags (prompt elements). "
        "Generate a new descriptive text (3-5 sentences) that incorporates these tags.\n\n"
        f"User Tags: {tags_text}\n"
        "Please generate a vivid, imaginative scene description."
    )

    try:
        if use_deepseek:
            # 调用 DeepSeek
            client = OpenAI(api_key=deepseek_key, base_url="https://api.deepseek.com")
            model = "deepseek-chat"
        else:
            # 调用 OpenAI
            client = OpenAI(api_key=api_key)
            model = "gpt-4"  # 或其他可用模型，比如 "gpt-3.5-turbo"
        response = client.chat.completions.create(
            messages=[
                {
                    "role": "system",
                    "content": "You are a creative assistant that generates imaginative scene descriptions..."
                },
                {
                    "role": "user",
                    "content": user_prompt,
                },
            ],
            model=model,
        )
        new_description = response.choices[0].message.content.strip()
    except Exception as e:
        new_description = f"(GPT generation failed: {e})"

    new_combined_output = f"{tags_text}\n\n{new_description}"
    return new_description, new_combined_output


# ========== 翻译功能：将 combined_output 翻译成用户选定语言 ==========
def translate_combined_output(combined_text, target_language, api_key, use_deepseek, deepseek_key):
    """
    使用 GPT 或 DeepSeek API，将 combined_text 翻译成 target_language。
    """
    if not api_key and not deepseek_key:
        return "(No API Key provided)"

    # 简单用 GPT 做翻译，也可改成其他翻译 API
    translation_prompt = (
        f"You are a professional translator. Please translate the following text into {target_language}.\n\n"
        f"{combined_text}"
    )

    try:
        if use_deepseek:
            # 调用 DeepSeek
            client = OpenAI(api_key=deepseek_key, base_url="https://api.deepseek.com")
            model = "deepseek-chat"
        else:
            # 调用 OpenAI
            client = OpenAI(api_key=api_key)
            model = "gpt-3.5-turbo"  # 或者别的模型
        response = client.chat.completions.create(
            messages=[
                {"role": "system", "content": "You are a professional translator."},
                {"role": "user", "content": translation_prompt},
            ],
            model=model,
        )
        translated_text = response.choices[0].message.content.strip()
    except Exception as e:
        translated_text = f"(Translation failed: {e})"

    return translated_text


# ========== 收藏功能：最多存 3 条 ==========
def add_to_favorites(combined_output, current_favorites):
    """
    将当前生成的 combined_output 添加到收藏列表中（最多存 3 条）。
    """
    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("## 【RPGAT】Random Prompt Generator with Adjustable Tags (V2)")

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

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

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

                use_deepseek = gr.Checkbox(label="Use DeepSeek API")
                dtr_enabled = gr.Checkbox(label="Enable DTR (当前无效-now plz ignore this)")

                with gr.Group():
                    gr.Markdown("**upload your wildcards(optional)--上传文件 (可选)：**")
                    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="排除excluded Tags (逗号分隔)",
                    placeholder="As如：angry, sword"
                )
                user_custom_tags = gr.Textbox(
                    label="自定义附加custom addition Tags (逗号分隔)",
                    placeholder="As如：glowing eyes, giant wings"
                )

                with gr.Group():
                    gr.Markdown("**Set the Number of the random counts--随机数量设置：**")
                    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="Waiting for generate-等待生成...",
                    lines=4,
                    interactive=True
                )
                description_output = gr.Textbox(
                    label="Generated Description",
                    placeholder="Waiting for generate-等待生成...",
                    lines=4,
                    interactive=True
                )
                combined_output = gr.Textbox(
                    label="Combined Output: Tags + Description",
                    placeholder="Waiting for generate-等待生成...",
                    lines=6
                )

                # 新增一个按钮，只更新 description 和 combined
                update_desc_button = gr.Button("Update Description Only")

                # 翻译相关
                with gr.Row():
                    target_language = gr.Dropdown(
                        choices=["English", "Chinese", "Arabic (language)", "Japanese", "Persian (language)", "Italian (language)", "Dutch (language)","Russian (language)","German (language)"],
                        value="English",
                        label="Target Language-目标语言"
                    )
                    translate_button = gr.Button("Translate to selected language")
                translated_output = gr.Textbox(
                    label="Translated Output",
                    placeholder="Waiting for the result of Translation--等待翻译...",
                    lines=6
                )

                # 收藏
                with gr.Row():
                    favorite_button = gr.Button("Favorites this result-收藏本次结果")
                    favorites_box = gr.Textbox(
                        label="Favorites Folder (MAX-Save 3 tags-- 最多 3 条)",
                        placeholder="暂无收藏",
                        lines=6
                    )

                # 点击“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, 
                        user_custom_tags, excluded_tags
                    ],
                    outputs=[tags_output, description_output, combined_output],
                )

                # 点击“Update Description Only”按钮
                update_desc_button.click(
                    update_description,
                    inputs=[
                        tags_output,     # 用户在文本框里编辑后的 Tags
                        api_key_input, 
                        use_deepseek,
                        deepseek_key_input,
                    ],
                    outputs=[description_output, combined_output],
                )

                # 点击“Translate to selected language”按钮
                translate_button.click(
                    fn=translate_combined_output,
                    inputs=[
                        combined_output,   # 要翻译的源文本
                        target_language,
                        api_key_input,
                        use_deepseek,
                        deepseek_key_input
                    ],
                    outputs=[translated_output],
                )

                # 收藏按钮点击事件
                favorite_button.click(
                    fn=add_to_favorites,
                    inputs=[combined_output, favorites_state],
                    outputs=[favorites_box, favorites_state],
                )

    return demo


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