Add comprehensive FastAPI endpoints for all face swap operations

fastapi_app.py

@@ -1,7 +1,9 @@
 import logging
 import os
 import shutil
+import subprocess
 import tempfile
+import zipfile
 from io import BytesIO
 
 import cv2
@@ -35,14 +37,44 @@ def decode_image(upload: UploadFile) -> np.ndarray:
     return bgr
 
 
-def swap_video_file(src_bgr: np.ndarray, src_idx: int, video_path: str, dst_idx: int) -> tuple[str, str]:
+def add_audio_to_video(original_video_path: str, video_no_audio_path: str, output_path: str) -> tuple[bool, str]:
+    """Add audio from original video to swapped video."""
+    cmd = [
+        "ffmpeg",
+        "-y",
+        "-i", video_no_audio_path,
+        "-i", original_video_path,
+        "-c:v", "copy",
+        "-c:a", "aac",
+        "-map", "0:v:0",
+        "-map", "1:a:0?",
+        "-shortest",
+        output_path
+    ]
+    try:
+        subprocess.run(cmd, check=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
+        return True, ""
+    except subprocess.CalledProcessError as e:
+        return False, e.stderr.decode() if e.stderr else "Unknown error"
+
+
+def swap_video_file(
+    src_bgr: np.ndarray,
+    src_idx: int,
+    video_path: str,
+    dst_idx: int,
+    add_audio: bool = True,
+    workdir: str | None = None
+) -> tuple[str, str]:
     """Run face swap on a video file. Returns (output_video_path, workdir)."""
-    workdir = tempfile.mkdtemp(prefix="swap_video_")
+    if workdir is None:
+        workdir = tempfile.mkdtemp(prefix="swap_video_")
     src_path = os.path.join(workdir, "data_src.jpg")
     dst_video_path = os.path.join(workdir, "data_dst.mp4")
     frames_dir = os.path.join(workdir, "video_frames")
     swapped_dir = os.path.join(workdir, "swapped_frames")
-    output_video_path = os.path.join(workdir, "output_swapped_video.mp4")
+    output_video_path = os.path.join(workdir, "output_tmp_output_video.mp4")
+    final_output_path = os.path.join(workdir, "output_with_audio.mp4")
 
     os.makedirs(frames_dir, exist_ok=True)
     os.makedirs(swapped_dir, exist_ok=True)
@@ -67,7 +99,6 @@ def swap_video_file(src_bgr: np.ndarray, src_idx: int, video_path: str, dst_idx:
             cv2.imwrite(out_path, swapped)
         except Exception as exc:
             logger.warning("Failed to swap frame %s: %s", idx, exc)
-            # Copy original frame so the video remains playable
             cv2.imwrite(out_path, cv2.imread(frame_path))
 
     cap = cv2.VideoCapture(dst_video_path)
@@ -75,6 +106,15 @@ def swap_video_file(src_bgr: np.ndarray, src_idx: int, video_path: str, dst_idx:
     cap.release()
     frames_to_video(swapped_dir, output_video_path, fps)
     logger.info("Created swapped video at %s", output_video_path)
+
+    if add_audio:
+        ok, audio_log = add_audio_to_video(dst_video_path, output_video_path, final_output_path)
+        if ok:
+            logger.info("Added audio to %s", final_output_path)
+            return final_output_path, workdir
+        else:
+            logger.warning("Audio muxing failed: %s", audio_log)
+            return output_video_path, workdir
     return output_video_path, workdir
 
 
@@ -91,6 +131,7 @@ async def swap_photo(
     destination_face_idx: int = Form(1, ge=1),
     _: None = Depends(verify_api_key),
 ) -> StreamingResponse:
+    """Swap a single face from source image to destination image."""
     src_bgr = decode_image(source_image)
     dst_bgr = decode_image(destination_image)
 
@@ -106,7 +147,7 @@ async def swap_photo(
                 swapped = face_swapper.swap_faces(src_path, source_face_idx, dst_path, 1)
                 logger.info("Fallback to destination_face_idx 1 after missing face")
             else:
-                raise
+                raise HTTPException(status_code=400, detail=str(err))
 
     ok, buffer = cv2.imencode(".jpg", swapped)
     if not ok:
@@ -114,6 +155,151 @@ async def swap_photo(
     return StreamingResponse(BytesIO(buffer.tobytes()), media_type="image/jpeg")
 
 
+@app.post("/swap/photo/single-src-multi-dst")
+async def swap_single_src_multi_dst(
+    source_image: UploadFile = File(..., description="Source face image"),
+    destination_images: list[UploadFile] = File(..., description="Destination images"),
+    destination_indices: str = Form(..., description="Comma-separated destination face indices (e.g., '1,1,2')"),
+    _: None = Depends(verify_api_key),
+) -> StreamingResponse:
+    """Swap a single source face onto multiple destination images."""
+    src_bgr = decode_image(source_image)
+    dst_indices_list = [int(idx.strip()) for idx in destination_indices.split(",") if idx.strip().isdigit()]
+
+    with tempfile.TemporaryDirectory(prefix="single_src_multi_dst_") as workdir:
+        src_path = os.path.join(workdir, "src.jpg")
+        cv2.imwrite(src_path, src_bgr)
+
+        zip_buffer = BytesIO()
+        with zipfile.ZipFile(zip_buffer, "w", zipfile.ZIP_DEFLATED) as zip_file:
+            for j, dst_file in enumerate(destination_images):
+                dst_bgr = decode_image(dst_file)
+                dst_path = os.path.join(workdir, f"dst_{j}.jpg")
+                output_path = os.path.join(workdir, f"output_{j}.jpg")
+                cv2.imwrite(dst_path, dst_bgr)
+
+                try:
+                    dst_idx = dst_indices_list[j] if j < len(dst_indices_list) else 1
+                    swapped = face_swapper.swap_faces(src_path, 1, dst_path, dst_idx)
+                    cv2.imwrite(output_path, swapped)
+                    zip_file.write(output_path, f"swapped_{j}.jpg")
+                except Exception as exc:
+                    logger.warning("Failed to swap destination %s: %s", j, exc)
+
+        zip_buffer.seek(0)
+        return StreamingResponse(zip_buffer, media_type="application/zip", headers={"Content-Disposition": "attachment; filename=swapped_images.zip"})
+
+
+@app.post("/swap/photo/multi-src-single-dst")
+async def swap_multi_src_single_dst(
+    source_images: list[UploadFile] = File(..., description="Source face images"),
+    destination_image: UploadFile = File(..., description="Destination image"),
+    destination_face_idx: int = Form(1, ge=1),
+    _: None = Depends(verify_api_key),
+) -> StreamingResponse:
+    """Swap multiple source faces onto a single destination image."""
+    dst_bgr = decode_image(destination_image)
+
+    with tempfile.TemporaryDirectory(prefix="multi_src_single_dst_") as workdir:
+        dst_path = os.path.join(workdir, "dst.jpg")
+        cv2.imwrite(dst_path, dst_bgr)
+
+        zip_buffer = BytesIO()
+        with zipfile.ZipFile(zip_buffer, "w", zipfile.ZIP_DEFLATED) as zip_file:
+            for i, src_file in enumerate(source_images):
+                src_bgr = decode_image(src_file)
+                src_path = os.path.join(workdir, f"src_{i}.jpg")
+                output_path = os.path.join(workdir, f"output_{i}.jpg")
+                cv2.imwrite(src_path, src_bgr)
+
+                try:
+                    swapped = face_swapper.swap_faces(src_path, 1, dst_path, destination_face_idx)
+                    cv2.imwrite(output_path, swapped)
+                    zip_file.write(output_path, f"swapped_{i}.jpg")
+                except Exception as exc:
+                    logger.warning("Failed to swap source %s: %s", i, exc)
+
+        zip_buffer.seek(0)
+        return StreamingResponse(zip_buffer, media_type="application/zip", headers={"Content-Disposition": "attachment; filename=swapped_images.zip"})
+
+
+@app.post("/swap/photo/multi-src-multi-dst")
+async def swap_multi_src_multi_dst(
+    source_images: list[UploadFile] = File(..., description="Source face images"),
+    destination_images: list[UploadFile] = File(..., description="Destination images"),
+    destination_indices: str = Form(..., description="Comma-separated destination face indices (e.g., '1,1,2')"),
+    _: None = Depends(verify_api_key),
+) -> StreamingResponse:
+    """Swap multiple source faces onto multiple destination images."""
+    dst_indices_list = [int(idx.strip()) for idx in destination_indices.split(",") if idx.strip().isdigit()]
+
+    with tempfile.TemporaryDirectory(prefix="multi_src_multi_dst_") as workdir:
+        zip_buffer = BytesIO()
+        with zipfile.ZipFile(zip_buffer, "w", zipfile.ZIP_DEFLATED) as zip_file:
+            for i, src_file in enumerate(source_images):
+                src_bgr = decode_image(src_file)
+                src_path = os.path.join(workdir, f"src_{i}.jpg")
+                cv2.imwrite(src_path, src_bgr)
+
+                for j, dst_file in enumerate(destination_images):
+                    dst_bgr = decode_image(dst_file)
+                    dst_path = os.path.join(workdir, f"dst_{j}.jpg")
+                    output_path = os.path.join(workdir, f"output_{i}_{j}.jpg")
+
+                    try:
+                        dst_idx = dst_indices_list[j] if j < len(dst_indices_list) else 1
+                        swapped = face_swapper.swap_faces(src_path, 1, dst_path, dst_idx)
+                        cv2.imwrite(output_path, swapped)
+                        zip_file.write(output_path, f"swapped_src{i}_dst{j}.jpg")
+                    except Exception as exc:
+                        logger.warning("Failed to swap src %s with dst %s: %s", i, j, exc)
+
+        zip_buffer.seek(0)
+        return StreamingResponse(zip_buffer, media_type="application/zip", headers={"Content-Disposition": "attachment; filename=swapped_images.zip"})
+
+
+@app.post("/swap/photo/custom-mapping")
+async def swap_faces_custom(
+    source_images: list[UploadFile] = File(..., description="Source face images"),
+    destination_image: UploadFile = File(..., description="Destination image"),
+    mapping: str = Form(..., description="Comma-separated mapping (e.g., '2,1,3' means face1->src2, face2->src1, face3->src3)"),
+    _: None = Depends(verify_api_key),
+) -> StreamingResponse:
+    """Swap faces with custom mapping: map destination faces to specific source images."""
+    dst_bgr = decode_image(destination_image)
+    mapping_list = [int(x.strip()) for x in mapping.split(",") if x.strip().isdigit()]
+
+    with tempfile.TemporaryDirectory(prefix="custom_swap_") as workdir:
+        dst_path = os.path.join(workdir, "dst.jpg")
+        temp_path = os.path.join(workdir, "temp.jpg")
+        output_path = os.path.join(workdir, "output.jpg")
+        cv2.imwrite(dst_path, dst_bgr)
+        shutil.copy(dst_path, temp_path)
+
+        src_paths = []
+        for i, src_file in enumerate(source_images):
+            src_bgr = decode_image(src_file)
+            src_path = os.path.join(workdir, f"src_{i+1}.jpg")
+            cv2.imwrite(src_path, src_bgr)
+            src_paths.append(src_path)
+
+        for face_idx, src_idx in enumerate(mapping_list, start=1):
+            if src_idx < 1 or src_idx > len(src_paths):
+                logger.warning("Invalid source index %s for face %s", src_idx, face_idx)
+                continue
+            try:
+                swapped_img = face_swapper.swap_faces(src_paths[src_idx-1], 1, temp_path, face_idx)
+                cv2.imwrite(temp_path, swapped_img)
+            except Exception as exc:
+                logger.warning("Failed to swap face %s with source %s: %s", face_idx, src_idx, exc)
+
+        shutil.copy(temp_path, output_path)
+        ok, buffer = cv2.imencode(".jpg", cv2.imread(output_path))
+        if not ok:
+            raise HTTPException(status_code=500, detail="Failed to encode swapped image")
+        return StreamingResponse(BytesIO(buffer.tobytes()), media_type="image/jpeg")
+
+
 @app.post("/swap/video")
 async def swap_video(
     background_tasks: BackgroundTasks,
@@ -121,26 +307,245 @@ async def swap_video(
     target_video: UploadFile = File(..., description="Target video (mp4)"),
     source_face_idx: int = Form(1, ge=1),
     destination_face_idx: int = Form(1, ge=1),
+    add_audio: bool = Form(True, description="Add audio from original video"),
     _: None = Depends(verify_api_key),
 ) -> StreamingResponse:
+    """Swap a single face from source image onto all frames of target video."""
     src_bgr = decode_image(source_image)
-    with tempfile.NamedTemporaryFile(delete=False, suffix=".mp4") as tmp_video:
+    workdir = tempfile.mkdtemp(prefix="swap_video_")
+    with tempfile.NamedTemporaryFile(delete=False, suffix=".mp4", dir=workdir) as tmp_video:
         video_bytes = await target_video.read()
         tmp_video.write(video_bytes)
         video_path = tmp_video.name
 
     try:
-        output_path, workdir = swap_video_file(src_bgr, source_face_idx, video_path, destination_face_idx)
+        output_path, _ = swap_video_file(src_bgr, source_face_idx, video_path, destination_face_idx, add_audio, workdir)
     except Exception as exc:
         logger.exception("Video swap failed: %s", exc)
+        shutil.rmtree(workdir, ignore_errors=True)
+        raise HTTPException(status_code=500, detail="Video swap failed") from exc
+
+    def cleanup() -> None:
+        shutil.rmtree(workdir, ignore_errors=True)
+        logger.info("Cleaned up workdir %s", workdir)
+
+    background_tasks.add_task(cleanup)
+    return StreamingResponse(open(output_path, "rb"), media_type="video/mp4")
+
+
+@app.post("/swap/video/all-faces")
+async def swap_video_all_faces(
+    background_tasks: BackgroundTasks,
+    source_image: UploadFile = File(..., description="Source face image"),
+    target_video: UploadFile = File(..., description="Target video (mp4)"),
+    num_faces_to_swap: int = Form(1, ge=1, description="Number of faces to swap"),
+    add_audio: bool = Form(True, description="Add audio from original video"),
+    _: None = Depends(verify_api_key),
+) -> StreamingResponse:
+    """Swap all faces in video with the same source face."""
+    src_bgr = decode_image(source_image)
+    workdir = tempfile.mkdtemp(prefix="swap_video_all_")
+    src_path = os.path.join(workdir, "data_src.jpg")
+    cv2.imwrite(src_path, src_bgr)
+
+    with tempfile.NamedTemporaryFile(delete=False, suffix=".mp4", dir=workdir) as tmp_video:
+        video_bytes = await target_video.read()
+        tmp_video.write(video_bytes)
+        video_path = tmp_video.name
+
+    try:
+        dst_video_path = os.path.join(workdir, "data_dst.mp4")
+        frames_dir = os.path.join(workdir, "video_frames")
+        swapped_dir = os.path.join(workdir, "swapped_frames")
+        temp_dir = os.path.join(swapped_dir, "temp_swap")
+        output_video_path = os.path.join(workdir, "output_tmp_output_video.mp4")
+        final_output_path = os.path.join(workdir, "output_with_audio.mp4")
+
+        os.makedirs(frames_dir, exist_ok=True)
+        os.makedirs(swapped_dir, exist_ok=True)
+        os.makedirs(temp_dir, exist_ok=True)
+
+        shutil.copy(video_path, dst_video_path)
+        frame_paths = extract_frames(dst_video_path, frames_dir)
+        logger.info("Extracted %s frames", len(frame_paths))
+
+        temp_frame_path = os.path.join(temp_dir, "temp.jpg")
+        for idx, frame_path in enumerate(frame_paths):
+            swapped_name = f"swapped_{idx:05d}.jpg"
+            out_path = os.path.join(swapped_dir, swapped_name)
+            try:
+                shutil.copy(frame_path, temp_frame_path)
+                for face_idx in range(1, num_faces_to_swap + 1):
+                    try:
+                        swapped_img = face_swapper.swap_faces(src_path, 1, temp_frame_path, face_idx)
+                        cv2.imwrite(temp_frame_path, swapped_img)
+                    except Exception as exc:
+                        logger.warning("Failed to swap face %s in frame %s: %s", face_idx, idx, exc)
+                shutil.copy(temp_frame_path, out_path)
+            except Exception as exc:
+                logger.warning("Failed to swap frame %s: %s", idx, exc)
+                cv2.imwrite(out_path, cv2.imread(frame_path))
+
+        cap = cv2.VideoCapture(dst_video_path)
+        fps = cap.get(cv2.CAP_PROP_FPS) or 25.0
+        cap.release()
+        frames_to_video(swapped_dir, output_video_path, fps)
+
+        if add_audio:
+            ok, audio_log = add_audio_to_video(dst_video_path, output_video_path, final_output_path)
+            if ok:
+                output_path = final_output_path
+            else:
+                logger.warning("Audio muxing failed: %s", audio_log)
+                output_path = output_video_path
+        else:
+            output_path = output_video_path
+
+    except Exception as exc:
+        logger.exception("Video swap all faces failed: %s", exc)
+        shutil.rmtree(workdir, ignore_errors=True)
+        raise HTTPException(status_code=500, detail="Video swap failed") from exc
+
+    def cleanup() -> None:
+        shutil.rmtree(workdir, ignore_errors=True)
+        logger.info("Cleaned up workdir %s", workdir)
+
+    background_tasks.add_task(cleanup)
+    return StreamingResponse(open(output_path, "rb"), media_type="video/mp4")
+
+
+@app.post("/swap/video/custom-mapping")
+async def swap_video_custom_mapping(
+    background_tasks: BackgroundTasks,
+    source_images: list[UploadFile] = File(..., description="Source face images"),
+    target_video: UploadFile = File(..., description="Target video (mp4)"),
+    mapping: str = Form(..., description="Comma-separated mapping (e.g., '2,1,3')"),
+    add_audio: bool = Form(True, description="Add audio from original video"),
+    _: None = Depends(verify_api_key),
+) -> StreamingResponse:
+    """Swap faces in video with custom mapping: map destination faces to specific source images."""
+    workdir = tempfile.mkdtemp(prefix="swap_video_custom_")
+    mapping_list = [int(x.strip()) for x in mapping.split(",") if x.strip().isdigit()]
+
+    src_paths = []
+    for i, src_file in enumerate(source_images):
+        src_bgr = decode_image(src_file)
+        src_path = os.path.join(workdir, f"src_{i+1}.jpg")
+        cv2.imwrite(src_path, src_bgr)
+        src_paths.append(src_path)
+
+    with tempfile.NamedTemporaryFile(delete=False, suffix=".mp4", dir=workdir) as tmp_video:
+        video_bytes = await target_video.read()
+        tmp_video.write(video_bytes)
+        video_path = tmp_video.name
+
+    try:
+        dst_video_path = os.path.join(workdir, "data_dst.mp4")
+        temp_dir = os.path.join(workdir, "temp")
+        frames_dir = os.path.join(workdir, "frames")
+        swapped_dir = os.path.join(workdir, "swapped_frames")
+        output_video_path = os.path.join(workdir, "output_tmp_output_video.mp4")
+        final_output_path = os.path.join(workdir, "output_with_audio.mp4")
+
+        os.makedirs(temp_dir, exist_ok=True)
+        os.makedirs(frames_dir, exist_ok=True)
+        os.makedirs(swapped_dir, exist_ok=True)
+
+        shutil.copy(video_path, dst_video_path)
+        frame_paths = extract_frames(dst_video_path, frames_dir)
+        logger.info("Extracted %s frames", len(frame_paths))
+
+        temp_frame_path = os.path.join(temp_dir, "temp.jpg")
+        for idx, frame_path in enumerate(frame_paths):
+            swapped_name = f"swapped_{idx:05d}.jpg"
+            out_path = os.path.join(swapped_dir, swapped_name)
+            try:
+                shutil.copy(frame_path, temp_frame_path)
+                for face_idx, src_idx in enumerate(mapping_list, start=1):
+                    if src_idx < 1 or src_idx > len(src_paths):
+                        logger.warning("Invalid source index %s for face %s in frame %s", src_idx, face_idx, idx)
+                        continue
+                    try:
+                        swapped_img = face_swapper.swap_faces(src_paths[src_idx-1], 1, temp_frame_path, face_idx)
+                        cv2.imwrite(temp_frame_path, swapped_img)
+                    except Exception as exc:
+                        logger.warning("Failed to swap face %s with source %s in frame %s: %s", face_idx, src_idx, idx, exc)
+                shutil.copy(temp_frame_path, out_path)
+            except Exception as exc:
+                logger.warning("Failed to swap frame %s: %s", idx, exc)
+                cv2.imwrite(out_path, cv2.imread(frame_path))
+
+        cap = cv2.VideoCapture(dst_video_path)
+        fps = cap.get(cv2.CAP_PROP_FPS) or 25.0
+        cap.release()
+        frames_to_video(swapped_dir, output_video_path, fps)
+
+        if add_audio:
+            ok, audio_log = add_audio_to_video(dst_video_path, output_video_path, final_output_path)
+            if ok:
+                output_path = final_output_path
+            else:
+                logger.warning("Audio muxing failed: %s", audio_log)
+                output_path = output_video_path
+        else:
+            output_path = output_video_path
+
+    except Exception as exc:
+        logger.exception("Video swap custom mapping failed: %s", exc)
+        shutil.rmtree(workdir, ignore_errors=True)
         raise HTTPException(status_code=500, detail="Video swap failed") from exc
 
     def cleanup() -> None:
         shutil.rmtree(workdir, ignore_errors=True)
-        if os.path.exists(video_path):
-            os.remove(video_path)
         logger.info("Cleaned up workdir %s", workdir)
 
     background_tasks.add_task(cleanup)
     return StreamingResponse(open(output_path, "rb"), media_type="video/mp4")
 
+
+@app.post("/swap/video/single-src-multi-video")
+async def swap_single_src_multi_video(
+    background_tasks: BackgroundTasks,
+    source_image: UploadFile = File(..., description="Source face image"),
+    target_videos: list[UploadFile] = File(..., description="Target videos (mp4)"),
+    destination_indices: str = Form(..., description="Comma-separated destination face indices (e.g., '1,2,1')"),
+    add_audio: bool = Form(True, description="Add audio from original video"),
+    _: None = Depends(verify_api_key),
+) -> StreamingResponse:
+    """Swap a single source face onto multiple videos."""
+    src_bgr = decode_image(source_image)
+    dst_indices_list = [int(idx.strip()) for idx in destination_indices.split(",") if idx.strip().isdigit()]
+    workdir = tempfile.mkdtemp(prefix="single_src_multi_video_")
+    src_path = os.path.join(workdir, "data_src.jpg")
+    cv2.imwrite(src_path, src_bgr)
+
+    video_paths = []
+    for i, video_file in enumerate(target_videos):
+        video_bytes = await video_file.read()
+        video_path = os.path.join(workdir, f"video_{i}.mp4")
+        with open(video_path, "wb") as f:
+            f.write(video_bytes)
+        video_paths.append((video_path, dst_indices_list[i] if i < len(dst_indices_list) else 1))
+
+    try:
+        zip_buffer = BytesIO()
+        with zipfile.ZipFile(zip_buffer, "w", zipfile.ZIP_DEFLATED) as zip_file:
+            for i, (video_path, dst_idx) in enumerate(video_paths):
+                try:
+                    output_path, _ = swap_video_file(src_bgr, 1, video_path, dst_idx, add_audio, workdir)
+                    zip_file.write(output_path, f"swapped_video_{i}.mp4")
+                except Exception as exc:
+                    logger.warning("Failed to swap video %s: %s", i, exc)
+
+        zip_buffer.seek(0)
+    except Exception as exc:
+        logger.exception("Single src multi video swap failed: %s", exc)
+        shutil.rmtree(workdir, ignore_errors=True)
+        raise HTTPException(status_code=500, detail="Video swap failed") from exc
+
+    def cleanup() -> None:
+        shutil.rmtree(workdir, ignore_errors=True)
+        logger.info("Cleaned up workdir %s", workdir)
+
+    background_tasks.add_task(cleanup)
+    return StreamingResponse(zip_buffer, media_type="application/zip", headers={"Content-Disposition": "attachment; filename=swapped_videos.zip"})