Update models/model_manager.py

models/model_manager.py

@@ -1,77 +1,35 @@
-# models/model_manager.py - 增强版
-import torch
-from PIL import Image
-from transformers import (
-    BlipProcessor, 
-    BlipForConditionalGeneration,
-    CLIPProcessor, 
-    CLIPModel
-)
-from diffusers import (
-    StableDiffusionPipeline,
-    StableDiffusionControlNetPipeline,
-    ControlNetModel,
-    EulerAncestralDiscreteScheduler
-)
-import numpy as np
-import gc
-import os
-import logging
-import time
-from typing import Optional, Dict, List, Tuple
+def load_caption_model(self):
+    """加载图像描述模型"""
+    if self.caption_model is None:
+        start_time = time.time()
+        logger.info("正在加载BLIP图像描述模型...")
+        
+        try:
+            self.caption_processor = BlipProcessor.from_pretrained(
+                self.model_config["caption_model"],
+                cache_dir=self.cache_dir
+            )
+            
+            self.caption_model = BlipForConditionalGeneration.from_pretrained(
+                self.model_config["caption_model"],
+                cache_dir=self.cache_dir,
+                torch_dtype=torch.float16 if self.device == "cuda" else torch.float32
+            ).to(self.device)
+            
+            self.load_times["caption_model"] = time.time() - start_time
+            self.last_used["caption_model"] = time.time()
+            logger.info(f"BLIP图像描述模型加载完成，用时 {self.load_times['caption_model']:.2f} 秒")
+        except Exception as e:
+            logger.error(f"加载BLIP图像描述模型失败: {e}")
 
-# 设置日志
-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}")
-        
-        # 初始化模型为空
-        self.caption_model = None
-        self.caption_processor = None
-        self.clip_model = None
-        self.clip_processor = None
-        self.sd_pipeline = None
-        self.controlnet_pipeline = None
-        self.controlnet = None
-        
-        # 模型配置 - 针对Spaces环境优化
-        self.model_config = {
-            "caption_model": "Salesforce/blip-image-captioning-base",
-            "clip_model": "openai/clip-vit-base-patch32",
-            "sd_model": "runwayml/stable-diffusion-v1-5",  # 使用更稳定的v1.5
-            "controlnet_model": "lllyasviel/sd-controlnet-openpose"
-        }
-        
-        # 创建缓存目录
-        self.cache_dir = "/tmp/models"
-        os.makedirs(self.cache_dir, exist_ok=True)
-        logger.info(f"模型缓存目录: {self.cache_dir}")
-        
-        # 加载统计
-        self.load_times = {}
-        self.last_used = {}
-        
-        # 预热标志
-        self.models_warmed = False
+def caption_image(self, image: Image.Image) -> str:
+    """对图像生成描述"""
+    if self.caption_model is None or self.caption_processor is None:
+        self.load_caption_model()
     
-    def load_caption_model(self):
-        """加载图像描述模型"""
-        if self.caption_model is None:
-            start_time = time.time()
-            logger.info("正在加载BLIP图像描述模型...")
-            
-            try:
-                self.caption_processor = BlipProcessor.from_pretrained(
-                    self.model_config["caption_model"],
-                    cache_dir=self.cache_dir
-                )
-                
-                self.caption_model = BlipForConditionalGeneration.from_pretrained(
-                    self.model_config["caption_model"],
-                    cache_dir=self.cache_dir,
-                    torch_dtype=torch.float16 if self.
+    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)