app.py

import gradio as gr
import os
import cv2
import numpy as np
import shutil
import subprocess
import time
from SinglePhoto import FaceSwapper
import argparse


# Replace the old line:
# swapper = FaceSwapper()

# With:
swapper = FaceSwapper(det_size=(1280, 1280), use_enhancer=False)  # set True if you want enhancer always on
# ------------------------------------------------------------------
# Global swapper (GPU if available)
# ------------------------------------------------------------------
print("Initializing FaceSwapper (this may take a moment on first run)...")
swapper = FaceSwapper(det_size=(640, 640), ctx_id=0)
print("FaceSwapper ready.")


# ------------------------------------------------------------------
# Helpers
# ------------------------------------------------------------------
def add_audio_to_video(original_video_path, video_no_audio_path, output_path):
    """Mux original audio onto the silent swapped video using ffmpeg."""
    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(errors="ignore")


def safe_rmtree(path):
    try:
        if os.path.exists(path):
            shutil.rmtree(path)
    except Exception:
        pass


def ensure_dirs(*paths):
    for p in paths:
        os.makedirs(p, exist_ok=True)


def _resolve_video_path(video):
    """Turn whatever Gradio gives us into a real file path."""
    if video is None:
        raise ValueError("No video provided")
    if isinstance(video, str) and os.path.exists(video):
        return video
    if hasattr(video, "name") and os.path.exists(getattr(video, "name", "")):
        return video.name
    if isinstance(video, dict) and "name" in video and os.path.exists(video["name"]):
        return video["name"]
    if isinstance(video, (list, tuple)) and len(video) > 0:
        return _resolve_video_path(video[0])
    raise ValueError(f"Could not resolve video path from type {type(video)}: {video}")


# ------------------------------------------------------------------
# Photo functions
# ------------------------------------------------------------------
def swap_single_photo(src_img, src_idx, dst_img, dst_idx, progress=gr.Progress(track_tqdm=True)):
    log = ""
    start = time.time()
    try:
        progress(0, desc="Preparing")
        src_path = "workdir/SinglePhoto/data_src.jpg"
        dst_path = "workdir/SinglePhoto/data_dst.jpg"
        out_path = "workdir/SinglePhoto/output_swapped.jpg"
        ensure_dirs(os.path.dirname(src_path), os.path.dirname(out_path))

        cv2.imwrite(src_path, cv2.cvtColor(src_img, cv2.COLOR_RGB2BGR))
        cv2.imwrite(dst_path, cv2.cvtColor(dst_img, cv2.COLOR_RGB2BGR))
        log += "Saved source & destination\n"

        progress(0.4, desc="Swapping on GPU")
        result = swapper.swap_faces(src_path, int(src_idx), dst_path, int(dst_idx))
        cv2.imwrite(out_path, result)
        log += f"Saved result → {out_path}\n"

        for p in (src_path, dst_path):
            if os.path.exists(p):
                os.remove(p)

        progress(1, desc="Done")
        log += f"Elapsed: {time.time()-start:.2f}s\n"
        return out_path, log
    except Exception as e:
        log += f"ERROR: {e}\n"
        return None, log


def swap_single_src_multi_dst(src_img, dst_imgs, dst_indices, progress=gr.Progress(track_tqdm=True)):
    log = ""
    results = []
    base = "workdir/SingleSrcMultiDst"
    src_dir = f"{base}/src"
    dst_dir = f"{base}/dst"
    out_dir = f"{base}/output"
    ensure_dirs(src_dir, dst_dir, out_dir)

    try:
        if isinstance(src_img, tuple):
            src_img = src_img[0]
        src_path = os.path.join(src_dir, "data_src.jpg")
        cv2.imwrite(src_path, cv2.cvtColor(src_img, cv2.COLOR_RGB2BGR))
        log += "Saved source image\n"

        if isinstance(dst_indices, str):
            idx_list = [int(x.strip()) for x in dst_indices.split(",") if x.strip().isdigit()]
        else:
            idx_list = [int(x) for x in dst_indices]

        total = max(1, len(dst_imgs))
        for j, dst_img in enumerate(dst_imgs):
            if isinstance(dst_img, tuple):
                dst_img = dst_img[0]
            if dst_img is None:
                results.append(None)
                continue
            dst_path = os.path.join(dst_dir, f"data_dst_{j}.jpg")
            out_path = os.path.join(out_dir, f"output_swapped_{j}.jpg")
            cv2.imwrite(dst_path, cv2.cvtColor(dst_img, cv2.COLOR_RGB2BGR))
            try:
                dst_idx = idx_list[j] if j < len(idx_list) else 1
                result = swapper.swap_faces(src_path, 1, dst_path, dst_idx)
                cv2.imwrite(out_path, result)
                results.append(out_path)
                log += f"OK: src → dst[{j}] (face {dst_idx})\n"
            except Exception as e:
                results.append(None)
                log += f"FAIL dst[{j}]: {e}\n"
            progress((j + 1) / total, desc=f"Dst {j+1}/{total}")
        return results, log
    except Exception as e:
        log += f"FATAL: {e}\n"
        return results, log


