File size: 2,288 Bytes
644bf93 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 |
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()
|