import gradio as gr
from transformers import pipeline
from diffusers import DiffusionPipeline

# 故事文字生成模型
text_pipeline = pipeline("text-generation", model="isarth/distill_gpt2_story_generator")

# 圖像生成模型
image_pipeline = DiffusionPipeline.from_pretrained("prompthero/openjourney")

# 文字生成函式
def generate_story(input_text):
    result = text_pipeline(input_text, max_length=200, do_sample=True)
    story = result[0]["generated_text"]
    return story

# 圖像生成函式
def generate_image_from_story(story_text):
    image = image_pipeline(story_text).images[0]
    return image

# Gradio 介面
with gr.Blocks() as demo:
    gr.Markdown("## ✨ AI 故事與圖片生成器")
    
    with gr.Row():
        input_text = gr.Textbox(label="輸入你的故事開頭", placeholder="從這裡開始你的冒險...")
        generate_btn = gr.Button("生成故事")

    story_output = gr.Textbox(label="生成的故事")
    
    with gr.Row():
        image_btn = gr.Button("根據故事生成圖片")
        image_output = gr.Image(label="生成的圖片")

    generate_btn.click(fn=generate_story, inputs=input_text, outputs=story_output)
    image_btn.click(fn=generate_image_from_story, inputs=story_output, outputs=image_output)

demo.launch()