def swap_multi_src_single_dst(src_imgs, dst_img, dst_idx, progress=gr.Progress(track_tqdm=True)):
    log = ""
    results = []
    base = "workdir/MultiSrcSingleDst"
    src_dir = f"{base}/src"
    dst_dir = f"{base}/dst"
    out_dir = f"{base}/output"
    ensure_dirs(src_dir, dst_dir, out_dir)

    try:
        if isinstance(dst_img, tuple):
            dst_img = dst_img[0]
        dst_path = os.path.join(dst_dir, "data_dst.jpg")
        cv2.imwrite(dst_path, cv2.cvtColor(dst_img, cv2.COLOR_RGB2BGR))
        log += "Saved destination\n"

        total = max(1, len(src_imgs))
        for i, src_img in enumerate(src_imgs):
            if isinstance(src_img, tuple):
                src_img = src_img[0]
            if src_img is None:
                results.append(None)
                continue
            src_path = os.path.join(src_dir, f"data_src_{i}.jpg")
            out_path = os.path.join(out_dir, f"output_swapped_{i}.jpg")
            cv2.imwrite(src_path, cv2.cvtColor(src_img, cv2.COLOR_RGB2BGR))
            try:
                result = swapper.swap_faces(src_path, 1, dst_path, int(dst_idx))
                cv2.imwrite(out_path, result)
                results.append(out_path)
                log += f"OK: src[{i}] → dst\n"
            except Exception as e:
                results.append(None)
                log += f"FAIL src[{i}]: {e}\n"
            progress((i + 1) / total, desc=f"Src {i+1}/{total}")
        return results, log
    except Exception as e:
        log += f"FATAL: {e}\n"
        return results, log


def swap_multi_src_multi_dst(src_imgs, dst_imgs, dst_indices, progress=gr.Progress(track_tqdm=True)):
    log = ""
    results = []
    base = "workdir/MultiSrcMultiDst"
    src_dir = f"{base}/src"
    dst_dir = f"{base}/dst"
    out_dir = f"{base}/output"
    ensure_dirs(src_dir, dst_dir, out_dir)

    try:
        if isinstance(dst_indices, str):
            idx_list = [int(x.strip()) for x in dst_indices.split(",") if x.strip().isdigit()]
        else:
            idx_list = [int(x) for x in dst_indices]

        total = max(1, len(src_imgs) * len(dst_imgs))
        count = 0
        for i, src_img in enumerate(src_imgs):
            if isinstance(src_img, tuple):
                src_img = src_img[0]
            if src_img is None:
                continue
            src_path = os.path.join(src_dir, f"data_src_{i}.jpg")
            cv2.imwrite(src_path, cv2.cvtColor(src_img, cv2.COLOR_RGB2BGR))

            for j, dst_img in enumerate(dst_imgs):
                if isinstance(dst_img, tuple):
                    dst_img = dst_img[0]
                if dst_img is None:
                    continue
                dst_path = os.path.join(dst_dir, f"data_dst_{j}.jpg")
                out_path = os.path.join(out_dir, f"output_swapped_{i}_{j}.jpg")
                cv2.imwrite(dst_path, cv2.cvtColor(dst_img, cv2.COLOR_RGB2BGR))
                try:
                    dst_idx = idx_list[j] if j < len(idx_list) else 1
                    result = swapper.swap_faces(src_path, 1, dst_path, dst_idx)
                    cv2.imwrite(out_path, result)
                    results.append(out_path)
                    log += f"OK src[{i}]→dst[{j}]\n"
                except Exception as e:
                    results.append(None)
                    log += f"FAIL {i}/{j}: {e}\n"
                count += 1
                progress(count / total, desc=f"{count}/{total}")
        return results, log
    except Exception as e:
        log += f"FATAL: {e}\n"
        return results, log


