app.py

##############################################
# app.py
##############################################
import os
import json
import gradio as gr
import requests

# 新版 openai>=1.0.0
from openai import OpenAI

##############################################################################
# 1. 从外部文件中加载 Furry 物种数据 + Gender 细节文本
##############################################################################

# 读取 Furry 物种分类的 JSON 文件
# 假设文件名为 furry_species.json，结构示例：
# {
#   "CANIDS": ["Dogs", "Wolves", "Foxes"],
#   "FELINES": ["Lions", "Tigers", "Cheetahs"],
#   ...
# }
try:
    with open("furry_species.json", "r", encoding="utf-8") as fs_file:
        furry_map = json.load(fs_file)
except:
    furry_map = {}  # 若文件不存在或出错，可给个空dict备用

def flatten_furry_species(map_data):
    """
    将 {分类: [物种列表]} 展开为 ["分类 - 物种"] 形式的扁平列表
    """
    result = []
    for category, species_list in map_data.items():
        for sp in species_list:
            result.append(f"{category} - {sp}")
    return sorted(result)

ALL_FURRY_SPECIES = flatten_furry_species(furry_map)

# 读取 Gender 细节长文本 (记忆库)
try:
    with open("gender_details.txt", "r", encoding="utf-8") as gd_file:
        GENDER_DETAILS_TEXT = gd_file.read()
except:
    GENDER_DETAILS_TEXT = (
        "Gender conversion rules are missing. Please provide gender_details.txt."
    )

##############################################################################
# 2. 核心函数：根据选择调用 GPT 或 DeepSeek
##############################################################################

def generate_description_with_tags(prompt, gender_option, furry_species, api_mode, api_key):
    """
    1) 构造 tags（包含 gender/furry/base_prompt）。
    2) 调用 GPT 或 DeepSeek，带上记忆库信息 (GENDER_DETAILS_TEXT) 作为系统提示。
    3) 返回 (tags + 自然语言描述) 的输出形式。
    """

    if not api_key:
        return "Error: No API Key provided."

    # 组装 tags
    tags = {}
    if gender_option == "Trans_to_Male":
        tags["gender"] = "male"
    elif gender_option == "Trans_to_Female":
        tags["gender"] = "female"
    elif gender_option == "Trans_to_Mannequin":
        tags["gender"] = "genderless"
    elif gender_option == "Trans_to_Intersex":
        tags["gender"] = "intersex"
    elif gender_option == "Trans_to_Furry":
        tags["gender"] = "furry"
        tags["furry_species"] = furry_species or "unknown"

    tags["base_prompt"] = prompt

    # 确定 base_url + model
    if api_mode == "GPT":
        base_url = None
        model_name = "gpt-3.5-turbo"  # or "gpt-4", "gpt-4o", etc.
    else:
        # DeepSeek
        base_url = "https://api.deepseek.com"
        model_name = "deepseek-chat"

    # 创建 OpenAI Client
    client = OpenAI(api_key=api_key)
    if base_url:
        client.base_url = base_url

    # 将tags拼成便于阅读的字符串
    tags_str = "\n".join([f"{k}: {v}" for k, v in tags.items() if v])

    # 准备对话消息：把 gender_details 放在系统提示中，以便 AI 参考
    system_prompt = (
        "You are a creative assistant that generates detailed and imaginative scene descriptions "
        "for AI generation prompts. Focus on the details provided, incorporate them into a cohesive narrative, "
        "and obey the following gender transformation & furry rules:\n\n"
        f"{GENDER_DETAILS_TEXT}\n\n"
        "After forming your final description, return it in no more than five sentences. Thank you!"
    )

    try:
        response = client.chat.completions.create(
            model=model_name,
            messages=[
                {"role": "system", "content": system_prompt},
                {
                    "role": "user",
                    "content": (
                        f"Here are the tags:\n{tags_str}\n"
                        f"Please generate a vivid, imaginative scene description."
                    ),
                },
            ],
        )
        description = response.choices[0].message.content.strip()

        # 输出形式：(tags + 自然语言描述)
        output_text = f"=== Tags ===\n{tags_str}\n\n=== Description ===\n{description}"
        return output_text

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

def translate_text(text, translate_language, api_mode, api_key):
    """
    仅做一次简单翻译。若不需要翻译可以省去。也可用系统提示让 AI 翻译。
    """
    if not api_key:
        return "Error: No API Key provided."
    if not text.strip():
        return ""

    # GPT vs DeepSeek
    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

    system_prompt = f"You are a professional translator. Translate the following text to {translate_language}:"
    try:
        resp = client.chat.completions.create(
            model=model_name,
            messages=[
                {"role": "system", "content": system_prompt},
                {"role": "user",   "content": text},
            ]
        )
        return resp.choices[0].message.content.strip()
    except Exception as e:
        return f"{api_mode} translation failed. Error: {e}"

##############################################################################
# 3. Gradio 界面
##############################################################################
def build_interface():
    with gr.Blocks() as demo:
        gr.Markdown("## Prompt Transformer (GPT / DeepSeek) - (tags + 自然语言描述)")

        with gr.Row():
            with gr.Column():
                # 选择 API
                api_mode = gr.Radio(
                    label="选择API (GPT or DeepSeek)",
                    choices=["GPT", "DeepSeek"],
                    value="GPT",
                )
                api_key = gr.Textbox(
                    label="API密钥 (API Key)",
                    type="password",
                    placeholder="在此输入 GPT / DeepSeek 的API Key"
                )
                # 性别 / Furry
                gender_option = gr.Radio(
                    label="转换目标",
                    choices=[
                        "Trans_to_Male",
                        "Trans_to_Female",
                        "Trans_to_Mannequin",
                        "Trans_to_Intersex",
                        "Trans_to_Furry",
                    ],
                    value="Trans_to_Male",
                )
                furry_species = gr.Dropdown(
                    label="Furry 物种选择",
                    choices=ALL_FURRY_SPECIES,
                    value=None,
                    visible=False
                )

                def show_furry_species(choice):
                    return gr.update(visible=(choice == "Trans_to_Furry"))

                gender_option.change(show_furry_species, inputs=[gender_option], outputs=[furry_species])

            with gr.Column():
                user_prompt = gr.Textbox(
                    label="提示词 (Prompt)",
                    lines=4,
                    placeholder="示例：一位穿着蓝色长裙的少女，坐在海边..."
                )
                output_tags_and_desc = gr.Textbox(
                    label="(tags + 自然语言描述)",
                    lines=10
                )

        with gr.Row():
            translate_language = gr.Dropdown(
                label="翻译语言 (Translation Language)",
                choices=["English", "Chinese", "Japanese", "French", "German", "Spanish"],
                value="English"
            )
            translate_output = gr.Textbox(
                label="翻译结果",
                lines=10
            )

        # 点击生成 -> 同时调用生成描述 & 翻译
        def on_generate(prompt, gender, f_species, mode, key, lang):
            # 1. 生成 (tags + 自然语言描述)
            result_text = generate_description_with_tags(prompt, gender, f_species, mode, key)
            # 2. 翻译
            translated = translate_text(result_text, lang, mode, key)
            return result_text, translated

        user_prompt.submit(
            fn=on_generate,
            inputs=[user_prompt, gender_option, furry_species, api_mode, api_key, translate_language],
            outputs=[output_tags_and_desc, translate_output]
        )
        generate_btn = gr.Button("生成 / Generate")
        generate_btn.click(
            fn=on_generate,
            inputs=[user_prompt, gender_option, furry_species, api_mode, api_key, translate_language],
            outputs=[output_tags_and_desc, translate_output]
        )

    return demo

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