Spaces:
Paused
Paused
| import os | |
| import time | |
| from pathlib import Path | |
| import cv2 | |
| import gradio as gr | |
| from optimized_swapper import FaceSwapEngine | |
| ENGINE = FaceSwapEngine() | |
| def _as_index(v, default=1): | |
| try: | |
| return max(1, int(v)) | |
| except Exception: | |
| return default | |
| def _file_to_path(item): | |
| if item is None: | |
| return None | |
| if isinstance(item, (str, Path)): | |
| return str(item) | |
| if hasattr(item, "name"): | |
| return item.name | |
| if isinstance(item, dict): | |
| for key in ("path", "name"): | |
| if key in item and item[key]: | |
| return str(item[key]) | |
| if isinstance(item, (list, tuple)) and item: | |
| first = item[0] | |
| if isinstance(first, (str, Path)): | |
| return str(first) | |
| if hasattr(first, "name"): | |
| return first.name | |
| return None | |
| def _normalize_video_input(video_value): | |
| if video_value is None: | |
| return None | |
| if isinstance(video_value, (str, Path)): | |
| return str(video_value) | |
| if isinstance(video_value, dict): | |
| if video_value.get("video"): | |
| return str(video_value["video"]) | |
| if video_value.get("path"): | |
| return str(video_value["path"]) | |
| if isinstance(video_value, (list, tuple)) and video_value: | |
| first = video_value[0] | |
| if isinstance(first, (str, Path)): | |
| return str(first) | |
| if isinstance(first, dict): | |
| return str(first.get("video") or first.get("path") or "") or None | |
| return None | |
| def _eta_text(start_time, progress_value): | |
| elapsed = max(0.001, time.time() - start_time) | |
| if progress_value <= 0: | |
| return f"Elapsed: {elapsed:.1f}s | ETA: estimating..." | |
| remaining = elapsed * (1.0 - progress_value) / progress_value | |
| if remaining >= 60: | |
| return f"Elapsed: {elapsed:.1f}s | ETA: {remaining / 60:.1f} min" | |
| return f"Elapsed: {elapsed:.1f}s | ETA: {remaining:.1f}s" | |
| def _status(progress_value, message, start_time): | |
| pct = int(max(0, min(100, round(progress_value * 100)))) | |
| return f"{message}\nProgress: {pct}%\n{_eta_text(start_time, progress_value)}" | |
| def swap_photo(source, source_idx, target, target_idx, det_size, progress=gr.Progress()): | |
| start = time.time() | |
| progress(0.05, desc="Preparing image swap") | |
| if source is None or target is None: | |
| raise gr.Error("Upload both source and target images.") | |
| progress(0.35, desc="Running face detection") | |
| result = ENGINE.swap_image( | |
| source, | |
| _as_index(source_idx), | |
| target, | |
| _as_index(target_idx), | |
| det_size=int(det_size), | |
| ) | |
| progress(1.0, desc="Done") | |
| return result, _status(1.0, "Photo swap complete.", start) | |
| def swap_video(source, source_idx, video_path, target_idx, det_size, detection_interval, jpeg_quality, audio, progress=gr.Progress()): | |
| start = time.time() | |
| progress(0.02, desc="Validating inputs") | |
| if source is None: | |
| raise gr.Error("Upload a source image.") | |
| video_path = _normalize_video_input(video_path) | |
| if not video_path: | |
| raise gr.Error("Upload a target video.") | |
| progress(0.10, desc="Preparing source face") | |
| progress(0.20, desc="Starting video processing") | |
| result = ENGINE.swap_video( | |
| source=source, | |
| source_idx=_as_index(source_idx), | |
| video_path=video_path, | |
| target_idx=_as_index(target_idx), | |
| det_size=int(det_size), | |
| detection_interval=max(1, int(detection_interval)), | |
| jpeg_quality=int(jpeg_quality), | |
| preserve_audio=bool(audio), | |
| ) | |
| progress(1.0, desc="Video complete") | |
| return result, _status(1.0, "Video swap complete.", start) | |
| def swap_multi_source_single(source_files, target, target_idx, det_size, progress=gr.Progress()): | |
| start = time.time() | |
| if not source_files or target is None: | |
| raise gr.Error("Upload source images and a target image.") | |
| results = [] | |
| total = len(source_files) | |
| for i, item in enumerate(source_files, start=1): | |
| progress(((i - 1) / max(total, 1)) * 0.9 + 0.05, desc=f"Processing source {i}/{total}") | |
| path = _file_to_path(item) | |
| if not path: | |
| continue | |
| img = cv2.imread(path) | |
| if img is None: | |
| results.append(f"Error: could not read source image {path}") | |
| continue | |
| try: | |
| result = ENGINE.swap_image( | |
| img, | |
| 1, | |
| target, | |
| _as_index(target_idx), | |
| int(det_size), | |
| ) | |
| results.append(result) | |
| except Exception as e: | |
| results.append(f"Error: {e}") | |
| if not results: | |
| raise gr.Error("No readable source images were uploaded.") | |
| progress(1.0, desc="Batch complete") | |
| return results, _status(1.0, f"Processed {len(results)} output(s).", start) | |
| def swap_multi_source_multi(source_files, target_files, target_indices, det_size, progress=gr.Progress()): | |
| start = time.time() | |
| if not source_files or not target_files: | |
| raise gr.Error("Upload source and target images.") | |
| indices = [x.strip() for x in str(target_indices).split(",") if x.strip()] | |
| results = [] | |
| targets = [] | |
| progress(0.05, desc="Loading target images") | |
| for item in target_files: | |
| path = _file_to_path(item) | |
| if not path: | |
| continue | |
| img = cv2.imread(path) | |
| if img is not None: | |
| targets.append(img) | |
| if not targets: | |
| raise gr.Error("No readable target images were uploaded.") | |
| total_jobs = max(1, len(source_files) * len(targets)) | |
| done = 0 | |
| for src_item in source_files: | |
| src_path = _file_to_path(src_item) | |
| if not src_path: | |
| continue | |
| src = cv2.imread(src_path) | |
| if src is None: | |
| results.append(f"Error: could not read source image {src_path}") | |
| continue | |
| try: | |
| ENGINE.prepare_source(src, 1, int(det_size)) | |
| except Exception as e: | |
| results.append(f"Error preparing source {src_path}: {e}") | |
| continue | |
| for j, dst in enumerate(targets): | |
| idx = _as_index(indices[j] if j < len(indices) else 1) | |
| done += 1 | |
| progress((done / total_jobs) * 0.95, desc=f"Processing pair {done}/{total_jobs}") | |
| try: | |
| results.append(ENGINE.swap_prepared_source(dst, idx, int(det_size))) | |
| except Exception as e: | |
| results.append(f"Error: {e}") | |
| if not results: | |
| raise gr.Error("No output images were generated.") | |
| progress(1.0, desc="Batch complete") | |
| return results, _status(1.0, f"Processed {len(results)} output(s).", start) | |
| with gr.Blocks(title="Fast Face Swap") as demo: | |
| gr.Markdown( | |
| "# Fast Face Swapping Suite\n" | |
| "CUDA/ONNX Runtime optimized photo and video face swapping with live progress status." | |
| ) | |
| with gr.Tab("Single Photo"): | |
| with gr.Row(): | |
| with gr.Column(): | |
| src = gr.Image(type="numpy", label="Source Image") | |
| src_idx = gr.Number(value=1, precision=0, label="Source Face Index") | |
| with gr.Column(): | |
| dst = gr.Image(type="numpy", label="Target Image") | |
| dst_idx = gr.Number(value=1, precision=0, label="Target Face Index") | |
| det_single = gr.Dropdown([256, 320, 384, 512], value=320, label="Detector Size") | |
| btn = gr.Button("Swap", variant="primary") | |
| out = gr.Image(type="numpy", label="Result") | |
| photo_status = gr.Textbox(label="Status / ETA", lines=3, interactive=False) | |
| btn.click( | |
| fn=swap_photo, | |
| inputs=[src, src_idx, dst, dst_idx, det_single], | |
| outputs=[out, photo_status], | |
| show_progress="full", | |
| ) | |
| with gr.Tab("Fast Video"): | |
| vsrc = gr.Image(type="numpy", label="Source Image") | |
| vsrc_idx = gr.Number(value=1, precision=0, visible=False, label="Source Face Index") | |
| with gr.Row(): | |
| vid = gr.Video(label="Target Video") | |
| vdst_idx = gr.Number(value=1, precision=0, label="Target Face Index") | |
| with gr.Row(): | |
| vdet = gr.Dropdown([256, 320, 384, 512], value=320, label="Detector Size") | |
| interval = gr.Slider( | |
| 1, | |
| 5, | |
| value=1, | |
| step=1, | |
| label="Face detection interval (1 = best tracking accuracy)", | |
| ) | |
| quality = gr.Slider(75, 98, value=92, step=1, label="JPEG fallback quality") | |
| audio = gr.Checkbox(value=True, label="Preserve original audio") | |
| vbtn = gr.Button("Swap Video", variant="primary") | |
| vout = gr.Video(label="Output Video") | |
| video_status = gr.Textbox(label="Status / ETA", lines=3, interactive=False) | |
| vbtn.click( | |
| fn=swap_video, | |
| inputs=[vsrc, vsrc_idx, vid, vdst_idx, vdet, interval, quality, audio], | |
| outputs=[vout, video_status], | |
| show_progress="full", | |
| ) | |
| with gr.Tab("Multi Source -> Single Target"): | |
| ms = gr.File(file_count="multiple", file_types=["image"], type="filepath", label="Source Images") | |
| md = gr.Image(type="numpy", label="Target Image") | |
| mi = gr.Number(value=1, precision=0, label="Target Face Index") | |
| det_multi_single = gr.Dropdown([256, 320, 384, 512], value=320, label="Detector Size") | |
| mb = gr.Button("Process", variant="primary") | |
| mo = gr.Gallery(label="Results", columns=3) | |
| multi_single_status = gr.Textbox(label="Status / ETA", lines=3, interactive=False) | |
| mb.click( | |
| fn=swap_multi_source_single, | |
| inputs=[ms, md, mi, det_multi_single], | |
| outputs=[mo, multi_single_status], | |
| show_progress="full", | |
| ) | |
| with gr.Tab("Multi Source -> Multi Target"): | |
| mss = gr.File(file_count="multiple", file_types=["image"], type="filepath", label="Source Images") | |
| mdd = gr.File(file_count="multiple", file_types=["image"], type="filepath", label="Target Images") | |
| mids = gr.Textbox(value="1", label="Target face indices, comma-separated") | |
| det_multi_multi = gr.Dropdown([256, 320, 384, 512], value=320, label="Detector Size") | |
| mdb = gr.Button("Process", variant="primary") | |
| mdo = gr.Gallery(label="Results", columns=3) | |
| multi_multi_status = gr.Textbox(label="Status / ETA", lines=3, interactive=False) | |
| mdb.click( | |
| fn=swap_multi_source_multi, | |
| inputs=[mss, mdd, mids, det_multi_multi], | |
| outputs=[mdo, multi_multi_status], | |
| show_progress="full", | |
| ) | |
| if __name__ == "__main__": | |
| demo.queue(default_concurrency_limit=1, max_size=8).launch( | |
| server_name="0.0.0.0", | |
| server_port=int(os.getenv("PORT", "7860")), | |
| show_error=True, | |
| ) | |