app.py

import os
import json
import gradio as gr
from openai import OpenAI

##############################################################################
# 1. 读取外部文件: furry_species.json & gender_rules.json
##############################################################################
try:
    with open("furry_species.json", "r", encoding="utf-8") as f:
        FURRY_DATA = json.load(f)
except:
    FURRY_DATA = {}

try:
    with open("gender_rules.json", "r", encoding="utf-8") as f:
        GENDER_RULES = json.load(f)
except:
    GENDER_RULES = {}

##############################################################################
# 2. 多级菜单构造函数
##############################################################################
def get_top_categories(furry_data):
    return sorted(list(furry_data.keys()))

def get_sub_categories(furry_data, top_category):
    if top_category in furry_data:
        return sorted(list(furry_data[top_category].keys()))
    return []

def get_species_list(furry_data, top_category, sub_category):
    if top_category in furry_data and sub_category in furry_data[top_category]:
        return sorted(furry_data[top_category][sub_category])
    return []

##############################################################################
# 3. 核心调用：GPT 或 DeepSeek
##############################################################################
def generate_transformed_output(
    prompt,        # 原始 Prompt，如 "1girl, butterfly, solo..."
    gender_option, # 转换目标
    top_cat, sub_cat, species_item,
    api_mode, api_key
):
    """
    根据指定的性别/物种等规则，让 GPT/DeepSeek 输出仅两段内容：
     (转化后tags)

     转化后描述 (3~6句)
    不展示原始 prompt/base_prompt/gender: ... 等信息。
    """
    if not api_key:
        return "Error: No API Key provided."

    # 根据 gender_option 选对应的 rule
    if gender_option == "Trans_to_Male":
        rule_text = GENDER_RULES.get("male", "")
    elif gender_option == "Trans_to_Female":
        rule_text = GENDER_RULES.get("female", "")
    elif gender_option == "Trans_to_Mannequin":
        rule_text = GENDER_RULES.get("genderless", "")
    elif gender_option == "Trans_to_Intersex":
        rule_text = GENDER_RULES.get("intersex", "")
    else:
        # Furry
        # 你可以综合 male/female/intersex/genderless，也可以有专门 furry 的说明
        rule_text = (
            GENDER_RULES.get("male", "") + "\n\n"
            + GENDER_RULES.get("female", "") + "\n\n"
            + GENDER_RULES.get("intersex", "") + "\n\n"
            + GENDER_RULES.get("genderless", "")
        )
        # 如果想在规则中附加选定的物种信息：
        if top_cat and sub_cat and species_item:
            rule_text += f"\nFurry species: {top_cat} > {sub_cat} > {species_item}\n"

    # 选定 GPT or DeepSeek
    if api_mode == "GPT":
        base_url = None
        model_name = "gpt-3.5-turbo"  # 可改成 "gpt-4"
    else:
        base_url = "https://api.deepseek.com"
        model_name = "deepseek-chat"

    client = OpenAI(api_key=api_key)
    if base_url:
        client.base_url = base_url

    # 构造 System Prompt：要求只输出两段；第一行(转化后tags)，空行，随后3~6句描述
    # 让它把 prompt 中的 tags 进行「合并、替换、去重、增加」等处理
    system_prompt = f"""
You are a creative assistant that modifies the user's base prompt tags 
to reflect the correct gender/furry transformations, following these rules:

{rule_text}

Steps:
1) The original prompt tags are: {prompt}
2) Convert them into NEW combined tags that reflect the requested transformation. 
   (Remove or replace conflicting tags, unify synonyms, add any essential tags 
    for {gender_option} or for the selected furry species.)
3) Output EXACTLY two parts:
   - First line: the final, consolidated tags in parentheses (e.g. (male, solo, ...)).
   - Then a blank line.
   - Then a short imaginative scene description (3 to 6 sentences). 
4) Do NOT include 'gender:' or 'base_prompt:' or any headings or extra lines. 
5) Output everything in English.
6) Do not reference these steps in the final answer.
""".strip()

    try:
        resp = client.chat.completions.create(
            model=model_name,
            messages=[
                {"role": "system", "content": system_prompt},
                {"role": "user", "content": "Generate the final tags and description now."},
            ],
        )
        # 结果中仅包含 (tags)\n\n(description)
        result = resp.choices[0].message.content.strip()
        return result

    except Exception as e:
        return f"{api_mode} generation failed. Error: {e}"

def translate_text(content, lang, api_mode, api_key):
    """
    对上一步的 (tags)\n\n(description) 做翻译，保持格式
    """
    if not api_key:
        return "Error: No API Key provided."
    if not content.strip():
        return ""

    if api_mode == "GPT":
        base_url = None
        model_name = "gpt-3.5-turbo"
    else:
        base_url = "https://api.deepseek.com"
        model_name = "deepseek-chat"

    client = OpenAI(api_key=api_key)
    if base_url:
        client.base_url = base_url

    translate_system_prompt = f"""
You are a translator. Translate the following text to {lang}, 
preserving the parentheses line and blank lines if present. 
Do not add extra headings.
""".strip()

    try:
        resp = client.chat.completions.create(
            model=model_name,
            messages=[
                {"role": "system", "content": translate_system_prompt},
                {"role": "user", "content": content},
            ],
        )
        return resp.choices[0].message.content.strip()
    except Exception as e:
        return f"{api_mode} translation failed. Error: {e}"

