app.py

# app.py

import gradio as gr
import os
import torch
from diffusers import StableDiffusionPipeline, DPMSolverMultistepScheduler
from transformers import pipeline
from openai import OpenAI
import time
from PIL import Image

# --- 配置 ---

# 尝试从 Hugging Face Secrets 或环境变量获取 API 密钥
# 重要：请将您的 API 密钥在 Hugging Face Space Secrets 中命名为 'OPENAI_API_KEY'！
OPENAI_API_KEY = os.getenv("OPENAI_API_KEY")
if not OPENAI_API_KEY:
    print("警告：未在环境变量/Secrets中找到OpenAI API密钥。")
    # 如果需要强制要求密钥存在，可以提供占位符或引发错误
    # raise ValueError("缺少 OpenAI API 密钥！请在 Hugging Face Space Secrets 中设置。")

# 初始化 OpenAI 客户端 (仅当找到密钥时)
openai_client = None
if OPENAI_API_KEY:
    try:
        openai_client = OpenAI(api_key=OPENAI_API_KEY)
    except Exception as e:
        print(f"初始化OpenAI客户端时出错: {e}")
        openai_client = None # 确保初始化失败时客户端为 None
else:
    print("因缺少API密钥，OpenAI客户端未初始化。")


# --- 模型加载 ---
# 在主函数外部加载模型，以防止每次请求都重新加载。
# Space 启动时可能需要一些时间。

# 1. Stable Diffusion 模型 (v1.5 以降低 CPU 资源使用)
# 使用 float32 以获得更好的 CPU 兼容性。在 CPU 上生成会很慢。
print("正在加载 Stable Diffusion 模型...")
start_time = time.time()
try:
    pipe = StableDiffusionPipeline.from_pretrained(
        "runwayml/stable-diffusion-v1-5",
        torch_dtype=torch.float32 # CPU 使用 float32
    )
    # 使用更快的调度器可以在 CPU 上有所帮助
    pipe.scheduler = DPMSolverMultistepScheduler.from_config(pipe.scheduler.config)
    # 如果在免费层遇到 RAM 问题，可能需要进一步优化
    # pipe.enable_attention_slicing() # 有助于略微减少内存使用
    print(f"Stable Diffusion 模型加载完成，耗时 {time.time() - start_time:.2f} 秒。")
except Exception as e:
    print(f"加载 Stable Diffusion 模型时出错: {e}")
    pipe = None # 如果加载失败，则将 pipe 设置为 None


# 2. 语音转文本模型 (Whisper - base 模型在 CPU 上是较好的平衡)
print("正在加载语音转文本模型...")
start_time = time.time()
try:
    # 使用 'base' 模型以在 CPU 上获得更快的推理速度
    # 如果需要，可以显式指定 device='cpu'，尽管它通常会正确默认
    asr_pipeline = pipeline("automatic-speech-recognition", model="openai/whisper-base", device=-1) # device=-1 强制使用 CPU
    print(f"语音转文本模型加载完成，耗时 {time.time() - start_time:.2f} 秒。")
except Exception as e:
    print(f"加载语音转文本模型时出错: {e}")
    asr_pipeline = None # 如果加载失败，则将 pipeline 设置为 None


# --- 核心函数 ---

