Update models/model_manager.py

models/model_manager.py

@@ -1,10 +1,186 @@
-def caption_image(self, image) -> str:
-    """对图像生成描述"""
-    if self.caption_model is None or self.caption_processor is None:
+# models/model_manager.py
+import torch
+from PIL import Image
+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()
-    
-    inputs = self.caption_processor(images=image, return_tensors="pt").to(self.device)
-    with torch.no_grad():
-        output = self.caption_model.generate(**inputs, max_length=50)
-    
-    return self.caption_processor.decode(output[0], skip_special_tokens=True)
+        self.load_clip_model()
+        self.load_sd_pipeline()
+        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...")
+            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  # 可按需配置安全检查器
+            )
+            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()
+        
+        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()
+
+        # 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()
+
+        # 输入的 image 应该是 PIL Image 格式的控制图（比如人体姿态图）
+        result = self.controlnet_pipeline(
+            prompt=prompt,
+            image=image,
+            negative_prompt=negative_prompt,
+            num_inference_steps=num_inference_steps,
+            guidance_scale=guidance_scale,
+        )
+        return result.images[0]
+
+    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
+            torch.cuda.empty_cache()
+            import gc
+            gc.collect()
+            logger.info("显存清理完成")
+        except Exception as e:
+            logger.error(f"清理显存失败: {e}")