def swap_faces_custom(src_imgs, dst_img, mapping_str, progress=gr.Progress(track_tqdm=True)):
    log = ""
    start = time.time()
    base = "workdir/CustomSwap"
    src_dir = f"{base}/src"
    temp_dir = f"{base}/temp"
    dst_path = f"{base}/data_dst.jpg"
    out_path = f"{base}/output_swapped.jpg"
    ensure_dirs(src_dir, temp_dir)

    try:
        cv2.imwrite(dst_path, cv2.cvtColor(dst_img, cv2.COLOR_RGB2BGR))
        log += "Saved destination\n"

        src_paths = []
        for i, src_img in enumerate(src_imgs):
            if isinstance(src_img, tuple):
                src_img = src_img[0]
            if src_img is None:
                continue
            p = os.path.join(src_dir, f"data_src_{i+1}.jpg")
            cv2.imwrite(p, cv2.cvtColor(src_img, cv2.COLOR_RGB2BGR))
            src_paths.append(p)
            log += f"Saved source {i+1}\n"

        try:
            mapping = [int(x.strip()) for x in mapping_str.split(",") if x.strip().isdigit()]
        except Exception as e:
            return None, f"Bad mapping: {e}"

        temp_dst = os.path.join(temp_dir, "temp_dst.jpg")
        shutil.copy(dst_path, temp_dst)

        for face_idx, src_idx in enumerate(mapping, start=1):
            if src_idx < 1 or src_idx > len(src_paths):
                log += f"Skip invalid src index {src_idx} for face {face_idx}\n"
                continue
            try:
                swapped = swapper.swap_faces(src_paths[src_idx-1], 1, temp_dst, face_idx)
                cv2.imwrite(temp_dst, swapped)
                log += f"Swapped face {face_idx} ← source {src_idx}\n"
            except Exception as e:
                log += f"Failed face {face_idx}: {e}\n"

        shutil.copy(temp_dst, out_path)
        safe_rmtree(temp_dir)
        log += f"Elapsed: {time.time()-start:.2f}s\n"
        return out_path, log
    except Exception as e:
        log += f"FATAL: {e}\n"
        return None, log


