Update models/model_manager.py

models/model_manager.py

@@ -1,195 +1,19 @@
-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
-import logging
-import time
-
-logging.basicConfig(level=logging.INFO)
-logger = logging.getLogger(__name__)
-
-class ModelManager:
-    def __init__(self):
-        self.device = "cuda" if torch.cuda.is_available() else "cpu"
-        logger.info(f"使用设备: {self.device}")
-
-        # 模型配置（更新了 SD 模型路径）
-        self.model_config = {
-            "caption_model": "Salesforce/blip-image-captioning-base",
-            "clip_model": "openai/clip-vit-base-patch32",
-            "sd_model": "runwayml/stable-diffusion-v1-5",
-            "controlnet_model": "lllyasviel/sd-controlnet-openpose"
-        }
-
-        # 模型容器
-        self.caption_processor = None
-        self.caption_model = None
-
-        self.clip_processor = None
-        self.clip_model = None
-
-        self.sd_pipeline = None
-        self.controlnet = None
-        self.controlnet_pipeline = None
-
-        # 预加载所有模型
-        self.load_all_models()
-
-    def load_all_models(self):
-        self.load_caption_model()
-        self.load_clip_model()
-        self.load_sd_pipeline()
+# 在ModelManager类中添加新方法
+def generate_controlnet_image(self, image, prompt, reference_image=None, 
+                             negative_prompt=None, num_inference_steps=30, 
+                             guidance_scale=8.0):
+    """使用ControlNet生成图像，支持参考图像"""
+    if self.controlnet_pipeline is None:
         self.load_controlnet_pipeline()
-
-    def load_caption_model(self):
-        try:
-            logger.info("加载 BLIP 图像描述模型...")
-            self.caption_processor = BlipProcessor.from_pretrained(
-                self.model_config["caption_model"],
-                cache_dir="/tmp/models"
-            )
-            self.caption_model = BlipForConditionalGeneration.from_pretrained(
-                self.model_config["caption_model"],
-                cache_dir="/tmp/models",
-                torch_dtype=torch.float16 if self.device=="cuda" else torch.float32
-            ).to(self.device)
-            logger.info("BLIP 模型加载完成")
-        except Exception as e:
-            logger.error(f"BLIP 模型加载失败: {e}")
-
-    def load_clip_model(self):
-        try:
-            logger.info("加载 CLIP 模型...")
-            self.clip_processor = CLIPProcessor.from_pretrained(
-                self.model_config["clip_model"],
-                cache_dir="/tmp/models"
-            )
-            self.clip_model = CLIPModel.from_pretrained(
-                self.model_config["clip_model"],
-                cache_dir="/tmp/models",
-                torch_dtype=torch.float16 if self.device=="cuda" else torch.float32
-            ).to(self.device)
-            logger.info("CLIP 模型加载完成")
-        except Exception as e:
-            logger.error(f"CLIP 模型加载失败: {e}")
-
-    def load_sd_pipeline(self):
-        try:
-            logger.info("加载 Stable Diffusion Pipeline...")
-            
-            # 尝试加载原始模型
-            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)
-            logger.info("Stable Diffusion Pipeline 加载完成")
-        except Exception as e:
-            logger.error(f"Stable Diffusion Pipeline 加载失败: {e}")
-
-    def load_controlnet_pipeline(self):
-        try:
-            logger.info("加载 ControlNet 模型和 Pipeline...")
-            self.controlnet = ControlNetModel.from_pretrained(
-                self.model_config["controlnet_model"],
-                cache_dir="/tmp/models",
-                torch_dtype=torch.float16 if self.device=="cuda" else torch.float32
-            ).to(self.device)
-
-            self.controlnet_pipeline = StableDiffusionControlNetPipeline.from_pretrained(
-                self.model_config["sd_model"],
-                controlnet=self.controlnet,
-                cache_dir="/tmp/models",
-                torch_dtype=torch.float16 if self.device=="cuda" else torch.float32,
-                safety_checker=None
-            ).to(self.device)
-
-            self.controlnet_pipeline.scheduler = EulerAncestralDiscreteScheduler.from_config(self.controlnet_pipeline.scheduler.config)
-            logger.info("ControlNet Pipeline 加载完成")
-        except Exception as e:
-            logger.error(f"ControlNet Pipeline 加载失败: {e}")
-
-    # 下面是真正调用模型的接口
-
-    def generate_caption(self, image):
-        if self.caption_model is None or self.caption_processor is None:
-            self.load_caption_model()
-
-        inputs = self.caption_processor(images=image, return_tensors="pt").to(self.device)
-        with torch.no_grad():
-            outputs = self.caption_model.generate(**inputs, max_length=50)
-        caption = self.caption_processor.decode(outputs[0], skip_special_tokens=True)
-        return caption
-
-    def analyze_style(self, image):
-        if self.clip_model is None or self.clip_processor is None:
-            self.load_clip_model()
+    
+    try:
+        # 如果有参考图像，将其融入提示词
+        if reference_image is not None:
+            # 简化的参考图像描述（实际应用中可用CLIP生成详细描述）
+            ref_desc = "参考设计风格"
+            prompt = f"{prompt}, {ref_desc}"
         
-        inputs = self.clip_processor(images=image, return_tensors="pt").to(self.device)
-        with torch.no_grad():
-            outputs = self.clip_model.get_image_features(**inputs)
-        features = outputs.cpu().numpy()[0]
-        # 简单归一化（范例）
-        norm = features / (np.linalg.norm(features) + 1e-10)
-        style_score = { "clip_feature_vector": norm }
-        return style_score
-
-    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(
-            prompt=prompt,
-            negative_prompt=negative_prompt,
-            num_inference_steps=num_inference_steps,
-            guidance_scale=guidance_scale,
-            height=height,
-            width=width
-        )
-        return result.images[0]
-
-    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(
             prompt=prompt,
             image=image,
@@ -198,28 +22,7 @@ class ModelManager:
             guidance_scale=guidance_scale,
         )
         return result.images[0]
-
-    def cleanup(self):
-        logger.info("释放模型占用显存和缓存...")
-        try:
-            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}")
+    except Exception as e:
+        logger.error(f"ControlNet生成失败: {e}")
+        # 创建占位图像
+        return Image.new('RGB', (512, 768), color=(180, 180, 180))