Update app.py

app.py

@@ -1,6 +1,6 @@
 import gradio as gr
 from transformers import pipeline, WhisperProcessor, WhisperForConditionalGeneration
-from diffusers import StableDiffusionPipeline
+from diffusers import DiffusionPipeline
 import torch
 
 # 使用BART模型生成文本描述
@@ -12,13 +12,11 @@ def generate_prompt(description: str) -> str:
     return prompt
 
 # 加载 ByteDance/SDXL-Lightning 模型
-sdxl_pipeline = StableDiffusionPipeline.from_pretrained("ByteDance/SDXL-Lightning")
-sdxl_pipeline.to("cpu")  # 使用 CPU
-
+pipe = DiffusionPipeline.from_pretrained("ByteDance/SDXL-Lightning")
 def generate_image(prompt: str):
     # 根据提示生成图像
-    image = sdxl_pipeline(prompt).images[0]
-    return image
+    prompt = "Astronaut in a jungle, cold color palette, muted colors, detailed, 8k"
+    image = pipe(prompt).images[0]
 
 # 使用Whisper模型进行语音转文本
 processor = WhisperProcessor.from_pretrained("openai/whisper-large")