##############################################################################
# 4. Gradio 前端
##############################################################################
def build_interface():
    with gr.Blocks() as demo:
        gr.Markdown("## Prompt Transformer - 不显示原Prompt与gender/base_prompt，只输出新tags和描述 (GPT/DeepSeek)")

        with gr.Row():
            with gr.Column():
                # 选择 GPT/DeepSeek
                api_mode = gr.Radio(
                    label="Choose API (GPT or DeepSeek)",
                    choices=["GPT", "DeepSeek"],
                    value="GPT"
                )
                # 输入 API Key
                api_key = gr.Textbox(
                    label="API Key",
                    type="password",
                    placeholder="Enter your GPT or DeepSeek Key here"
                )

                # 性别/Furry
                gender_option = gr.Radio(
                    label="Gender / Furry Conversion",
                    choices=[
                        "Trans_to_Male",
                        "Trans_to_Female",
                        "Trans_to_Mannequin",
                        "Trans_to_Intersex",
                        "Trans_to_Furry"
                    ],
                    value="Trans_to_Male"
                )

                # 若选 Furry -> 多级菜单
                top_cat_dd = gr.Dropdown(
                    label="Furry: Top Category",
                    choices=get_top_categories(FURRY_DATA),
                    value=None,
                    visible=False
                )
                sub_cat_dd = gr.Dropdown(
                    label="Furry: Sub Category",
                    choices=[],
                    value=None,
                    visible=False
                )
                species_dd = gr.Dropdown(
                    label="Furry: Species",
                    choices=[],
                    value=None,
                    visible=False
                )

                # 当性别选项切到 Furry 时，显示下拉，否则隐藏
                def show_furry_options(chosen):
                    if chosen == "Trans_to_Furry":
                        return (gr.update(visible=True),
                                gr.update(visible=True),
                                gr.update(visible=True))
                    else:
                        return (gr.update(visible=False),
                                gr.update(visible=False),
                                gr.update(visible=False))
                gender_option.change(
                    fn=show_furry_options,
                    inputs=[gender_option],
                    outputs=[top_cat_dd, sub_cat_dd, species_dd]
                )

                # 主分类 -> 子分类
                def on_top_cat_select(selected):
                    subs = get_sub_categories(FURRY_DATA, selected)
                    return gr.update(choices=subs, value=None)
                top_cat_dd.change(
                    fn=on_top_cat_select,
                    inputs=[top_cat_dd],
                    outputs=[sub_cat_dd]
                )

                # 子分类 -> 物种
                def on_sub_cat_select(top_c, sub_c):
                    sp = get_species_list(FURRY_DATA, top_c, sub_c)
                    return gr.update(choices=sp, value=None)
                sub_cat_dd.change(
                    fn=on_sub_cat_select,
                    inputs=[top_cat_dd, sub_cat_dd],
                    outputs=[species_dd]
                )

            # 输入 prompt & 输出
            with gr.Column():
                user_prompt = gr.Textbox(
                    label="原始 Prompt (Base Tags)",
                    lines=5,
                    placeholder="e.g. 1girl, butterfly, solo, flower..."
                )
                final_output = gr.Textbox(
                    label="(转化后Tags)\n\n(转化后描述)",
                    lines=10
                )

        # 翻译
        with gr.Row():
            translate_lang = gr.Dropdown(
                label="翻译语言",
                choices=["English", "Chinese", "Japanese", "French", "German", "Spanish"],
                value="English"
            )
            translated_result = gr.Textbox(
                label="翻译后的结果",
                lines=10
            )

        ######################################################################
        # 事件
        ######################################################################
        def on_generate(prompt, gender, tc, sc, sp, mode, key, lang):
            # 1) 生成新的 (tags) + 描述
            merged = generate_transformed_output(prompt, gender, tc, sc, sp, mode, key)
            # 2) 翻译
            trans = translate_text(merged, lang, mode, key)
            return merged, trans

        # 回车提交
        user_prompt.submit(
            fn=on_generate,
            inputs=[user_prompt, gender_option, top_cat_dd, sub_cat_dd, species_dd, api_mode, api_key, translate_lang],
            outputs=[final_output, translated_result]
        )
        # 点击按钮
        btn = gr.Button("生成 / Generate")
        btn.click(
            fn=on_generate,
            inputs=[user_prompt, gender_option, top_cat_dd, sub_cat_dd, species_dd, api_mode, api_key, translate_lang],
            outputs=[final_output, translated_result]
        )

    return demo


if __name__ == "__main__":
    demo = build_interface()
    demo.launch()