Update app.py

app.py

@@ -2,7 +2,7 @@ import os
 import torch
 import gradio as gr
 from transformers import pipeline
-from diffusers import StableDiffusionPipeline
+from diffusers import StableDiffusionPipeline, DPMSolverMultistepScheduler
 
 # 如果需要使用 Hugging Face 访问令牌，取消下面一行的注释并设置环境变量 HUGGINGFACE_TOKEN
 # from huggingface_hub import login
@@ -15,35 +15,28 @@ llm = pipeline(
     device=0 if torch.cuda.is_available() else -1
 )
 
-# Step 2: 加载 Stable Diffusion 模型并优化以加速推理
+# Step 2: 加载并量化 Stable Diffusion 模型以加速推理
+# 使用 8-bit 量化和自动设备映射
+device = "cuda" if torch.cuda.is_available() else "cpu"
+load_kwargs = {
+    "torch_dtype": torch.float16 if device == "cuda" else torch.float32,
+    "device_map": "auto",
+    "load_in_8bit": True    # 需要安装 bitsandbytes
+}
+
 # SD v1.5
 sd_v15 = StableDiffusionPipeline.from_pretrained(
     "runwayml/stable-diffusion-v1-5",
-    torch_dtype=torch.float16
+    **load_kwargs
 )
-# 开启注意力切片，减少显存峰值
-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_v15.scheduler = DPMSolverMultistepScheduler.from_config(sd_v15.scheduler.config)
 
 # SD XL
 sd_xl = StableDiffusionPipeline.from_pretrained(
     "stabilityai/stable-diffusion-xl-base-1.0",
-    torch_dtype=torch.float16
+    **load_kwargs
 )
-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")
+sd_xl.scheduler = DPMSolverMultistepScheduler.from_config(sd_xl.scheduler.config)
 
 # 可选：语音输入模块，使用 Whisper
 asr = pipeline(
@@ -57,7 +50,7 @@ def transcribe(audio_path):
     return text
 
 
-def generate(description, model_choice, guidance_scale, negative_prompt, style):
+def generate(description, model_choice, guidance_scale, negative_prompt, style, steps):
     # 构造给 LLM 的指令
     instruction = (
         f"请将以下简短描述扩展为 Stable Diffusion 友好的提示词，包含细节和风格：\n"
@@ -71,7 +64,8 @@ def generate(description, model_choice, guidance_scale, negative_prompt, style):
     image = pipeline_model(
         prompt,
         guidance_scale=guidance_scale,
-        negative_prompt=negative_prompt
+        negative_prompt=negative_prompt,
+        num_inference_steps=steps
     ).images[0]
     return prompt, image
 
@@ -94,6 +88,10 @@ with gr.Blocks(title="Prompt-to-Image Generator") as demo:
                 minimum=0, maximum=20, step=0.5, value=7.5, 
                 label="Guidance Scale"
             )
+            steps_slider = gr.Slider(
+                minimum=1, maximum=50, step=1, value=20,
+                label="推理步数 (步数减少可加速)"
+            )
             neg_text = gr.Textbox(
                 label="反向提示词", 
                 placeholder="排除内容（如：低分辨率、水印）"
@@ -117,7 +115,7 @@ with gr.Blocks(title="Prompt-to-Image Generator") as demo:
     # 点击按钮生成提示词并绘图
     generate_btn.click(
         fn=generate, 
-        inputs=[desc_input, model_radio, guidance_slider, neg_text, style_dropdown], 
+        inputs=[desc_input, model_radio, guidance_slider, neg_text, style_dropdown, steps_slider], 
         outputs=[prompt_output, image_output]
     )