Update models/image_models.py

models/image_models.py

@@ -78,7 +78,7 @@ class ImageModelManager:
             self.logger.info(f"Loading advanced image model: {self.advanced_model_name}")
             self.advanced_processor = Blip2Processor.from_pretrained(self.advanced_model_name)
             self.advanced_model = Blip2ForConditionalGeneration.from_pretrained(
-                self.advanced_model_name, torch_dtype=torch.float16)
+                self.advanced_model_name, torch_dtype=torch.float32)
                 
             self.initialized["advanced"] = True
             self.logger.info("Advanced image model initialized successfully")
@@ -211,6 +211,7 @@ class ImageModelManager:
         # Generate caption
         with torch.no_grad():
             if model_type == "advanced":
+                pixel_values = inputs.pixel_values.to(torch.float32)
                 generated_ids = model.generate(
                     pixel_values=inputs.pixel_values,
                     max_new_tokens=50,  # Using max_new_tokens instead of max_length