import gradio as gr
from transformers import pipeline
from diffusers import StableDiffusionPipeline
import torch
from PIL import Image

# === 模型定义 ===
llm = pipeline("text2text-generation", model="google/flan-t5-large")
stt = pipeline("automatic-speech-recognition", model="openai/whisper-small")

# 加载 Stable Diffusion v1.5 模型（去掉 float16）
pipe = StableDiffusionPipeline.from_pretrained(
    "runwayml/stable-diffusion-v1-5",
    safety_checker=None,  # 可选：避免屏蔽图像
)
pipe = pipe.to("cpu")  # 或者 "cuda" 如果你在 GPU 上部署

# === Prompt-to-Prompt ===
def refine_prompt(user_input):
    prompt = f"请将这句话改写为适合图像生成的英文提示词：'{user_input}'"
    result = llm(prompt, max_new_tokens=50)[0]['generated_text']
    return result

# === Prompt-to-Image ===
def generate_image(prompt, model_choice, num_steps, guidance):
    image = pipe(prompt, num_inference_steps=num_steps, guidance_scale=guidance).images[0]
    return prompt, image

# === 总流程 ===
def process_all(user_input, model_choice, steps, guidance):
    refined = refine_prompt(user_input)
    prompt, image = generate_image(refined, model_choice, steps, guidance)
    return refined, image

# === 语音识别 ===
def transcribe_audio(audio):
    text = stt(audio)["text"]
    return text

# === Gradio UI ===
with gr.Blocks() as demo:
    gr.Markdown("## 🪄 Prompt-to-Image Generator")

    with gr.Row():
        with gr.Column():
            audio_input = gr.Audio(type="filepath", label="🎤 或点击录音")
            voice_btn = gr.Button("使用语音输入")
            textbox = gr.Textbox(label="描述一句画面", placeholder="比如：空中的魔法树屋")
            model_choice = gr.Radio(["SD v1.4", "SDXL"], value="SDXL", label="选择模型")
            steps_slider = gr.Slider(10, 50, value=30, step=5, label="推理步数")
            guidance_slider = gr.Slider(5, 15, value=7.5, step=0.5, label="引导系数")
            submit_btn = gr.Button("生成图像 🎨")

        with gr.Column():
            refined_out = gr.Textbox(label="生成的提示词", lines=2)
            image_out = gr.Image(label="生成图像", type="pil")

    voice_btn.click(fn=transcribe_audio, inputs=audio_input, outputs=textbox)
    submit_btn.click(fn=process_all, inputs=[textbox, model_choice, steps_slider, guidance_slider], outputs=[refined_out, image_out])

demo.launch()