# 步骤 1: Prompt-to-Prompt 模块
def generate_detailed_prompt(short_prompt, style_preference="电影感", neg_prompt=""):
    """使用 LLM 生成详细提示词。"""
    if not openai_client:
        print("OpenAI 客户端不可用。")
        gr.Warning("OpenAI API Key 未配置。将使用基础提示词。")
        # 回退：如果没有 LLM，则仅使用简短提示词和风格
        return f"{short_prompt}，风格：{style_preference}", neg_prompt or "丑陋, 模糊, 低质量"

    # 为 LLM 构建更具指导性的提示 (保持英文以确保模型理解)
    system_message = "You are an expert prompt generator for Stable Diffusion."
    user_message = (
        f"Based on the user's simple idea: '{short_prompt}', "
        f"generate a detailed and structured prompt for Stable Diffusion v1.5. "
        f"Incorporate the desired style: '{style_preference}'. "
        f"Include details about the subject, setting, lighting, composition, and quality keywords (like 'photorealistic', 'highly detailed', '4k', 'masterpiece'). "
        f"Also, suggest a relevant negative prompt focusing on common image issues (like 'ugly, deformed, blurry, low quality, extra limbs, text, words'). "
        f"If the user provided a negative prompt ('{neg_prompt}'), incorporate its essence or add to it. "
        f"Format the output clearly, separating the main prompt and the negative prompt, perhaps using '### Prompt:' and '### Negative Prompt:' labels."
    )

    try:
        print(f"向OpenAI发送请求以生成提示词： {short_prompt}")
        response = openai_client.chat.completions.create(
            model="gpt-3.5-turbo", # 一个相对快速且有能力的模型
            messages=[
                {"role": "system", "content": system_message},
                {"role": "user", "content": user_message},
            ],
            max_tokens=150, # 限制 token 使用
            temperature=0.7,
        )
        generated_text = response.choices[0].message.content.strip()
        print(f"收到来自OpenAI的响应： {generated_text}")

        # 假设 LLM 遵循请求的格式进行基本解析
        prompt_marker = "### Prompt:"
        neg_prompt_marker = "### Negative Prompt:"

        final_prompt = generated_text # 默认为全部文本
        final_neg_prompt = neg_prompt # 默认为用户输入的反向提示词

        if prompt_marker in generated_text:
            start_idx = generated_text.find(prompt_marker) + len(prompt_marker)
            end_idx = generated_text.find(neg_prompt_marker)
            if end_idx != -1:
                 final_prompt = generated_text[start_idx:end_idx].strip()
            else:
                 final_prompt = generated_text[start_idx:].strip()

        if neg_prompt_marker in generated_text:
            start_idx = generated_text.find(neg_prompt_marker) + len(neg_prompt_marker)
            final_neg_prompt = generated_text[start_idx:].strip()
            # 如果用户提供了反向提示词，则合并
            if neg_prompt and neg_prompt not in final_neg_prompt:
                final_neg_prompt = f"{neg_prompt}, {final_neg_prompt}"


        # 清理可能未完美解析的残留部分
        final_prompt = final_prompt.replace("### Prompt:", "").strip()
        final_neg_prompt = final_neg_prompt.replace("### Negative Prompt:", "").strip()

        # 确保不返回空字符串
        if not final_prompt:
             final_prompt = f"{short_prompt}，风格：{style_preference}，高度详细"
        if not final_neg_prompt:
             final_neg_prompt = "丑陋, 模糊, 低质量" # 默认中文反向提示词

        return final_prompt, final_neg_prompt

    except Exception as e:
        print(f"调用OpenAI API时出错: {e}")
        gr.Warning(f"OpenAI API 错误: {e}。将使用基础提示词。")
        # 出错时回退
        return f"{short_prompt}，风格：{style_preference}，高度详细", neg_prompt or "丑陋, 模糊, 低质量"

# 步骤 2: Prompt-to-Image 模块
def generate_image(prompt, neg_prompt, guidance, steps):
    """使用 Stable Diffusion 生成图像。"""
    if not pipe:
        gr.Error("Stable Diffusion 模型加载失败。无法生成图像。")
        # 返回占位符图像或 None
        return Image.new('RGB', (512, 512), color = 'grey'), "错误：SD 模型未加载。"

    print(f"正在使用提示词生成图像： {prompt}")
    print(f"反向提示词: {neg_prompt}, 引导系数: {guidance}, 步数: {steps}")
    start_time = time.time()
    try:
        # 在 CPU 上生成图像 (会很慢!)
        with torch.no_grad(): # 确保不计算梯度
             image = pipe(
                 prompt,
                 negative_prompt=neg_prompt,
                 guidance_scale=float(guidance),
                 num_inference_steps=int(steps),
                 num_images_per_prompt=1,
                 # height=512, # v1.5 默认值
                 # width=512,  # v1.5 默认值
             ).images[0]
        print(f"图像生成完成，耗时 {time.time() - start_time:.2f} 秒。")
        return image, prompt # 返回图像和使用的提示词
    except Exception as e:
        print(f"图像生成过程中出错: {e}")
        gr.Error(f"图像生成失败: {e}")
        return Image.new('RGB', (512, 512), color = 'red'), prompt # 返回错误占位符


# 附加功能: 语音转文本函数
def transcribe_audio(audio_filepath):
    """使用 Whisper 模型转录音频。"""
    if not asr_pipeline:
        gr.Warning("语音转文本模型未加载。")
        return "错误：ASR模型未加载。"
    if audio_filepath is None:
        gr.Warning("未提供音频输入。")
        return "" # 如果没有音频，则返回空字符串

    print(f"正在转录音频文件： {audio_filepath}")
    start_time = time.time()
    try:
        transcript = asr_pipeline(audio_filepath)
        print(f"音频转录完成，耗时 {time.time() - start_time:.2f} 秒。")
        print(f"转录结果： {transcript['text']}")
        return transcript["text"]
    except Exception as e:
        print(f"音频转录过程中出错: {e}")
        gr.Error(f"音频转录失败: {e}")
        return f"音频转录过程中出错: {e}"

