Update models/model_manager.py

models/model_manager.py

@@ -1,28 +1,205 @@
-# 在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:
+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}")
+
+        # 模型配置 - 使用更精细的3D模型
+        self.model_config = {
+            "caption_model": "Salesforce/blip-image-captioning-large",
+            "clip_model": "openai/clip-vit-large-patch14",
+            "sd_model": "runwayml/stable-diffusion-v1-5",
+            "controlnet_model": "lllyasviel/control_v11p_sd15_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()
         self.load_controlnet_pipeline()
-    
-    try:
+
+    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"],
+                torch_dtype=torch.float16 if self.device=="cuda" else torch.float32,
+                cache_dir="/tmp/models",
+                safety_checker=None,
+                use_safetensors=True
+            )
+            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):
+        """使用BLIP模型生成图像描述"""
+        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):
+        """使用CLIP模型分析服装风格"""
+        if self.clip_model is None or self.clip_processor is None:
+            self.load_clip_model()
+        
+        # 定义服装风格类别
+        styles = ["商务正装", "休闲风", "运动风", "时尚潮流", "复古风", "街头风", "优雅风"]
+        
+        # 准备输入
+        inputs = self.clip_processor(
+            text=styles,
+            images=image,
+            return_tensors="pt",
+            padding=True
+        ).to(self.device)
+        
+        # 获取特征
+        with torch.no_grad():
+            outputs = self.clip_model(**inputs)
+            logits_per_image = outputs.logits_per_image
+            probs = logits_per_image.softmax(dim=1).cpu().numpy()[0]
+        
+        # 转换为分数字典
+        style_scores = {style: float(prob) for style, prob in zip(styles, probs)}
+        return style_scores
+
+    def generate_image(self, prompt, negative_prompt=None, num_inference_steps=25, guidance_scale=7.5, width=512, height=512):
+        """使用Stable Diffusion生成设计图像"""
+        if self.sd_pipeline is None:
+            self.load_sd_pipeline()
+            if self.sd_pipeline is None:
+                logger.error("无法生成图像：Stable Diffusion 模型未加载")
+                return self.create_placeholder_image(width, height)
+
+        # 生成图像
+        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, reference_image=None, negative_prompt=None, num_inference_steps=30, guidance_scale=8.0):
+        """使用ControlNet生成3D试穿效果 - 更精细的模型"""
+        if self.controlnet_pipeline is None:
+            self.load_controlnet_pipeline()
+            if self.controlnet_pipeline is None:
+                logger.error("无法生成3D试穿：ControlNet 模型未加载")
+                return self.create_placeholder_image(512, 768)
+        
         # 如果有参考图像，将其融入提示词
         if reference_image is not None:
-            # 简化的参考图像描述（实际应用中可用CLIP生成详细描述）
-            ref_desc = "参考设计风格"
-            prompt = f"{prompt}, {ref_desc}"
+            # 这里可以添加将参考图像融入提示词的逻辑
+            prompt = f"{prompt}, based on reference design"
         
         # 生成图像
         result = self.controlnet_pipeline(
             prompt=prompt,
-            image=image,
+            image=image,  # 控制图像（如人体姿态）
             negative_prompt=negative_prompt,
             num_inference_steps=num_inference_steps,
             guidance_scale=guidance_scale,
         )
         return result.images[0]
-    except Exception as e:
-        logger.error(f"ControlNet生成失败: {e}")
-        # 创建占位图像
-        return Image.new('RGB', (512, 768), color=(180, 180, 180))
+
+    def create_placeholder_image(self, width, height):
+        """创建占位图像"""
+        color = (random.randint(120, 200), random.randint(120, 200), random.randint(120, 200)
+        return Image.new('RGB', (width, height), color=color)
+
+    def cleanup(self):
+        logger.info("释放模型占用显存和缓存...")
+        try:
+            # ...清理代码不变...
+        except Exception as e:
+            logger.error(f"清理显存失败: {e}")