Update model_pipelines.py

model_pipelines.py

@@ -1,28 +1,48 @@
+from diffusers import StableDiffusionPipeline, DiffusionPipeline
 import torch
-from diffusers import StableDiffusionPipeline
+from PIL import Image
+import time
 
-device = "cuda" if torch.cuda.is_available() else "cpu"
+device = "cpu"
+torch_dtype = torch.float32  # CPU上使用float32更稳定
 
 def load_pipelines():
-    model_ids = {
-        "sd_v1_5": "runwayml/stable-diffusion-v1-5",
-        "openjourney_v4": "prompthero/openjourney-v4",
-        "realistic_vision": "SG161222/Realistic_Vision_V5.1"
+    print("正在加载模型(CPU模式)...")
+    
+    # 使用较小的模型或优化配置
+    models = {
+        "sd_v1_5": StableDiffusionPipeline.from_pretrained(
+            "runwayml/stable-diffusion-v1-5",
+            torch_dtype=torch_dtype,
+            safety_checker=None,
+            requires_safety_checker=False
+        ).to(device),
+        
+        "openjourney_v4": DiffusionPipeline.from_pretrained(
+            "prompthero/openjourney-v4",
+            torch_dtype=torch_dtype,
+            safety_checker=None
+        ).to(device),
+        
+        "ldm_256": DiffusionPipeline.from_pretrained(
+            "CompVis/ldm-text2im-large-256",
+            torch_dtype=torch_dtype
+        ).to(device)
     }
-    pipes = {}
-    for name, mid in model_ids.items():
-        pipe = StableDiffusionPipeline.from_pretrained(
-            mid,
-            torch_dtype=torch.float16 if device == "cuda" else torch.float32
-        )
-        pipe = pipe.to(device)
-        pipe.enable_attention_slicing()
-        pipes[name] = pipe
-    return pipes
+    
+    # 启用内存优化
+    for model in models.values():
+        model.enable_attention_slicing()
+        
+    return models
 
-def generate_all(pipes, prompt):
+def generate_all(pipelines, prompt, steps=20):
     results = {}
-    for name, pipe in pipes.items():
-        img = pipe(prompt, guidance_scale=7.5, num_inference_steps=30).images[0]
-        results[name] = img
-    return results
+    for name, pipe in pipelines.items():
+        print(f"正在用 {name} 生成图像...")
+        start = time.time()
+        result = pipe(prompt, num_inference_steps=steps).images[0]
+        gen_time = time.time() - start
+        print(f"{name} 生成完成，耗时 {gen_time:.2f}秒")
+        results[name] = result
+    return results