Update app.py

app.py

@@ -1,13 +1,40 @@
 import gradio as gr
 from transformers import pipeline
+from diffusers import DiffusionPipeline
 
-pipeline = pipeline("text-generation", model="isarth/distill_gpt2_story_generator")
-# pipeline = pipeline("text-generation", model="tohur/natsumura-storytelling-rp-1.0-llama-3.1-8b")
+# 故事文字生成模型
+text_pipeline = pipeline("text-generation", model="isarth/distill_gpt2_story_generator")
 
-def text_generate(input_text):
-    result = pipeline(input_text)
-    result = result[0]["generated_text"]
-    return result
+# 圖像生成模型
+image_pipeline = DiffusionPipeline.from_pretrained("prompthero/openjourney")
+image_pipeline.to("cuda")  # 若部署環境支援 GPU，這樣可加速生成
+
+# 文字生成函式
+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 = gr.Interface(fn=text_generate, inputs="text", outputs="text")
 demo.launch()