Update app.py

app.py

@@ -172,6 +172,39 @@ def decode_base64_file(data_url: str) -> tuple[str, str, str]:
         logger.error(f"Failed to parse data URL: {e}")
         return None, None, None
 
+def format_image_for_replicate(base64_data: str) -> str:
+    """
+    将 base64 图片数据格式化为 Replicate 期望的格式
+    """
+    # 检查 base64 数据是否已经包含 data URL 前缀
+    if base64_data.startswith("data:"):
+        return base64_data
+    
+    # 如果没有前缀，添加默认的 JPEG data URL 前缀
+    # 但首先尝试检测实际的图片格式
+    try:
+        # 解码 base64 数据的前几个字节来检测格式
+        decoded_bytes = base64.b64decode(base64_data[:100])
+        
+        if decoded_bytes.startswith(b'\xff\xd8\xff'):
+            # JPEG
+            return f"data:image/jpeg;base64,{base64_data}"
+        elif decoded_bytes.startswith(b'\x89PNG\r\n\x1a\n'):
+            # PNG
+            return f"data:image/png;base64,{base64_data}"
+        elif decoded_bytes.startswith(b'GIF87a') or decoded_bytes.startswith(b'GIF89a'):
+            # GIF
+            return f"data:image/gif;base64,{base64_data}"
+        elif decoded_bytes.startswith(b'RIFF') and b'WEBP' in decoded_bytes[:20]:
+            # WebP
+            return f"data:image/webp;base64,{base64_data}"
+        else:
+            # 默认使用 JPEG
+            return f"data:image/jpeg;base64,{base64_data}"
+    except Exception as e:
+        logger.warning(f"Failed to detect image format: {e}, using JPEG as default")
+        return f"data:image/jpeg;base64,{base64_data}"
+
 def extract_content_from_message(message: Dict[str, Any]) -> tuple[str, List[str], List[Dict[str, str]]]:
     """
     从消息中提取文本内容、图片和文件
@@ -204,10 +237,16 @@ def extract_content_from_message(message: Dict[str, Any]) -> tuple[str, List[str
                         try:
                             if ";base64," in url:
                                 base64_data = url.split(";base64,")[1]
-                                images.append(base64_data)
+                                # 格式化为 Replicate 期望的格式
+                                formatted_image = format_image_for_replicate(base64_data)
+                                images.append(formatted_image)
                                 logger.info(f"Found base64 image, size: {len(base64_data)} chars")
+                            else:
+                                logger.warning(f"Image URL format not supported: {url[:100]}...")
                         except Exception as e:
                             logger.error(f"Error processing image: {e}")
+                    else:
+                        logger.warning(f"External image URLs not supported: {url}")
                 
                 elif item_type == "file" or (item_type == "image_url" and not item.get("image_url", {}).get("url", "").startswith("data:image/")):
                     # 处理文件上传
@@ -221,7 +260,8 @@ def extract_content_from_message(message: Dict[str, Any]) -> tuple[str, List[str
                             
                             if file_ext in SUPPORTED_IMAGE_EXTENSIONS and mime_type.startswith("image/"):
                                 # 图片文件
-                                images.append(file_content)
+                                formatted_image = format_image_for_replicate(file_content)
+                                images.append(formatted_image)
                                 logger.info(f"Found image file: {filename}")
                             elif file_ext in SUPPORTED_TEXT_EXTENSIONS or mime_type.startswith("text/"):
                                 # 文本文件
@@ -278,6 +318,7 @@ def transform_openai_to_replicate(openai_request: Dict[str, Any], model_override
         has_images = False
         has_files = False
         all_files = []
+        primary_image = None
         
         for message in messages:
             if message.get("role") == "system":
@@ -297,6 +338,10 @@ def transform_openai_to_replicate(openai_request: Dict[str, Any], model_override
                 
                 if image_list:
                     has_images = True
+                    # 使用最后一个用户消息中的第一张图片作为主要图片
+                    if message.get("role") == "user":
+                        primary_image = image_list[0]
+                
                 if file_list:
                     has_files = True
                     all_files.extend(file_list)
@@ -359,15 +404,10 @@ def transform_openai_to_replicate(openai_request: Dict[str, Any], model_override
         
         replicate_input["prompt"] = prompt
         
-        # 处理图片（只使用第一张图片）
-        if has_images:
-            # 找到最后一个包含图片的用户消息
-            for msg in reversed(user_messages):
-                if msg["role"] == "user" and msg["images"]:
-                    primary_image = msg["images"][0]
-                    replicate_input["image"] = f"data:image/jpeg;base64,{primary_image}"
-                    logger.info(f"Added primary image to request")
-                    break
+        # 处理图片 - 使用正确的 data URL 格式
+        if has_images and primary_image:
+            replicate_input["image"] = primary_image
+            logger.info(f"Added primary image to request: {primary_image[:100]}...")
         
         # 只在有 system_prompt 时才添加
         if system_prompt:
@@ -437,7 +477,8 @@ async def create_replicate_prediction(session: aiohttp.ClientSession, model: str
         log_data = data.copy()
         if "input" in log_data:
             if "image" in log_data["input"]:
-                log_data["input"]["image"] = f"[IMAGE_DATA_{len(log_data['input']['image'])}]"
+                image_data = log_data["input"]["image"]
+                log_data["input"]["image"] = f"[IMAGE_DATA_{len(image_data)}]"
             if "prompt" in log_data["input"] and len(log_data["input"]["prompt"]) > 1000:
                 log_data["input"]["prompt"] = log_data["input"]["prompt"][:1000] + "...[TRUNCATED]"
         logger.info(f"Request data: {json.dumps(log_data, indent=2)}")
@@ -522,7 +563,7 @@ async def root():
         "message": "Replicate API Proxy for LobeChat with Vision and File Support",
         "status": "running",
         "replicate_token_configured": bool(REPLICATE_API_TOKEN),
-        "version": "1.1.0",
+        "version": "1.1.1",
         "supported_models": list(MODEL_CONFIGS.keys()),
         "vision_support": True,
         "file_support": True,