import gradio as gr
import pandas as pd
import torch
from diffusers import StableDiffusionXLPipeline
from PIL import Image
import os
import zipfile

# 加载模型（GPU环境）
pipe = StableDiffusionXLPipeline.from_pretrained(
    "stabilityai/stable-diffusion-xl-base-1.0",
    torch_dtype=torch.float16,
    variant="fp16"
).to("cuda")

# 文本模板生成函数
def generate_text(prompt):
    return f"这是一幅描绘：{prompt} 的画面，小朋友可以根据图像发挥想象力哦！"

# 模式一：处理CSV并生成图文ZIP
def process_csv(file):
    df = pd.read_csv(file.name)
    output_dir = "output"
    os.makedirs(output_dir, exist_ok=True)

    for idx, row in df.iterrows():
        prompt = row["prompt"]
        image = pipe(prompt=prompt).images[0]
        image_path = os.path.join(output_dir, f"image_{idx+1}.png")
        image.save(image_path)

        text = generate_text(prompt)
        text_path = os.path.join(output_dir, f"text_{idx+1}.txt")
        with open(text_path, "w") as f:
            f.write(text)

    # 打包ZIP
    zip_path = "output.zip"
    with zipfile.ZipFile(zip_path, "w") as zipf:
        for file_name in os.listdir(output_dir):
            file_path = os.path.join(output_dir, file_name)
            zipf.write(file_path, arcname=file_name)
    return zip_path

# 模式二：文本转图片
def text_to_image(prompt):
    image = pipe(prompt=prompt).images[0]
    return image

# Gradio界面
with gr.Blocks() as demo:
    gr.Markdown("## 🤖 AI 图文生成器 - 批量 & 单图模式")

    with gr.Tab("📂 批量生成（上传CSV）"):
        csv_input = gr.File(label="上传CSV文件", file_types=[".csv"])
        csv_output = gr.File(label="生成图文ZIP包")
        csv_btn = gr.Button("开始生成")
        csv_btn.click(fn=process_csv, inputs=csv_input, outputs=csv_output)

    with gr.Tab("🖼️ 单图生成（文本转图片）"):
        prompt_input = gr.Textbox(label="输入提示词", placeholder="比如：一只飞翔在太空的小猫咪")
        image_output = gr.Image(label="生成的图像", type="pil")
        single_btn = gr.Button("立即生成")
        single_btn.click(fn=text_to_image, inputs=prompt_input, outputs=image_output)

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