Update model_pipelines.py

model_pipelines.py

@@ -2,47 +2,63 @@ from diffusers import StableDiffusionPipeline, DiffusionPipeline
 import torch
 from PIL import Image
 import time
+import logging
+from tqdm import tqdm
 
-device = "cpu"
-torch_dtype = torch.float32  # CPU上使用float32更稳定
+logging.basicConfig(format='%(asctime)s - %(message)s', level=logging.INFO)
+logger = logging.getLogger(__name__)
 
-def load_pipelines():
-    print("正在加载模型(CPU模式)...")
-    
-    # 使用新版API加载模型
-    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),
+class ModelWrapper:
+    def __init__(self, pipe, name):
+        self.pipe = pipe
+        self.name = name
         
-        "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)
-    }
+    def generate(self, prompt, steps=20):
+        logger.info(f"Generating with {self.name}...")
+        try:
+            start = time.time()
+            result = self.pipe(
+                prompt,
+                num_inference_steps=steps,
+                guidance_scale=7.5,
+                output_type="pil"
+            ).images[0]
+            logger.info(f"{self.name} generated in {time.time()-start:.2f}s")
+            return result
+        except Exception as e:
+            logger.error(f"{self.name} failed: {str(e)}")
+            return Image.new("RGB", (512, 512), color="#FF0000")
+
+def load_pipelines():
+    configs = [
+        ("sd_v1_5", "runwayml/stable-diffusion-v1-5", StableDiffusionPipeline, 
+         {"safety_checker": None}),
+        ("openjourney_v4", "prompthero/openjourney-v4", DiffusionPipeline, {}),
+        ("ldm_256", "CompVis/ldm-text2im-large-256", DiffusionPipeline, {})
+    ]
     
-    # 内存优化配置
-    for model in models.values():
-        model.enable_attention_slicing()
-        
+    models = {}
+    for name, repo, pipe_cls, kwargs in tqdm(configs, desc="Loading models"):
+        try:
+            pipe = pipe_cls.from_pretrained(
+                repo,
+                torch_dtype=torch.float32,
+                **kwargs
+            )
+            pipe = pipe.to("cpu")
+            pipe.enable_attention_slicing()
+            models[name] = ModelWrapper(pipe, name)
+        except Exception as e:
+            logger.error(f"Failed to load {name}: {str(e)}")
+            continue
+            
     return models
 
-def generate_all(pipelines, prompt, steps=20):
-    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
+def generate_all(models, prompt, steps=20):
+    if not models:
+        raise ValueError("No models loaded")
+        
+    return {
+        name: model.generate(prompt, steps)
+        for name, model in models.items()
+    }