# ------------------------------------------------------------------
# Video functions – robust Gradio path handling
# ------------------------------------------------------------------
def swap_video(
    src_img,
    src_idx,
    video,
    dst_idx,
    enhance=True,                    # NEW: enable face enhancement
    delete_frames_dir=True,
    add_audio=True,
    copy_to_drive=False,
    progress=gr.Progress(track_tqdm=True)
):
    """
    High-quality video face swap optimized for A100.
    """
    log = ""
    start_time = time.time()

    # Paths
    base_dir = "VideoSwapping"
    src_path = os.path.join(base_dir, "data_src.jpg")
    dst_video_path = os.path.join(base_dir, "data_dst.mp4")
    frames_dir = os.path.join(base_dir, "video_frames")
    swapped_dir = os.path.join(base_dir, "swapped_frames")
    output_video_path = os.path.join(base_dir, "output_tmp_output_video.mp4")
    final_output_path = os.path.join(base_dir, "output_with_audio.mp4")

    os.makedirs(base_dir, exist_ok=True)
    os.makedirs(frames_dir, exist_ok=True)
    os.makedirs(swapped_dir, exist_ok=True)

    try:
        # ---------- 1. Save source image ----------
        progress(0.02, desc="Saving source image...")
        if isinstance(src_img, tuple):
            src_img = src_img[0]
        src_img_bgr = cv2.cvtColor(src_img, cv2.COLOR_RGB2BGR)
        cv2.imwrite(src_path, src_img_bgr, [int(cv2.IMWRITE_JPEG_QUALITY), 95])
        log += f"Saved source image → {src_path}\n"

        # ---------- 2. Prepare target video ----------
        progress(0.05, desc="Preparing target video...")
        if isinstance(video, str) and os.path.exists(video):
            shutil.copy(video, dst_video_path)
            log += f"Copied video → {dst_video_path}\n"
        else:
            # Gradio sometimes passes a temporary path
            dst_video_path = video
            log += f"Using video path: {dst_video_path}\n"

        # ---------- 3. Extract frames (high quality JPG) ----------
        progress(0.08, desc="Extracting frames...")
        from VideoSwapping import extract_frames, frames_to_video

        frame_paths = extract_frames(dst_video_path, frames_dir, quality=95)
        total_frames = len(frame_paths)
        log += f"Extracted {total_frames} frames\n"
        progress(0.15, desc=f"Extracted {total_frames} frames")

        if total_frames == 0:
            raise ValueError("No frames extracted from video")

        # ---------- 4. Face swap loop ----------
        swapped_files = set(os.listdir(swapped_dir))
        start_loop_time = time.time()

        for idx, frame_path in enumerate(frame_paths):
            swapped_name = f"swapped_{idx:05d}.jpg"
            out_path = os.path.join(swapped_dir, swapped_name)

            # Skip already processed frames (resume support)
            if swapped_name in swapped_files and os.path.exists(out_path):
                log += f"Frame {idx+1}: already swapped, skipping\n"
            else:
                try:
                    try:
                        swapped = swapper.swap_faces(
                            source_path=src_path,
                            source_face_idx=int(src_idx),
                            target_path=frame_path,
                            target_face_idx=int(dst_idx),
                            enhance=enhance
                        )
                    except ValueError as ve:
                        # Fallback to face index 1 if requested index is missing
                        if int(dst_idx) != 1 and "Target image contains" in str(ve):
                            swapped = swapper.swap_faces(
                                source_path=src_path,
                                source_face_idx=int(src_idx),
                                target_path=frame_path,
                                target_face_idx=1,
                                enhance=enhance
                            )
                            log += f"Frame {idx+1}: dst_idx {dst_idx} not found → used 1\n"
                        else:
                            raise ve

                    cv2.imwrite(out_path, swapped, [int(cv2.IMWRITE_JPEG_QUALITY), 95])
                    log += f"Frame {idx+1}/{total_frames} swapped\n"

                except Exception as e:
                    # Keep original frame on failure
                    shutil.copy(frame_path, out_path)
                    log += f"Frame {idx+1} FAILED: {e} → kept original\n"

            # Progress + ETA
            elapsed = time.time() - start_loop_time
            avg_time = elapsed / (idx + 1)
            remaining = avg_time * (total_frames - (idx + 1))
            mins, secs = divmod(int(remaining), 60)
            progress(
                0.15 + 0.65 * (idx + 1) / total_frames,
                desc=f"Swapping {idx+1}/{total_frames} | ETA {mins:02d}:{secs:02d}"
            )

        # ---------- 5. Rebuild video ----------
        progress(0.82, desc="Building video...")
        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, use_ffmpeg=True, crf=17)
        log += f"Video written → {output_video_path}\n"

        # Optional: copy to Google Drive (Colab-style)
        if copy_to_drive:
            drive_path = "/content/drive/MyDrive/" + os.path.basename(output_video_path)
            try:
                shutil.copy(output_video_path, drive_path)
                log += f"Copied to Google Drive: {drive_path}\n"
            except Exception as e:
                log += f"Google Drive copy failed: {e}\n"

        # ---------- 6. Mux audio ----------
        progress(0.90, desc="Muxing audio...")
        if add_audio:
            ok, audio_log = add_audio_to_video(dst_video_path, output_video_path, final_output_path)
            if ok:
                log += f"Audio added → {final_output_path}\n"
            else:
                log += f"Audio mux failed: {audio_log}\n"
                final_output_path = output_video_path
        else:
            final_output_path = output_video_path
            log += "Audio skipped (user request)\n"

        # ---------- 7. Cleanup ----------
        progress(0.97, desc="Cleaning up...")
        try:
            if os.path.exists(src_path):
                os.remove(src_path)
            if os.path.exists(dst_video_path) and dst_video_path != video:
                os.remove(dst_video_path)

            if delete_frames_dir and os.path.exists(frames_dir):
                shutil.rmtree(frames_dir)
                log += "Deleted frames directory\n"
            else:
                log += "Kept frames directory\n"

            if os.path.exists(swapped_dir):
                shutil.rmtree(swapped_dir)
                log += "Deleted swapped frames\n"
        except Exception as e:
            log += f"Cleanup warning: {e}\n"

        # ---------- Done ----------
        progress(1.0, desc="Done")
        elapsed = time.time() - start_time
        log += f"\nTotal time: {elapsed:.1f} seconds ({elapsed/60:.1f} min)\n"
        log += f"Enhancement: {'ON' if enhance else 'OFF'}\n"

        return final_output_path, log

    except Exception as e:
        log += f"\nFATAL ERROR: {e}\n"
        progress(1.0, desc="Error")
        elapsed = time.time() - start_time
        log += f"Elapsed: {elapsed:.1f}s\n"
        return None, log 

