Update app.py

app.py

@@ -15,17 +15,34 @@ llm = pipeline(
     device=0 if torch.cuda.is_available() else -1
 )
 
-# Step 2: 加载 Stable Diffusion 模型
-# 移除无效的 revision 参数，仅使用 torch_dtype 加速加载
+# Step 2: 加载 Stable Diffusion 模型并优化以加速推理
+# SD v1.5
 sd_v15 = StableDiffusionPipeline.from_pretrained(
     "runwayml/stable-diffusion-v1-5",
     torch_dtype=torch.float16
 )
+# 开启注意力切片，减少显存峰值
+sd_v15.enable_attention_slicing()
+# 如果安装了 xformers，启用更高效的注意力实现
+try:
+    sd_v15.enable_xformers_memory_efficient_attention()
+except Exception:
+    pass
+# 启用CPU内存卸载，减轻GPU显存压力
+sd_v15.enable_model_cpu_offload()
 sd_v15 = sd_v15.to("cuda" if torch.cuda.is_available() else "cpu")
 
+# SD XL
 sd_xl = StableDiffusionPipeline.from_pretrained(
-    "stabilityai/stable-diffusion-xl-base-1.0"
+    "stabilityai/stable-diffusion-xl-base-1.0",
+    torch_dtype=torch.float16
 )
+sd_xl.enable_attention_slicing()
+try:
+    sd_xl.enable_xformers_memory_efficient_attention()
+except Exception:
+    pass
+sd_xl.enable_model_cpu_offload()
 sd_xl = sd_xl.to("cuda" if torch.cuda.is_available() else "cpu")
 
 # 可选：语音输入模块，使用 Whisper
@@ -82,7 +99,6 @@ with gr.Blocks(title="Prompt-to-Image Generator") as demo:
                 placeholder="排除内容（如：低分辨率、水印）"
             )
             use_voice = gr.Checkbox(label="启用语音输入（加分项）")
-            # 移除 'source' 参数以兼容 Gradio 版本
             audio_input = gr.Audio(type="filepath", label="语音输入")
             generate_btn = gr.Button("生成图像")
         with gr.Column():