Update app.py

app.py

@@ -16,13 +16,22 @@ llm = pipeline(
 )
 
 # 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
-}
+# 根据硬件环境选择加载方式：
+# - GPU 环境：8-bit 量化 + 自动设备映射
+# - CPU 环境：浮点32 + 平衡设备映射
+if torch.cuda.is_available():
+    device = "cuda"
+    load_kwargs = {
+        "torch_dtype": torch.float16,
+        "device_map": "auto",
+        "load_in_8bit": True  # 需要安装 bitsandbytes
+    }
+else:
+    device = "cpu"
+    load_kwargs = {
+        "torch_dtype": torch.float32,
+        "device_map": "balanced"
+    }
 
 # SD v1.5
 sd_v15 = StableDiffusionPipeline.from_pretrained(
@@ -30,6 +39,7 @@ sd_v15 = StableDiffusionPipeline.from_pretrained(
     **load_kwargs
 )
 sd_v15.scheduler = DPMSolverMultistepScheduler.from_config(sd_v15.scheduler.config)
+sd_v15 = sd_v15.to(device)
 
 # SD XL
 sd_xl = StableDiffusionPipeline.from_pretrained(
@@ -37,6 +47,7 @@ sd_xl = StableDiffusionPipeline.from_pretrained(
     **load_kwargs
 )
 sd_xl.scheduler = DPMSolverMultistepScheduler.from_config(sd_xl.scheduler.config)
+sd_xl = sd_xl.to(device)
 
 # 可选：语音输入模块，使用 Whisper
 asr = pipeline(
@@ -122,3 +133,4 @@ with gr.Blocks(title="Prompt-to-Image Generator") as demo:
 # Step 4: 启动应用
 if __name__ == "__main__":
     demo.launch(server_name="0.0.0.0", server_port=7860)
+