def swap_video_all_faces(src_img, video, num_faces_to_swap,
                         delete_frames_dir=True, add_audio=True,
                         progress=gr.Progress()):
    """Swap the same source face onto the first N faces of every frame."""
    log = ""
    start = time.time()
    base = "workdir/VideoAllFaces"
    src_path = f"{base}/data_src.jpg"
    dst_video = f"{base}/data_dst.mp4"
    frames_dir = f"{base}/video_frames"
    swapped_dir = f"{base}/swapped_frames"
    temp_dir = f"{base}/temp"
    tmp_video = f"{base}/tmp_no_audio.mp4"
    final_video = f"{base}/output_with_audio.mp4"

    ensure_dirs(base, frames_dir, swapped_dir, temp_dir)

    try:
        cv2.imwrite(src_path, cv2.cvtColor(src_img, cv2.COLOR_RGB2BGR))
        log += "Saved source\n"

        video_path = _resolve_video_path(video)
        shutil.copy(video_path, dst_video)
        log += f"Copied video → {dst_video}\n"

        from VideoSwapping import extract_frames, frames_to_video
        frame_paths = extract_frames(dst_video, frames_dir)
        total = len(frame_paths)
        log += f"Extracted {total} frames\n"
        progress(0.1, desc=f"{total} frames")

        already = set(os.listdir(swapped_dir))
        loop_start = time.time()
        n_faces = max(1, int(num_faces_to_swap))

        for idx, frame_path in enumerate(frame_paths):
            name = f"swapped_{idx:05d}.jpg"
            out_p = os.path.join(swapped_dir, name)
            if name in already and os.path.exists(out_p):
                pass
            else:
                temp_frame = os.path.join(temp_dir, "t.jpg")
                shutil.copy(frame_path, temp_frame)
                for face_i in range(1, n_faces + 1):
                    try:
                        swapped = swapper.swap_faces(src_path, 1, temp_frame, face_i)
                        cv2.imwrite(temp_frame, swapped)
                    except Exception as e:
                        log += f"Frame {idx} face {face_i}: {e}\n"
                        break
                shutil.copy(temp_frame, out_p)
                if os.path.exists(temp_frame):
                    os.remove(temp_frame)

            elapsed = time.time() - loop_start
            avg = elapsed / (idx + 1)
            remain = avg * (total - idx - 1)
            m, s = divmod(int(remain), 60)
            progress(0.1 + 0.75 * (idx + 1) / total,
                     desc=f"{idx+1}/{total} | ETA {m:02d}:{s:02d}")

        safe_rmtree(temp_dir)

        cap = cv2.VideoCapture(dst_video)
        fps = cap.get(cv2.CAP_PROP_FPS) or 25.0
        cap.release()
        frames_to_video(swapped_dir, tmp_video, fps)

        if add_audio:
            ok, err = add_audio_to_video(dst_video, tmp_video, final_video)
            result_path = final_video if ok else tmp_video
            if not ok:
                log += f"Audio fail: {err}\n"
        else:
            result_path = tmp_video

        if delete_frames_dir:
            safe_rmtree(frames_dir)
        safe_rmtree(swapped_dir)

        log += f"Total: {time.time()-start:.1f}s\n"
        progress(1, desc="Done")
        return result_path, log

    except Exception as e:
        log += f"FATAL: {e}\n"
        return None, log


