Update models/model_manager.py

models/model_manager.py

@@ -0,0 +1,130 @@
+# models/model_manager.py
+from transformers import BlipProcessor, BlipForConditionalGeneration, pipeline
+from diffusers import StableDiffusionPipeline, ControlNetModel, StableDiffusionControlNetPipeline
+import torch
+
+class ModelManager:
+    def __init__(self):
+        self.device = "cuda" if torch.cuda.is_available() else "cpu"
+        self.init_models()
+    
+    def init_models(self):
+        print("正在加载模型...")
+        
+        # 修复1: 使用兼容的 BLIP 模型
+        print("加载图像理解模型...")
+        self.blip_processor = BlipProcessor.from_pretrained(
+            "Salesforce/blip-image-captioning-base",
+            # 添加兼容性参数
+        )
+        self.blip_model = BlipForConditionalGeneration.from_pretrained(
+            "Salesforce/blip-image-captioning-base",
+        ).to(self.device)
+        
+        # 修复2: 文本生成模型 - 添加错误处理
+        print("加载文本生成模型...")
+        try:
+            self.text_generator = pipeline(
+                "text-generation", 
+                model="microsoft/DialoGPT-medium",
+                device=0 if self.device == "cuda" else -1
+            )
+        except Exception as e:
+            print(f"DialoGPT 加载失败，使用备选模型: {e}")
+            self.text_generator = pipeline(
+                "text-generation",
+                model="gpt2",
+                device=0 if self.device == "cuda" else -1
+            )
+        
+        # 修复3: Stable Diffusion 模型 - 添加内存优化
+        print("加载 Stable Diffusion 模型...")
+        self.sd_pipeline = StableDiffusionPipeline.from_pretrained(
+            "runwayml/stable-diffusion-v1-5",
+            torch_dtype=torch.float16 if self.device == "cuda" else torch.float32,
+            use_safetensors=True,
+            variant="fp16" if self.device == "cuda" else None
+        )
+        
+        # 内存优化
+        if self.device == "cuda":
+            self.sd_pipeline.enable_model_cpu_offload()
+            self.sd_pipeline.enable_xformers_memory_efficient_attention()
+        else:
+            self.sd_pipeline = self.sd_pipeline.to(self.device)
+        
+        # 修复4: ControlNet 模型 - 添加错误处理
+        print("加载 ControlNet 模型...")
+        try:
+            self.controlnet = ControlNetModel.from_pretrained(
+                "lllyasviel/sd-controlnet-openpose",
+                torch_dtype=torch.float16 if self.device == "cuda" else torch.float32,
+                use_safetensors=True
+            )
+            
+            self.controlnet_pipeline = StableDiffusionControlNetPipeline.from_pretrained(
+                "runwayml/stable-diffusion-v1-5",
+                controlnet=self.controlnet,
+                torch_dtype=torch.float16 if self.device == "cuda" else torch.float32,
+                use_safetensors=True,
+                variant="fp16" if self.device == "cuda" else None
+            )
+            
+            if self.device == "cuda":
+                self.controlnet_pipeline.enable_model_cpu_offload()
+                self.controlnet_pipeline.enable_xformers_memory_efficient_attention()
+            else:
+                self.controlnet_pipeline = self.controlnet_pipeline.to(self.device)
+                
+        except Exception as e:
+            print(f"ControlNet 加载失败: {e}")
+            self.controlnet = None
+            self.controlnet_pipeline = None
+        
+        print("所有模型加载完成！")
+    
+    def generate_caption(self, image):
+        """生成图像描述"""
+        inputs = self.blip_processor(image, return_tensors="pt").to(self.device)
+        with torch.no_grad():
+            out = self.blip_model.generate(**inputs, max_length=50)
+        return self.blip_processor.decode(out[0], skip_special_tokens=True)
+    
+    def generate_text(self, prompt, max_length=100):
+        """生成文本"""
+        try:
+            result = self.text_generator(
+                prompt, 
+                max_length=max_length, 
+                num_return_sequences=1,
+                temperature=0.7,
+                do_sample=True,
+                pad_token_id=self.text_generator.tokenizer.eos_token_id
+            )
+            return result[0]['generated_text']
+        except Exception as e:
+            print(f"文本生成错误: {e}")
+            return f"生成失败: {str(e)}"
+    
+    def generate_image(self, prompt, negative_prompt="", num_inference_steps=20):
+        """生成图像"""
+        try:
+            with torch.autocast(self.device):
+                image = self.sd_pipeline(
+                    prompt=prompt,
+                    negative_prompt=negative_prompt,
+                    num_inference_steps=num_inference_steps,
+                    guidance_scale=7.5
+                ).images[0]
+            return image
+        except Exception as e:
+            print(f"图像生成错误: {e}")
+            return None
+    
+    def cleanup(self):
+        """清理 GPU ���存"""
+        if hasattr(self, 'sd_pipeline'):
+            del self.sd_pipeline
+        if hasattr(self, 'controlnet_pipeline'):
+            del self.controlnet_pipeline
+        torch.cuda.empty_cache() if torch.cuda.is_available() else None