Update models/model_manager.py

models/model_manager.py

@@ -6,6 +6,7 @@ from diffusers import StableDiffusionPipeline, ControlNetModel, StableDiffusionC
 import os
 import logging
 import time
+import random  # 补充导入
 
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
@@ -129,10 +130,8 @@ class ModelManager:
         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,
@@ -140,13 +139,11 @@ class ModelManager:
             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
 
@@ -158,7 +155,6 @@ class ModelManager:
                 logger.error("无法生成图像：Stable Diffusion 模型未加载")
                 return self.create_placeholder_image(width, height)
 
-        # 生成图像
         result = self.sd_pipeline(
             prompt=prompt,
             negative_prompt=negative_prompt,
@@ -177,15 +173,12 @@ class ModelManager:
                 logger.error("无法生成3D试穿：ControlNet 模型未加载")
                 return self.create_placeholder_image(512, 768)
         
-        # 如果有参考图像，将其融入提示词
         if reference_image is not None:
-            # 这里可以添加将参考图像融入提示词的逻辑
             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,
@@ -197,20 +190,19 @@ class ModelManager:
         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:
-        # 清理模型引用
-        del self.caption_model
-        del self.clip_model
-        del self.sd_pipeline
-        del self.controlnet
-        del self.controlnet_pipeline
-        
-        # 清理缓存
-        if torch.cuda.is_available():
-            torch.cuda.empty_cache()
-        
-        logger.info("显存和缓存清理完成")
-    except Exception as e:
-        logger.error(f"清理显存失败: {e}")
+    def cleanup(self):
+        """释放模型占用显存和缓存"""
+        logger.info("释放模型占用显存和缓存...")
+        try:
+            del self.caption_model
+            del self.clip_model
+            del self.sd_pipeline
+            del self.controlnet
+            del self.controlnet_pipeline
+            
+            if torch.cuda.is_available():
+                torch.cuda.empty_cache()
+            
+            logger.info("显存和缓存清理完成")
+        except Exception as e:
+            logger.error(f"清理显存失败: {e}")