def swap_video_custom_mapping(src_imgs, video, mapping_str,
                              delete_frames_dir=True, add_audio=True,
                              progress=gr.Progress()):
    log = ""
    start = time.time()
    base = "workdir/CustomVideo"
    src_dir = f"{base}/src"
    frames_dir = f"{base}/frames"
    swapped_dir = f"{base}/swapped"
    temp_dir = f"{base}/temp"
    dst_video = f"{base}/data_dst.mp4"
    tmp_video = f"{base}/tmp_no_audio.mp4"
    final_video = f"{base}/output_with_audio.mp4"

    ensure_dirs(src_dir, frames_dir, swapped_dir, temp_dir)

    try:
        # save sources
        src_paths = []
        for i, img in enumerate(src_imgs):
            if isinstance(img, tuple):
                img = img[0]
            if img is None:
                continue
            p = os.path.join(src_dir, f"src_{i+1}.jpg")
            cv2.imwrite(p, cv2.cvtColor(img, cv2.COLOR_RGB2BGR))
            src_paths.append(p)
        log += f"Saved {len(src_paths)} source faces\n"

        try:
            mapping = [int(x.strip()) for x in mapping_str.split(",") if x.strip().isdigit()]
        except Exception as e:
            return None, f"Bad mapping: {e}"

        video_path = _resolve_video_path(video)
        shutil.copy(video_path, dst_video)
        log += f"Copied video → {dst_video}\n"

        from VideoSwapping import extract_frames, frames_to_video
        frame_paths = extract_frames(dst_video, frames_dir)
        total = len(frame_paths)
        log += f"Extracted {total} frames\n"
        progress(0.08, desc=f"{total} frames")

        already = set(os.listdir(swapped_dir))
        loop_start = time.time()
        temp_frame = os.path.join(temp_dir, "t.jpg")

        for idx, frame_path in enumerate(frame_paths):
            name = f"swapped_{idx:05d}.jpg"
            out_p = os.path.join(swapped_dir, name)
            if name in already and os.path.exists(out_p):
                pass
            else:
                shutil.copy(frame_path, temp_frame)
                for face_idx, src_idx in enumerate(mapping, start=1):
                    if src_idx < 1 or src_idx > len(src_paths):
                        continue
                    try:
                        swapped = swapper.swap_faces(
                            src_paths[src_idx-1], 1, temp_frame, face_idx
                        )
                        cv2.imwrite(temp_frame, swapped)
                    except Exception as e:
                        log += f"F{idx} face{face_idx}: {e}\n"
                shutil.copy(temp_frame, out_p)

            elapsed = time.time() - loop_start
            avg = elapsed / (idx + 1)
            remain = avg * (total - idx - 1)
            m, s = divmod(int(remain), 60)
            progress(0.08 + 0.75 * (idx + 1) / total,
                     desc=f"{idx+1}/{total} | ETA {m:02d}:{s:02d}")

        safe_rmtree(temp_dir)

        cap = cv2.VideoCapture(dst_video)
        fps = cap.get(cv2.CAP_PROP_FPS) or 25.0
        cap.release()
        frames_to_video(swapped_dir, tmp_video, fps)

        if add_audio:
            ok, err = add_audio_to_video(dst_video, tmp_video, final_video)
            result_path = final_video if ok else tmp_video
            if not ok:
                log += f"Audio fail: {err}\n"
        else:
            result_path = tmp_video

        if delete_frames_dir:
            safe_rmtree(frames_dir)
        safe_rmtree(swapped_dir)

        log += f"Total: {time.time()-start:.1f}s\n"
        progress(1, desc="Done")
        return result_path, log

    except Exception as e:
        log += f"FATAL: {e}\n"
        return None, log


