'commit'

dsLightRag/Routes/QWenImageRoute.py

@@ -10,7 +10,7 @@ from fastapi import APIRouter, HTTPException
 from pydantic import BaseModel, Field
 
 from Config.Config import OBS_PREFIX, OBS_BUCKET, OBS_SERVER
-# 导入QwenImageGenerator类
+from QWenImage.QWenImageEditKit import QwenImageEditor
 from QWenImage.QwenImageKit import QwenImageGenerator
 from Util.ObsUtil import ObsUploader
 
@@ -20,8 +20,9 @@ router = APIRouter(prefix="/api/qwenImage", tags=["千问生图"])
 # 配置日志
 logger = logging.getLogger(__name__)
 
-# 初始化图片生成器
+# 初始化图片生成器和编辑器
 image_generator = QwenImageGenerator()
+image_editor = QwenImageEditor()
 
 
 class GenerateImageRequest(BaseModel):
@@ -31,6 +32,23 @@ class GenerateImageRequest(BaseModel):
     size: str = Field(default="1328*1328", description="图片尺寸")
     api_key: Optional[str] = Field(default=None, description="自定义API密钥")
 
+# 添加图片编辑请求模型
+class EditImageRequest(BaseModel):
+    """编辑图片请求模型"""
+    prompt: str = Field(..., description="编辑提示词")
+    size: str = Field(default="1328*1328", description="图片尺寸")
+    api_key: Optional[str] = Field(default=None, description="自定义API密钥")
+    # 支持URL或Base64两种格式，二选一
+    image_url: Optional[str] = Field(default=None, description="原始图片URL")
+    image_base64: Optional[str] = Field(default=None, description="原始图片Base64编码")
+
+    @root_validator(pre=True)
+    def check_image_source(cls, values):
+        """验证图片来源，确保提供了URL或Base64中的一种"""
+        if not values.get("image_url") and not values.get("image_base64"):
+            raise ValueError("必须提供image_url或image_base64中的一种")
+        return values
+
 
 @router.post("/generate")
 async def generate_image(request: GenerateImageRequest):
@@ -116,6 +134,100 @@ async def generate_image(request: GenerateImageRequest):
         )
 
 
+@router.post("/edit")
+async def edit_image(request: EditImageRequest):
+    """编辑图片API接口
+    
+    Args:
+        request: 包含编辑图片参数的请求对象
+        
+    Returns:
+        dict: 包含编辑结果的字典
+    """
+    try:
+        logger.info(f"接收到图片编辑请求: image_url={request.image_url[:50]}..., prompt={request.prompt[:50]}...")
+        
+        # 如果提供了自定义API密钥，创建新的编辑器实例
+        editor = QwenImageEditor(api_key=request.api_key) if request.api_key else image_editor
+
+        # 调用图片编辑API
+        result = editor.edit_image(
+            image_url=request.image_url,
+            prompt=request.prompt,
+            negative_prompt=request.negative_prompt
+        )
+        
+        # 处理结果
+        if result['success']:
+            logger.info(f"图片编辑成功")
+            
+            # 上传到OBS
+            obs_urls = []
+            uploader = ObsUploader()
+            
+            try:
+                # 下载编辑后的图片
+                import requests
+                response_img = requests.get(result['image_url'], timeout=10)
+                response_img.raise_for_status()
+                bytes_data = response_img.content
+
+                # 生成UUID文件名并上传
+                jpg_file_name = f"{str(uuid.uuid4())}.jpg"
+                object_key = f"{OBS_PREFIX}/QWen3Image/edited/{jpg_file_name}"
+                success, upload_result = uploader.upload_base64_image(object_key, bytes_data)
+
+                if success:
+                    obs_url = f"https://{OBS_BUCKET}.{OBS_SERVER}/{object_key}"
+                    obs_urls.append(obs_url)
+                    logger.info(f"编辑图片上传OBS成功: {obs_url}")
+                else:
+                    logger.error(f"编辑图片上传OBS失败: {upload_result}")
+            except Exception as e:
+                logger.error(f"处理编辑图片时出错: {str(e)}")
+                raise HTTPException(
+                    status_code=500,
+                    detail={
+                        "code": 500,
+                        "message": "图片处理失败",
+                        "error_detail": str(e)
+                    }
+                )
+
+            # 构造返回响应
+            response = {
+                "code": 200,
+                "message": "图片编辑成功",
+                "data": {
+                    "original_image": request.image_url,
+                    "edited_image": result['image_url'],
+                    "obs_url": obs_urls[0] if obs_urls else None
+                }
+            }
+            return response
+        else:
+            logger.error(f"图片编辑失败: {result['error_msg']}")
+            raise HTTPException(
+                status_code=500,
+                detail={
+                    "code": 500,
+                    "message": "图片编辑失败",
+                    "error_detail": result['error_msg']
+                }
+            )
+            
+    except Exception as e:
+        logger.exception(f"处理图片编辑请求时发生异常: {str(e)}")
+        raise HTTPException(
+            status_code=500,
+            detail={
+                "code": 500,
+                "message": "处理请求时发生异常",
+                "error_detail": str(e)
+            }
+        )
+
+
 @router.get("/config")
 async def get_image_config():
     """获取图片生成配置信息