Update models/model_manager.py

models/model_manager.py

@@ -6,7 +6,9 @@ from diffusers import StableDiffusionPipeline, ControlNetModel, StableDiffusionC
 import os
 import logging
 import time
-import random  # 补充导入
+import random
+import gc
+from functools import lru_cache
 
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
@@ -16,10 +18,10 @@ class ModelManager:
         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",
+            "clip_model": "openai/clip-vit-large-patch14", 
             "sd_model": "runwayml/stable-diffusion-v1-5",
             "controlnet_model": "lllyasviel/control_v11p_sd15_openpose"
         }
@@ -32,244 +34,301 @@ class ModelManager:
         self.sd_pipeline = None
         self.controlnet = None
         self.controlnet_pipeline = None
-
+        
+        # 性能优化设置
+        self.torch_dtype = torch.float16 if self.device == "cuda" else torch.float32
+        self.enable_attention_slicing = True
+        self.enable_cpu_offload = False  # 16GB显存应该够用
+        
         # 预加载所有模型
         self.load_all_models()
 
+    def optimize_memory_usage(self):
+        """内存优化设置"""
+        if torch.cuda.is_available():
+            # 启用内存优化
+            torch.backends.cudnn.benchmark = True
+            torch.backends.cuda.matmul.allow_tf32 = True
+            torch.backends.cudnn.allow_tf32 = True
+
     def load_all_models(self):
-        self.load_caption_model()
-        self.load_clip_model()
-        self.load_sd_pipeline()
-        self.load_controlnet_pipeline()
+        """按顺序加载所有模型，优化显存使用"""
+        self.optimize_memory_usage()
+        
+        try:
+            self.load_caption_model()
+            self.load_clip_model()
+            self.load_sd_pipeline()
+            self.load_controlnet_pipeline()
+            
+            logger.info("所有模型加载完成")
+            if torch.cuda.is_available():
+                logger.info(f"GPU显存使用: {torch.cuda.memory_allocated()/1024**3:.2f}GB / {torch.cuda.max_memory_allocated()/1024**3:.2f}GB")
+                
+        except Exception as e:
+            logger.error(f"模型加载过程中出错: {e}")
+            raise
 
     def load_caption_model(self):
+        """加载BLIP图像描述模型"""
         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
+                torch_dtype=self.torch_dtype,
+                low_cpu_mem_usage=True
             ).to(self.device)
+            
+            # 启用内存优化
+            if hasattr(self.caption_model, 'enable_attention_slicing'):
+                self.caption_model.enable_attention_slicing()
+                
+            self.caption_model.eval()
             logger.info("BLIP 模型加载完成")
+            
         except Exception as e:
             logger.error(f"BLIP 模型加载失败: {e}")
+            self.caption_model = None
+            self.caption_processor = None
 
     def load_clip_model(self):
+        """加载CLIP风格分析模型"""
         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
+                torch_dtype=self.torch_dtype
             ).to(self.device)
+            
+            self.clip_model.eval()
             logger.info("CLIP 模型加载完成")
+            
         except Exception as e:
             logger.error(f"CLIP 模型加载失败: {e}")
+            self.clip_model = None
+            self.clip_processor = None
 
     def load_sd_pipeline(self):
+        """加载Stable Diffusion Pipeline"""
         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,
+                torch_dtype=self.torch_dtype,
                 cache_dir="/tmp/models",
                 safety_checker=None,
-                use_safetensors=True
+                requires_safety_checker=False,
+                use_safetensors=True,
+                low_cpu_mem_usage=True
             )
+            
+            # 优化设置
             self.sd_pipeline = self.sd_pipeline.to(self.device)
-            self.sd_pipeline.scheduler = EulerAncestralDiscreteScheduler.from_config(self.sd_pipeline.scheduler.config)
+            self.sd_pipeline.scheduler = EulerAncestralDiscreteScheduler.from_config(
+                self.sd_pipeline.scheduler.config
+            )
+            
+            # 启用内存优化
+            if self.enable_attention_slicing:
+                self.sd_pipeline.enable_attention_slicing()
+            
+            # 启用内存高效attention（如果可用）
+            try:
+                self.sd_pipeline.enable_xformers_memory_efficient_attention()
+                logger.info("启用了xformers内存优化")
+            except:
+                logger.info("xformers不可用，使用默认attention")
+            
+            # 启用VAE slicing以节省显存
+            self.sd_pipeline.enable_vae_slicing()
+            
             logger.info("Stable Diffusion Pipeline 加载完成")
