app.py

import os
import pandas as pd
import gradio as gr
from openai import OpenAI

# 初始化 OpenAI 客户端
client = OpenAI(
    api_key=os.getenv("DEEPSEEK_API_KEY"),
    base_url="https://api.deepseek.com"
)

def read_csv_prefix(file, n_lines=128):
    """读取 CSV 文件的前 n 行并提取统一前缀"""
    df = pd.read_csv(file.name)
    if 'prompt' in df.columns:
        return df['prompt'].head(n_lines).tolist()
    else:
        raise ValueError("CSV 文件中没有 'prompt' 列")

def generate_prompts(user_input, csv_prompts, prefix, num_examples=3, language='English'):
    """基于用户输入和 CSV 提示生成多个新提示词"""
    # 构建系统提示
    '''
    system_prompt = (
        f"你是一位自然诗人，请将以下内容转化为{'英文' if language == 'English' else '中文'}的风景描写。"
        "请确保风格与提供的 CSV 中的描述一致，并以统一前缀 '{prefix}' 开头。"
        "请生成 {num_examples} 条不同风格的结果，每条结果不要加编号，且保持简洁自然。注意：输出语言必须与要求一致。"
    )
    '''

    system_prompt = (
    "你是一位专业的动漫视觉艺术家助手，负责根据用户输入生成符合 Wan 2.1 格式的文本提示词。\n"
    "请严格按照以下结构组织输出内容：\n\n"

    "1. 风格指定（Style Specification）：明确画面风格，例如 'anime style' 或特定艺术风格如 'in the style of Kinich'。\n"
    "2. 场景描述（Scene Description）：详细描述环境与背景，包括光影、天气、氛围等元素。\n"
    "3. 主体特征（Subject Features）：描述主要人物或对象的外貌、服装、动作等细节。\n"
    "4. 动态元素（Motion Elements）：如果是视频，请加入动作或场景变化，如角色移动、镜头推进等。\n"
    "5. 情感与氛围（Mood & Atmosphere）：强调整体情绪，如宁静、浪漫、神秘等。\n"
    "6. 特殊效果（Special Effects）：添加增强表现力的细节，如光效、粒子、天气动态等。\n\n"

    "以下是几个高质量提示词示例供你参考：\n\n"

    "【示例 1】\n"
    "{prefix} Golden light filters through the canopy, illuminating soft moss and fallen leaves. Wildflowers bloom nearby, and glowing fireflies hover in the air. A gentle stream flows in the background, its murmur blending with birdsong. The scene radiates tranquility and natural charm.\n\n"

    "【示例 2】\n"
    "{prefix} A majestic castle under a starry sky with silvery moonlight casting soft shadows. Towering spires rise into the night, their peaks adorned with glowing orbs that mimic the stars above. Fireflies dance around the castle’s ivy-covered arches, adding a touch of magic to the scene.\n\n"

    "【示例 3】\n"
    "{prefix} The video shifts to a breathtaking coastal scene. The turquoise sea stretches endlessly, its waves gently lapping against the golden sandy shore. Rocky cliffs rise along the coastline, their jagged edges softened by patches of green vegetation. Seafoam glistens as it washes ashore, and the air is filled with the soothing sound of the tide.\n\n"

    "【示例 4】\n"
    "{prefix} As the scene progresses, a pink lunar disc gradually emerges from pitch-black night, with violet clouds rolling like waves, the unearthly moonlight casting shifting halos across the drifting cloud sea.\n\n"

    "【示例 5】\n"
    "{prefix} A young male character with teal and dark blue tousled hair adorned with geometric neon patterns charges forward — left hand crackling with a swirling energy orb, right arm outstretched, trailing luminous particles. Behind him, shattered urban ruins float amidst explosive shockwaves.\n\n"

    "请确保每条提示词以 '{prefix}' 开头，使用 {language} 输出，且每条提示词之间用空行分隔。\n"
    "不要编号、不加额外解释，保持自然流畅。"
    )

    # 拼接用户输入和 CSV 内容
    user_content = f"转换以下内容：\n{user_input}\n\n参考示例：\n" + "\n".join(csv_prompts[:num_examples])
    
    # 替换为：
    language_label = "英文" if language == "English" else "中文"
    
    response = client.chat.completions.create(
        model="deepseek-chat",
        messages=[
            {"role": "system", "content": system_prompt.format(prefix=prefix, num_examples=num_examples, language=language_label)},
            {"role": "user", "content": user_content}
        ],
        temperature=0.7,
        max_tokens=300 * num_examples
    )

    # 将模型输出按换行分割为多个提示词
    raw_output = response.choices[0].message.content.strip()
    prompts = [line.strip() for line in raw_output.split('\n') if line.strip()]

    # 限制数量
    prompts = prompts[:num_examples]

    # 处理后文本：每行以指定前缀开头（先判断是否已包含）
    processed_prompts = []
    for p in prompts:
        if prefix not in p:
            processed_prompts.append(f"{prefix} {p}")
        else:
            processed_prompts.append(p)

    # 输出格式
    output_text = "\n\n".join(prompts)
    processed_text = "\n".join(processed_prompts)

    return output_text, processed_text

# Gradio 界面定义
with gr.Blocks(title="提示词生成器", theme=gr.themes.Soft()) as app:
    gr.Markdown("## 🎨 基于 CSV 的提示词生成器")
    
    with gr.Row():
        with gr.Column():
            input_text = gr.Textbox(
                label="输入文本",
                placeholder="例如：森林里的阳光透过树叶洒落...",
                lines=5
            )
            csv_file = gr.File(label=".csv 文件", file_types=[".csv"])
            line_count = gr.Number(label="读取行数", value=128, minimum=1, maximum=10000)
            prefix_input = gr.Textbox(label="提示词前缀", value="In the style of anime landscape,", lines=1)
            num_examples = gr.Number(label="生成示例数量", value=3, minimum=1, maximum=10)
            language_choice = gr.Radio(choices=["English", "Chinese"], label="输出语言", value="English")
            btn = gr.Button("生成提示词", variant="primary")
        with gr.Column():
            output_text = gr.Textbox(
                label="生成的提示词",
                lines=8,
                interactive=False
            )
            processed_text = gr.Textbox(
                label="处理后文本（每行以指定前缀开头）",
                lines=8,
                interactive=False
            )
    
    # 示例数据
    examples = gr.Examples(
        examples=[
            ["清晨的阳光洒在湖面上"],
            ["保留角色性别和外貌，描述角色的动作场面"]
        ],
        inputs=[input_text],
        label="点击试试示例"
    )

    # 事件绑定
    btn.click(
        fn=lambda text, file, lines, prefix, num, lang: generate_prompts(text, read_csv_prefix(file, lines), prefix, int(num), lang),
        inputs=[input_text, csv_file, line_count, prefix_input, num_examples, language_choice],
        outputs=[output_text, processed_text]
    )

app.launch(share=True)