Update models/model_manager.py

models/model_manager.py

@@ -1,6 +1,6 @@
-# models/model_manager.py
 import torch
 from PIL import Image
+import numpy as np  # 添加缺失的导入
 from transformers import BlipProcessor, BlipForConditionalGeneration, CLIPProcessor, CLIPModel
 from diffusers import StableDiffusionPipeline, ControlNetModel, StableDiffusionControlNetPipeline, EulerAncestralDiscreteScheduler
 import os
@@ -19,7 +19,7 @@ class ModelManager:
         self.model_config = {
             "caption_model": "Salesforce/blip-image-captioning-base",
             "clip_model": "openai/clip-vit-base-patch32",
-            "sd_model": "runwayml/stable-diffusion-v1-5",  # 这里用原版，可替换为镜像
+            "sd_model": "runwayml/stable-diffusion-v1-5",
             "controlnet_model": "lllyasviel/sd-controlnet-openpose"
         }
 
@@ -78,13 +78,35 @@ class ModelManager:
     def load_sd_pipeline(self):
         try:
             logger.info("加载 Stable Diffusion Pipeline...")
-            self.sd_pipeline = StableDiffusionPipeline.from_pretrained(
-                self.model_config["sd_model"],
-                revision="fp16" if self.device=="cuda" else None,
-                torch_dtype=torch.float16 if self.device=="cuda" else torch.float32,
-                cache_dir="/tmp/models",
-                safety_checker=None  # 可按需配置安全检查器
-            )
+            
+            # 尝试加载原始模型
+            try:
+                self.sd_pipeline = StableDiffusionPipeline.from_pretrained(
+                    self.model_config["sd_model"],
+                    torch_dtype=torch.float16 if self.device=="cuda" else torch.float32,
+                    cache_dir="/tmp/models",
+                    safety_checker=None,
+                    use_safetensors=True
+                )
+            except Exception as e:
+                logger.warning(f"原始模型加载失败: {e}")
+                logger.info("尝试加载本地缓存的模型...")
+                
+                # 定义本地模型路径
+                local_model_path = "./local_models/stable-diffusion-v1-5"
+                
+                # 检查本地模型是否存在
+                if os.path.exists(local_model_path):
+                    self.sd_pipeline = StableDiffusionPipeline.from_pretrained(
+                        local_model_path,
+                        torch_dtype=torch.float16 if self.device=="cuda" else torch.float32,
+                        safety_checker=None
+                    )
+                    logger.info("使用本地缓存的 Stable Diffusion 模型")
+                else:
+                    logger.error("没有可用的本地模型")
+                    raise
+            
             self.sd_pipeline = self.sd_pipeline.to(self.device)
             # 用更高效的调度器
             self.sd_pipeline.scheduler = EulerAncestralDiscreteScheduler.from_config(self.sd_pipeline.scheduler.config)
@@ -142,6 +164,11 @@ class ModelManager:
     def generate_image(self, prompt, negative_prompt=None, num_inference_steps=25, guidance_scale=7.5, width=512, height=512):
         if self.sd_pipeline is None:
             self.load_sd_pipeline()
+            if self.sd_pipeline is None:
+                logger.error("无法生成图像：Stable Diffusion 模型未加载")
+                # 创建占位图像
+                color = (180, 180, 180)
+                return Image.new('RGB', (width, height), color=color)
 
         # Stable Diffusion 生成图像
         result = self.sd_pipeline(
@@ -157,6 +184,10 @@ class ModelManager:
     def generate_controlnet_image(self, image, prompt, negative_prompt=None, num_inference_steps=30, guidance_scale=8.0):
         if self.controlnet_pipeline is None:
             self.load_controlnet_pipeline()
+            if self.controlnet_pipeline is None:
+                logger.error("无法生成图像：ControlNet 模型未加载")
+                # 创建占位图像
+                return Image.new('RGB', (512, 768), color=(180, 180, 180))
 
         # 输入的 image 应该是 PIL Image 格式的控制图（比如人体姿态图）
         result = self.controlnet_pipeline(
@@ -171,16 +202,24 @@ class ModelManager:
     def cleanup(self):
         logger.info("释放模型占用显存和缓存...")
         try:
-            del self.caption_model
-            del self.caption_processor
-            del self.clip_model
-            del self.clip_processor
-            del self.sd_pipeline
-            del self.controlnet
-            del self.controlnet_pipeline
+            if hasattr(self, 'caption_model') and self.caption_model is not None:
+                del self.caption_model
+            if hasattr(self, 'caption_processor') and self.caption_processor is not None:
+                del self.caption_processor
+            if hasattr(self, 'clip_model') and self.clip_model is not None:
+                del self.clip_model
+            if hasattr(self, 'clip_processor') and self.clip_processor is not None:
+                del self.clip_processor
+            if hasattr(self, 'sd_pipeline') and self.sd_pipeline is not None:
+                del self.sd_pipeline
+            if hasattr(self, 'controlnet') and self.controlnet is not None:
+                del self.controlnet
+            if hasattr(self, 'controlnet_pipeline') and self.controlnet_pipeline is not None:
+                del self.controlnet_pipeline
+            
             torch.cuda.empty_cache()
             import gc
             gc.collect()
             logger.info("显存清理完成")
         except Exception as e:
-            logger.error(f"清理显存失败: {e}")
+            logger.error(f"清理显存失败: {e}")