def swap_single_src_multi_video(src_img, dst_videos, dst_indices,
                                delete_frames_dir=True, add_audio=True,
                                progress=gr.Progress(track_tqdm=True)):
    log = ""
    results = []
    start = time.time()
    base = "workdir/SingleSrcMultiVideo"
    ensure_dirs(base)

    try:
        if isinstance(dst_indices, str):
            idx_list = [int(x.strip()) for x in dst_indices.split(",") if x.strip().isdigit()]
        else:
            idx_list = [int(x) for x in dst_indices]

        src_path = os.path.join(base, "data_src.jpg")
        cv2.imwrite(src_path, cv2.cvtColor(src_img, cv2.COLOR_RGB2BGR))
        log += "Saved source\n"

        from VideoSwapping import extract_frames, frames_to_video

        # Gradio File(file_count="multiple") can return list of paths or list of objects
        if not isinstance(dst_videos, (list, tuple)):
            dst_videos = [dst_videos]

        n_videos = len(dst_videos)

        for i, video in enumerate(dst_videos):
            dst_idx = idx_list[i] if i < len(idx_list) else 1
            v_base = os.path.join(base, f"v{i}")
            frames_dir = os.path.join(v_base, "frames")
            swapped_dir = os.path.join(v_base, "swapped")
            dst_v = os.path.join(v_base, "dst.mp4")
            tmp_v = os.path.join(v_base, "tmp.mp4")
            final_v = os.path.join(v_base, "out.mp4")
            ensure_dirs(frames_dir, swapped_dir)

            try:
                video_path = _resolve_video_path(video)
                shutil.copy(video_path, dst_v)
            except Exception as e:
                log += f"Video {i}: cannot resolve path: {e}\n"
                results.append(None)
                continue

            frame_paths = extract_frames(dst_v, frames_dir)
            total = len(frame_paths)
            log += f"Video {i}: {total} frames\n"
            progress(i / max(1, n_videos), desc=f"Video {i+1}/{n_videos}")

            for idx, fp in enumerate(frame_paths):
                out_p = os.path.join(swapped_dir, f"swapped_{idx:05d}.jpg")
                try:
                    swapped = swapper.swap_faces(src_path, 1, fp, dst_idx)
                    cv2.imwrite(out_p, swapped)
                except Exception:
                    shutil.copy(fp, out_p)

            cap = cv2.VideoCapture(dst_v)
            fps = cap.get(cv2.CAP_PROP_FPS) or 25.0
            cap.release()
            frames_to_video(swapped_dir, tmp_v, fps)

            if add_audio:
                ok, _ = add_audio_to_video(dst_v, tmp_v, final_v)
                results.append(final_v if ok else tmp_v)
            else:
                results.append(tmp_v)

            if delete_frames_dir:
                safe_rmtree(frames_dir)
            safe_rmtree(swapped_dir)

        log += f"All videos done in {time.time()-start:.1f}s\n"
        return results, log

    except Exception as e:
        log += f"FATAL: {e}\n"
        return results, log


# ------------------------------------------------------------------
# Gradio UI
# ------------------------------------------------------------------
welcome = """
# Face Swapping Suite (GPU-ready)
All-in-one face swapping for photos **and videos**.  
Optimised for rented GPUs (CUDA). First run downloads models (~300 MB).
"""