# --- 组合工作流函数 ---
def process_input(short_prompt_text, audio_input, style, neg_prompt, guidance, steps):
    """按钮点击触发的主函数。"""
    final_short_prompt = short_prompt_text
    status = "正在开始..."

    # 1. 确定输入：如果可用，优先使用音频
    if audio_input is not None:
        status = "正在转录音频..."
        print(f"检测到音频输入： {audio_input}")
        # 如果需要在 Gradio 中更新状态（需要生成器函数）
        # yield status, None, None # 更新状态，此时还没有提示词/图像
        transcription = transcribe_audio(audio_input)
        if "错误" not in transcription and transcription.strip(): # 检查是否包含错误并且不是空转录
            final_short_prompt = transcription
            status = f"使用转录文本： '{transcription[:50]}...'"
            print(status)
        elif "错误" in transcription:
             status = "音频转录失败。如文本框有内容将使用文本输入。"
             print(status)
             # 如果文本输入也为空，需要明确错误
             if not final_short_prompt:
                 gr.Error("音频转录失败且文本框为空。")
                 return "错误：无有效输入。", None, None # 状态, 提示词, 图像
        else: # 转录成功但为空（例如，静音）
            status = "音频转录结果为空。如文本框有内容将使用文本输入。"
            print(status)
            # 如果文本输入也为空
            if not final_short_prompt:
                 gr.Error("音频转录结果为空且文本框为空。")
                 return "错误：无有效输入。", None, None


    if not final_short_prompt:
        gr.Error("请输入简短描述或使用语音输入。")
        return "错误：输入为空。", None, None # 状态, 提示词, 图像

    # 2. 生成详细提示词
    status = "正在生成详细提示词..."
    print(status)
    # yield status, None, None
    # 使用中文风格偏好调用
    detailed_prompt, final_neg_prompt = generate_detailed_prompt(final_short_prompt, style, neg_prompt)
    status = "正在生成图像..."
    print(f"使用详细提示词： {detailed_prompt}")
    print(f"使用反向提示词： {final_neg_prompt}")
    # yield status, detailed_prompt, None # 显示详细提示词

    # 3. 生成图像
    image, used_prompt = generate_image(detailed_prompt, final_neg_prompt, guidance, steps)
    status = "图像生成完成！"
    print(status)
    return status, used_prompt, image # 最终状态, 提示词, 图像


# --- Gradio 界面 ---
with gr.Blocks(theme=gr.themes.Soft()) as demo:
    gr.Markdown("# Stable Diffusion 提示词优化与图像生成器")
    gr.Markdown(
        "输入简短描述（或使用语音输入！），选择风格，获取详细提示词和生成的图像。\n"
        "*注意：图像生成在 CPU 上运行，将会**很慢**（可能需要几分钟）。*\n"
        "*提示词优化功能需要在 Hugging Face Space Secrets 中设置 `OPENAI_API_KEY`。*"
        )

    # 可以使用 Accordion 进行更好的组织，或仅使用 Row/Column
    with gr.Row():
        with gr.Column(scale=1):
            # 输入控件
            inp_short_prompt = gr.Textbox(label="简短描述", placeholder="例如：天空中的魔法树屋")
            inp_audio = gr.Audio(sources=["microphone"], type="filepath", label="或者录制语音输入（可选）") # Type='filepath' 通常对 pipeline 更好
            inp_style = gr.Dropdown(
                label="图像风格",
                choices=["照片写实", "电影感", "奇幻艺术", "动漫", "水彩", "像素艺术", "赛博朋克"],
                value="电影感" # 默认值
            ) # 控件类型 1: Dropdown
            inp_neg_prompt = gr.Textbox(label="反向提示词（可选）", placeholder="例如：文字, 词语, 模糊, 变形") # 控件类型 2: Textbox
            inp_guidance = gr.Slider(minimum=1, maximum=15, step=0.5, value=7.5, label="引导系数 (CFG)") # 控件类型 3: Slider
            inp_steps = gr.Slider(minimum=10, maximum=50, step=1, value=20, label="推理步数（越低越快，细节可能减少）") # 控件类型 4: 另一个 Slider

            generate_button = gr.Button("生成图像", variant="primary") # 控件类型 5: Button

            # 状态显示
            out_status = gr.Textbox(label="状态", value="准备就绪", interactive=False)

        with gr.Column(scale=1):
            # 输出控件
            out_detailed_prompt = gr.Textbox(label="生成的详细提示词", interactive=False, lines=5) # 控件类型 6: Textbox (输出)
            out_image = gr.Image(label="生成的图像", type="pil") # 控件类型 7: Image 输出


    # 连接组件
    generate_button.click(
        fn=process_input,
        inputs=[
            inp_short_prompt,
            inp_audio,
            inp_style,
            inp_neg_prompt,
            inp_guidance,
            inp_steps
        ],
        outputs=[
            out_status,
            out_detailed_prompt,
            out_image
        ]
    )

# --- 启动应用 ---
if __name__ == "__main__":
    # 启动前检查模型加载状态
    if not pipe:
        print("\n严重错误：Stable Diffusion 模型加载失败。Gradio 应用可能无法正常工作。")
    if not asr_pipeline:
        print("\n警告：语音转文本模型加载失败。语音输入将无法工作。")
    if not openai_client:
        print("\n警告：OpenAI客户端未初始化（检查API密钥）。提示词优化将是基础模式。")

    demo.queue() # 为处理多个用户/较长任务启用队列
    demo.launch(debug=False) # 如果需要本地更详细的日志，设置 debug=True