Update models/model_manager.py

models/model_manager.py

@@ -191,18 +191,85 @@ class ModelManager:
         return Image.new('RGB', (width, height), color=color)
 
     def cleanup(self):
-        """释放模型占用显存和缓存"""
-        logger.info("释放模型占用显存和缓存...")
+        """仅清理显存缓存，保留模型以避免重新加载"""
+        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("显存和缓存清理完成")
+            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')
+            
+            if torch.cuda.is_available():
+                torch.cuda.empty_cache()
+            
+            logger.info("模型已移动到CPU")
+        except Exception as e:
+            logger.error(f"移动模型到CPU失败: {e}")
+
+    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}")
+
+    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 torch.cuda.is_available():
+                torch.cuda.empty_cache()
+            
+            # 重新加载
+            self.load_all_models()
+            logger.info("所有模型重新加载完成")
+            
+        except Exception as e:
+            logger.error(f"强制重新加载模型失败: {e}")