with gr.Blocks(title="FaceSwap GPU") as demo:
    gr.Markdown(welcome)

    with gr.Tab("Single Photo"):
        gr.Interface(
            fn=swap_single_photo,
            inputs=[
                gr.Image(label="Source Image", type="numpy"),
                gr.Number(value=1, label="Source Face Index (1-based)"),
                gr.Image(label="Destination Image", type="numpy"),
                gr.Number(value=1, label="Destination Face Index (1-based)"),
            ],
            outputs=[
                gr.Image(label="Result"),
                gr.Textbox(label="Log", lines=6, interactive=False)
            ],
            api_name="single_photo"
        )

    with gr.Tab("Single Src → Multi Dst"):
        gr.Interface(
            fn=swap_single_src_multi_dst,
            inputs=[
                gr.Image(label="Source Image", type="numpy"),
                gr.Gallery(label="Destination Images", columns=3, type="numpy"),
                gr.Textbox(label="Dst face indices (comma-sep, e.g. 1,1,2)", value="1"),
            ],
            outputs=[
                gr.Gallery(label="Results"),
                gr.Textbox(label="Log", lines=6, interactive=False)
            ],
            api_name="single_src_multi_dst"
        )

    with gr.Tab("Multi Src → Single Dst"):
        gr.Interface(
            fn=swap_multi_src_single_dst,
            inputs=[
                gr.Gallery(label="Source Images", columns=3, type="numpy"),
                gr.Image(label="Destination Image", type="numpy"),
                gr.Number(value=1, label="Destination Face Index"),
            ],
            outputs=[
                gr.Gallery(label="Results"),
                gr.Textbox(label="Log", lines=6, interactive=False)
            ],
            api_name="multi_src_single_dst"
        )

    with gr.Tab("Multi Src → Multi Dst"):
        gr.Interface(
            fn=swap_multi_src_multi_dst,
            inputs=[
                gr.Gallery(label="Source Images", columns=3, type="numpy"),
                gr.Gallery(label="Destination Images", columns=3, type="numpy"),
                gr.Textbox(label="Dst face indices (comma-sep)", value="1"),
            ],
            outputs=[
                gr.Gallery(label="Results"),
                gr.Textbox(label="Log", lines=6, interactive=False)
            ],
            api_name="multi_src_multi_dst"
        )

    with gr.Tab("Custom Face Mapping (Photo)"):
        gr.Interface(
            fn=swap_faces_custom,
            inputs=[
                gr.Gallery(label="Source Images (order = indices)", columns=3, type="numpy"),
                gr.Image(label="Destination Image", type="numpy"),
                gr.Textbox(label="Mapping (e.g. 2,1,3 means face1←src2, face2←src1 …)", value="1"),
            ],
            outputs=[
                gr.Image(label="Result"),
                gr.Textbox(label="Log", lines=8, interactive=False)
            ],
            api_name="custom_photo"
        )

    with gr.Tab("Video Swapping (single face)"):
        gr.Interface(
            fn=swap_video,
            inputs=[
                gr.Image(label="Source Face Image", type="numpy"),
                gr.Number(value=1, label="Source Face Index"),
                gr.Video(label="Target Video"),
                gr.Number(value=1, label="Destination Face Index"),
                gr.Checkbox(label="Delete extracted frames after finish", value=True),
                gr.Checkbox(label="Add original audio", value=True),
            ],
            outputs=[
                gr.Video(label="Swapped Video"),
                gr.Textbox(label="Log", lines=10, interactive=False)
            ],
            api_name="video_single"
        )

    with gr.Tab("Video – All Faces"):
        gr.Interface(
            fn=swap_video_all_faces,
            inputs=[
                gr.Image(label="Source Face Image", type="numpy"),
                gr.Video(label="Target Video"),
                gr.Number(value=1, label="How many faces to swap per frame", precision=0),
                gr.Checkbox(label="Delete extracted frames after finish", value=True),
                gr.Checkbox(label="Add original audio", value=True),
            ],
            outputs=[
                gr.Video(label="Swapped Video"),
                gr.Textbox(label="Log", lines=10, interactive=False)
            ],
            api_name="video_all_faces"
        )

    with gr.Tab("Video – Custom Mapping"):
        gr.Interface(
            fn=swap_video_custom_mapping,
            inputs=[
                gr.Gallery(label="Source Images (order = indices)", columns=3, type="numpy"),
                gr.Video(label="Target Video"),
                gr.Textbox(label="Mapping (e.g. 2,1,3)", value="1"),
                gr.Checkbox(label="Delete extracted frames after finish", value=True),
                gr.Checkbox(label="Add original audio", value=True),
            ],
            outputs=[
                gr.Video(label="Swapped Video"),
                gr.Textbox(label="Log", lines=10, interactive=False)
            ],
            api_name="video_custom"
        )

    with gr.Tab("Single Src → Multi Video"):
        gr.Interface(
            fn=swap_single_src_multi_video,
            inputs=[
                gr.Image(label="Source Face Image", type="numpy"),
                gr.File(label="Target Videos (multi-select)", file_count="multiple", type="filepath"),
                gr.Textbox(label="Dst face indices per video (comma-sep)", value="1"),
                gr.Checkbox(label="Delete extracted frames after each video", value=True),
                gr.Checkbox(label="Add original audio", value=True),
            ],
            outputs=[
                gr.Gallery(label="Swapped Videos", type="filepath"),
                gr.Textbox(label="Log", lines=10, interactive=False)
            ],
            api_name="single_src_multi_video"
        )


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--share", action="store_true", help="Create a public Gradio link")
    parser.add_argument("--server-name", default="0.0.0.0")
    parser.add_argument("--server-port", type=int, default=7860)
    args = parser.parse_args()
    demo.queue(max_size=4).launch(
        share=args.share,
        server_name=args.server_name,
        server_port=args.server_port,
        show_error=True,
        ssr_mode=False,   # prevents the SvelteKit 405 / Content-Length errors
    )