sync server add presigned url params (#1033)

Co-authored-by: yihuiwen <yihuiwen@sensetime.com>

Author: black-eleven

lightx2v/server/api/tasks/image.py

@@ -30,11 +30,7 @@ async def _wait_task_and_stream_result(task_id: str, timeout_seconds: int, poll_
         if status == TaskStatus.COMPLETED.value:
             result_png = task_manager.get_task_result_png(task_id)
             if result_png:
-                return Response(
-                    content=result_png,
-                    media_type="image/png",
-                    headers={"Content-Disposition": 'inline; filename="result.png"'},
-                )
+                return result_png
             raise HTTPException(status_code=500, detail=f"Task completed but no in-memory image found: {task_id}")
 
         if status == TaskStatus.FAILED.value:
@@ -50,6 +46,39 @@ async def _wait_task_and_stream_result(task_id: str, timeout_seconds: int, poll_
         await asyncio.sleep(poll_interval_seconds)
 
 
+def _build_png_response(result_png: bytes) -> Response:
+    return Response(
+        content=result_png,
+        media_type="image/png",
+        headers={"Content-Disposition": 'inline; filename="result.png"'},
+    )
+
+
+async def _upload_sync_result_if_needed(message: ImageTaskRequest, result_png: bytes):
+    presigned_url = (getattr(message, "presigned_url", "") or "").strip()
+    if not presigned_url:
+        return None
+
+    services = get_services()
+    assert services.file_service is not None, "File service is not initialized"
+
+    try:
+        await services.file_service.upload_to_presigned_url(
+            presigned_url=presigned_url,
+            file_content=result_png,
+            content_type="image/png",
+        )
+    except ValueError as e:
+        raise HTTPException(status_code=502, detail=f"Failed to upload sync result to presigned URL: {str(e)}")
+
+    return {
+        "task_id": message.task_id,
+        "task_status": "completed",
+        "uploaded_to_presigned_url": True,
+        "presigned_url": presigned_url,
+    }
+
+
 async def _watch_client_disconnect(request: Request, task_id: str, poll_interval_seconds: float = 0.2) -> bool:
     while True:
         if await request.is_disconnected():
@@ -105,6 +134,9 @@ async def create_image_task_sync(
         if hasattr(message, "image_mask_path") and message.image_mask_path and message.image_mask_path.startswith("http"):
             if not await validate_url_async(message.image_mask_path):
                 raise HTTPException(status_code=400, detail=f"Image mask URL is not accessible: {message.image_mask_path}")
+        if hasattr(message, "presigned_url") and message.presigned_url:
+            if not message.presigned_url.startswith(("http://", "https://")):
+                raise HTTPException(status_code=400, detail=f"Invalid presigned_url: {message.presigned_url}")
 
         message.prefer_memory_result = True
         task_id = task_manager.create_task(message)
@@ -124,7 +156,11 @@ async def create_image_task_sync(
                 await asyncio.gather(wait_task, return_exceptions=True)
             raise HTTPException(status_code=499, detail=f"Client disconnected, task {task_id} cancelled")
 
-        return wait_task.result()
+        result_png = wait_task.result()
+        upload_result = await _upload_sync_result_if_needed(message, result_png)
+        if upload_result is not None:
+            return upload_result
+        return _build_png_response(result_png)
 
     except asyncio.CancelledError:
         if task_id:

lightx2v/server/schema.py

@@ -35,6 +35,7 @@ class BaseTaskRequest(DisaggOverrideRequest):
     image_path: str = Field("", description="Base64 encoded image or URL")
     image_mask_path: str = Field("", description="Mask image path (supports URL, base64, or local path)")
     save_result_path: str = Field("", description="Save result path (optional, defaults to task_id, suffix auto-detected)")
+    presigned_url: str = Field("", description="Optional presigned URL for uploading final sync result")
     infer_steps: int = Field(5, description="Inference steps")
     seed: int = Field(default_factory=generate_random_seed, description="Random seed (auto-generated if not set)")
     target_shape: list[int] = Field([], description="Return video or image shape")

lightx2v/server/services/file_service.py

@@ -130,6 +130,67 @@ async def download_image(self, image_url: str) -> Path:
     async def download_audio(self, audio_url: str) -> Path:
         return await self.download_media(audio_url, "audio")
 
+    async def upload_to_presigned_url(
+        self,
+        presigned_url: str,
+        file_content: bytes,
+        content_type: str = "application/octet-stream",
+        max_retries: Optional[int] = None,
+    ) -> None:
+        if max_retries is None:
+            max_retries = self.max_retries
+
+        parsed_url = urlparse(presigned_url)
+        if not parsed_url.scheme or not parsed_url.netloc:
+            raise ValueError(f"Invalid presigned URL format: {presigned_url}")
+
+        last_exception = None
+        retry_delay = self.retry_delay
+
+        for attempt in range(max_retries):
+            try:
+                client = await self._get_http_client()
+                response = await client.put(
+                    presigned_url,
+                    content=file_content,
+                    headers={"Content-Type": content_type},
+                )
+                if 200 <= response.status_code < 300:
+                    logger.info(f"Successfully uploaded result to presigned URL, status: {response.status_code}")
+                    return
+                if response.status_code >= 500:
+                    logger.warning(f"Presigned upload server error {response.status_code}, attempt {attempt + 1}/{max_retries}")
+                    last_exception = httpx.HTTPStatusError(
+                        f"Server returned {response.status_code}",
+                        request=response.request,
+                        response=response,
+                    )
+                else:
+                    raise httpx.HTTPStatusError(
+                        f"Client error {response.status_code}",
+                        request=response.request,
+                        response=response,
+                    )
+            except (httpx.ConnectError, httpx.TimeoutException, httpx.NetworkError) as e:
+                logger.warning(f"Connection error uploading to presigned URL, attempt {attempt + 1}/{max_retries}: {str(e)}")
+                last_exception = e
+            except httpx.HTTPStatusError as e:
+                if e.response and e.response.status_code < 500:
+                    raise ValueError(f"Failed to upload to presigned URL: HTTP {e.response.status_code}")
+                last_exception = e
+            except Exception as e:
+                logger.error(f"Unexpected error uploading to presigned URL: {str(e)}")
+                last_exception = e
+
+            if attempt < max_retries - 1:
+                await asyncio.sleep(retry_delay)
+                retry_delay = min(retry_delay * 2, self.max_retry_delay)
+
+        error_msg = f"All {max_retries} upload attempts failed for presigned URL"
+        if last_exception:
+            error_msg += f": {str(last_exception)}"
+        raise ValueError(error_msg)
+
     def save_uploaded_file(self, file_content: bytes, filename: str) -> Path:
         file_extension = Path(filename).suffix
         unique_filename = f"{uuid.uuid4()}{file_extension}"

lightx2v/server/services/generation/image.py

@@ -45,6 +45,7 @@ async def generate_with_stop_event(self, message: Any, stop_event) -> Optional[A
             task_data["seed"] = message.seed
             prefer_memory_result = bool(getattr(message, "prefer_memory_result", False))
             task_data.pop("prefer_memory_result", None)
+            task_data.pop("presigned_url", None)
             task_data["return_result_tensor"] = prefer_memory_result
 
             result = await self.inference_service.submit_task_async(task_data)