+            
         except Exception as e:
             logger.error(f"Stable Diffusion Pipeline 加载失败: {e}")
+            self.sd_pipeline = None
 
     def load_controlnet_pipeline(self):
+        """加载ControlNet Pipeline"""
         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
+                torch_dtype=self.torch_dtype,
+                low_cpu_mem_usage=True
             ).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
+                torch_dtype=self.torch_dtype,
+                safety_checker=None,
+                requires_safety_checker=False,
+                low_cpu_mem_usage=True
             ).to(self.device)
 
-            self.controlnet_pipeline.scheduler = EulerAncestralDiscreteScheduler.from_config(self.controlnet_pipeline.scheduler.config)
+            # 优化设置
+            self.controlnet_pipeline.scheduler = EulerAncestralDiscreteScheduler.from_config(
+                self.controlnet_pipeline.scheduler.config
+            )
+            
+            # 内存优化
+            if self.enable_attention_slicing:
+                self.controlnet_pipeline.enable_attention_slicing()
+            
+            try:
+                self.controlnet_pipeline.enable_xformers_memory_efficient_attention()
+                logger.info("ControlNet启用了xformers内存优化")
+            except:
+                logger.info("ControlNet使用默认attention")
+                
+            self.controlnet_pipeline.enable_vae_slicing()
+            
             logger.info("ControlNet Pipeline 加载完成")
+            
         except Exception as e:
             logger.error(f"ControlNet Pipeline 加载失败: {e}")
+            self.controlnet = None
+            self.controlnet_pipeline = None
 
-    # 下面是真正调用模型的接口
-
+    @torch.no_grad()  # 禁用梯度计算节省显存
     def generate_caption(self, image):
         """使用BLIP模型生成图像描述"""
         if self.caption_model is None or self.caption_processor is None:
             self.load_caption_model()
+            if self.caption_model is None:
+                return "时尚服装设计作品"
+        
+        try:
+            # 预处理图像
+            if image.mode != 'RGB':
+                image = image.convert('RGB')
+            
+            # 调整图像大小以节省显存
+            if image.width > 512 or image.height > 512:
+                image.thumbnail((512, 512), Image.Resampling.LANCZOS)
+            
+            inputs = self.caption_processor(images=image, return_tensors="pt").to(self.device)
+            
+            # 生成描述
+            outputs = self.caption_model.generate(
+                **inputs, 
+                max_length=50,
+                num_beams=4,
+                temperature=0.7,
+                do_sample=True
+            )
+            
+            caption = self.caption_processor.decode(outputs[0], skip_special_tokens=True)
+            
+            # 清理显存
+            del inputs, outputs
+            if torch.cuda.is_available():
+                torch.cuda.empty_cache()
+            
+            return caption
+            
+        except Exception as e:
+            logger.error(f"图像描述生成失败: {e}")
+            return "时尚服装设计作品"
 
-        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
-
+    @torch.no_grad()
     def analyze_style(self, image):
         """使用CLIP模型分析服装风格"""
         if self.clip_model is None or self.clip_processor is None:
             self.load_clip_model()
+            if self.clip_model is None:
+                return {"时尚潮流": 0.8, "现代风格": 0.6}
         
-        styles = ["商务正装", "休闲风", "运动风", "时尚潮流", "复古风", "街头风", "优雅风"]
-        
-        inputs = self.clip_processor(
-            text=styles,
-            images=image,
-            return_tensors="pt",
-            padding=True
-        ).to(self.device)
-        
-        with torch.no_grad():
+        try:
+            # 风格标签 - 使用英文避免token问题
+            style_labels = [
+                "business formal suit professional attire",
+                "casual comfortable everyday wear",
+                "athletic sportswear activewear",
+                "fashion trendy modern stylish",
+                "vintage retro classic style", 
+                "streetwear urban contemporary",
+                "elegant sophisticated refined"
+            ]
+            
+            style_names = ["商务正装", "休闲风", "运动风", "时尚潮流", "复古风", "街头风", "优雅风"]
+            
+            # 预处理图像
+            if image.mode != 'RGB':
+                image = image.convert('RGB')
+            
+            # 调整图像大小
+            if image.width > 224 or image.height > 224:
+                image.thumbnail((224, 224), Image.Resampling.LANCZOS)
+            
+            # 处理输入
+            inputs = self.clip_processor(
+                text=style_labels,
+                images=image,
+                return_tensors="pt",
+                padding=True,
+                truncation=True,
+                max_length=77  # CLIP的最大长度
+            ).to(self.device)
+            
+            # 获取相似度分数
             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:
-            prompt = f"{prompt}, based on reference design"
-        
-        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 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:
-            # 只清理缓存，不删除模型
-            if torch.cuda.is_available():
-                torch.cuda.empty_cache()
             
-            logger.info("显存缓存清理完成")
-            
-        except Exception as e:
-            logger.error(f"清理显存失败: {e}")
-
-    def move_models_to_cpu(self):
-        """将模型移动到CPU以释放显存"""
-        try:
-            logger.info("将模型移动到CPU...")
-            if self.caption_model is not None:
-                self.caption_model = self.caption_model.to('cpu')
-            if self.clip_model is not None:
-                self.clip_model = self.clip_model.to('cpu')
-            if self.sd_pipeline is not None:
-                self.sd_pipeline = self.sd_pipeline.to('cpu')
-            if self.controlnet_pipeline is not None:
-                self.controlnet_pipeline = self.controlnet_pipeline.to('cpu')
+            # 构建结果
+            style_scores = {name: float(prob) for name, prob in zip(style_names, probs)}
             
+            # 清理显存
+            del inputs, outputs
             if torch.cuda.is_available():
                 torch.cuda.empty_cache()
             
-            logger.info("模型已移动到CPU")
+            return style_scores
+            
         except Exception as e:
-            logger.error(f"移动模型到CPU失败: {e}")
+            logger.error(f"风格分析失败: {e}")
+            return {"时尚潮流": 0.8, "现代风格": 0.6}
 
-    def move_models_to_gpu(self):
-        """将模型移回GPU"""
-        try:
-            logger.info("将模型移动到GPU...")
-            if self.caption_model is not None:
-                self.caption_model = self.caption_model.to(self.device)
-            if self.clip_model is not None:
-                self.clip_model = self.clip_model.to(self.device)
-            if self.sd_pipeline is not None:
-                self.sd_pipeline = self.sd_pipeline.to(self.device)
-            if self.controlnet_pipeline is not None:
-                self.controlnet_pipeline = self.controlnet_pipeline.to(self.device)
-            
-            logger.info("模型已移回GPU")
-        except Exception as e:
-            logger.error(f"移动模型到GPU失败: {e}")
+    @torch.no_grad()
+    def generate_image(self, prompt, negative_prompt=None, num_inference_steps=25, guidance_scale=7.5, width=512, height=512, **kwargs):
+        """使用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)
 
-    def force_reload_all_models(self):
-        """强制重新加载所有模型"""
-        logger.info("强制重新加载所有模型...")
         try:
-            # 先清理
-            if hasattr(self, 'caption_model') and self.caption_model is not None:
-                del self.caption_model
-                del self.caption_processor
-                self.caption_model = None
-                self.caption_processor = None
-            if hasattr(self, 'clip_model') and self.clip_model is not None:
-                del self.clip_model
-                del self.clip_processor
-                self.clip_model = None
-                self.clip_processor = None
-            if hasattr(self, 'sd_pipeline') and self.sd_pipeline is not None:
-                del self.sd_pipeline
-                self.sd_pipeline = None
-            if hasattr(self, 'controlnet_pipeline') and self.controlnet_pipeline is not None:
-                del self.controlnet
-                del self.controlnet_pipeline
-                self.controlnet = None
-                self.controlnet_pipeline = None
+            # 优化参数
+            if negative_prompt is None:
+                negative_prompt = "blurry, low quality, distorted, text, watermark, ugly, deformed"
             
-            if torch.cuda.is_available():
-                torch.cuda.empty_cache()
+            # 确保尺寸是8的倍数
+            width = (width // 8) * 8
+            height = (height // 8) * 8
             
-            # 重新加载
-            self.load_all_models()
-            logger.info("所有模型重新加载完成")
-            
-        except Exception as e:
-            logger.error(f"强制重新加载模型失败: {e}")
+            # 生成图像
+            result = self.sd_pipeline(
+                prompt=prompt,
+                negative_prompt=negative_prompt,
+                num_inference_steps=num_inference_steps,