Spaces:
Sleeping
Sleeping
| import gradio as gr | |
| import os | |
| import cv2 | |
| import numpy as np | |
| import shutil | |
| import subprocess | |
| import time | |
| import argparse | |
| import zipfile | |
| # --- LOAD ENGINE --- | |
| try: | |
| from SinglePhoto import FaceSwapper | |
| swapper = FaceSwapper() | |
| print("FaceSwapper loaded successfully.") | |
| except ImportError: | |
| print("WARNING: 'SinglePhoto' module not found. UI will load, but swapping will fail.") | |
| class FaceSwapper: | |
| def swap_faces(self, *args, **kwargs): | |
| raise ImportError("SinglePhoto.py not found.") | |
| swapper = FaceSwapper() | |
| # --- CSS & THEME CONFIGURATION --- | |
| custom_css = """ | |
| body { background-color: #0b0f19; } | |
| .gradio-container { background-color: #0b0f19 !important; } | |
| h1.title-header { | |
| color: #00ff41; | |
| font-family: 'Courier New', Courier, monospace; | |
| text-align: center; | |
| font-weight: bold; | |
| text-shadow: 0px 0px 10px rgba(0, 255, 65, 0.5); | |
| margin-bottom: 20px; | |
| padding: 20px; | |
| border: 1px solid #00ff41; | |
| border-radius: 8px; | |
| background: #000; | |
| } | |
| .gr-button-primary { | |
| background: linear-gradient(90deg, #004d1a 0%, #00b33c 100%) !important; | |
| border: 1px solid #00ff41 !important; | |
| color: white !important; | |
| } | |
| """ | |
| shinyy_theme = gr.themes.Soft( | |
| primary_hue="green", | |
| neutral_hue="slate", | |
| font=[gr.themes.GoogleFont("Source Code Pro"), "ui-sans-serif", "system-ui", "sans-serif"], | |
| ).set( | |
| body_background_fill="#0b0f19", | |
| block_background_fill="#111827", | |
| block_border_width="1px", | |
| block_title_text_color="#00ff41", | |
| block_label_text_color="#a3a3a3", | |
| input_background_fill="#1f2937", | |
| button_primary_background_fill="#00ff41", | |
| button_primary_text_color="#000000", | |
| ) | |
| # --- UTILITY: ZIP PROJECT --- | |
| def zip_project_directory(): | |
| """Zips the current working directory for download.""" | |
| zip_filename = "project_backup.zip" | |
| # Files/Folders to ignore to keep zip clean/small | |
| ignore_folders = {'.git', '__pycache__', 'venv', 'env', '.ipynb_checkpoints'} | |
| with zipfile.ZipFile(zip_filename, 'w', zipfile.ZIP_DEFLATED) as zipf: | |
| for root, dirs, files in os.walk('.'): | |
| # Remove ignored directories from traversal | |
| dirs[:] = [d for d in dirs if d not in ignore_folders] | |
| for file in files: | |
| if file == zip_filename: continue # Don't zip the zip itself | |
| file_path = os.path.join(root, file) | |
| # Archive name is the path relative to current directory | |
| arcname = os.path.relpath(file_path, '.') | |
| zipf.write(file_path, arcname) | |
| return zip_filename | |
| # --- LOGIC FUNCTIONS --- | |
| def swap_single_photo(src_img, src_idx, dst_img, dst_idx, include_hair, progress=gr.Progress(track_tqdm=True)): | |
| log = "" | |
| start_time = time.time() | |
| try: | |
| progress(0, desc="Preparing files") | |
| if src_img is None or dst_img is None: | |
| return None, "Error: Images missing." | |
| src_path = "SinglePhoto/data_src.jpg" | |
| dst_path = "SinglePhoto/data_dst.jpg" | |
| output_path = "SinglePhoto/output_swapped.jpg" | |
| os.makedirs(os.path.dirname(src_path), exist_ok=True) | |
| os.makedirs(os.path.dirname(dst_path), exist_ok=True) | |
| os.makedirs(os.path.dirname(output_path), exist_ok=True) | |
| cv2.imwrite(src_path, cv2.cvtColor(src_img, cv2.COLOR_RGB2BGR)) | |
| cv2.imwrite(dst_path, cv2.cvtColor(dst_img, cv2.COLOR_RGB2BGR)) | |
| progress(0.5, desc="Swapping faces" + (" & hair" if include_hair else "")) | |
| # --- ENGINE CALL --- | |
| result = swapper.swap_faces(src_path, int(src_idx), dst_path, int(dst_idx), swap_hair=include_hair) | |
| cv2.imwrite(output_path, result) | |
| log += f"Swapped{' (with Hair)' if include_hair else ''} and saved.\n" | |
| progress(1, desc="Done") | |
| result_rgb = cv2.cvtColor(result, cv2.COLOR_BGR2RGB) | |
| return result_rgb, 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 = [] | |
| src_dir = "SingleSrcMultiDst/src" | |
| dst_dir = "SingleSrcMultiDst/dst" | |
| output_dir = "SingleSrcMultiDst/output" | |
| os.makedirs(src_dir, exist_ok=True) | |
| os.makedirs(dst_dir, exist_ok=True) | |
| os.makedirs(output_dir, exist_ok=True) | |
| if src_img is None: return [], "Error: Source image missing" | |
| 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)) | |
| if isinstance(dst_indices, str): | |
| dst_indices_list = [int(idx.strip()) for idx in dst_indices.split(",") if idx.strip().isdigit()] | |
| else: | |
| dst_indices_list = [int(idx) for idx in dst_indices] if dst_indices else [] | |
| if not dst_imgs: return [], "Error: No destination images." | |
| for j, dst_img in enumerate(dst_imgs): | |
| if isinstance(dst_img, tuple): dst_img = dst_img[0] | |
| dst_path = os.path.join(dst_dir, f"data_dst_{j}.jpg") | |
| output_path = os.path.join(output_dir, f"output_swapped_{j}.jpg") | |
| cv2.imwrite(dst_path, cv2.cvtColor(dst_img, cv2.COLOR_RGB2BGR)) | |
| try: | |
| dst_idx = dst_indices_list[j] if j < len(dst_indices_list) else 1 | |
| result = swapper.swap_faces(src_path, 1, dst_path, int(dst_idx)) | |
| cv2.imwrite(output_path, result) | |
| results.append(output_path) | |
| except Exception as e: | |
| results.append(None) | |
| log += f"Error dst {j}: {e}\n" | |
| progress((j+1)/len(dst_imgs), desc=f"Processing {j+1}/{len(dst_imgs)}") | |
| return results, log | |
| def swap_multi_src_single_dst(src_imgs, dst_img, dst_idx, progress=gr.Progress(track_tqdm=True)): | |
| log = "" | |
| results = [] | |
| src_dir = "MultiSrcSingleDst/src" | |
| dst_dir = "MultiSrcSingleDst/dst" | |
| output_dir = "MultiSrcSingleDst/output" | |
| os.makedirs(src_dir, exist_ok=True) | |
| os.makedirs(dst_dir, exist_ok=True) | |
| os.makedirs(output_dir, exist_ok=True) | |
| if dst_img is None: return [], "Error: Dest image missing" | |
| 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)) | |
| if not src_imgs: return [], "Error: No source images." | |
| for i, src_img in enumerate(src_imgs): | |
| if isinstance(src_img, tuple): src_img = src_img[0] | |
| src_path = os.path.join(src_dir, f"data_src_{i}.jpg") | |
| output_path = os.path.join(output_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(output_path, result) | |
| results.append(output_path) | |
| except Exception as e: | |
| results.append(None) | |
| log += f"Error src {i}: {e}\n" | |
| progress((i+1)/len(src_imgs)) | |
| return results, log | |
| def swap_multi_src_multi_dst(src_imgs, dst_imgs, dst_indices, progress=gr.Progress(track_tqdm=True)): | |
| log = "" | |
| results = [] | |
| src_dir = "MultiSrcMultiDst/src" | |
| dst_dir = "MultiSrcMultiDst/dst" | |
| output_dir = "MultiSrcMultiDst/output" | |
| os.makedirs(src_dir, exist_ok=True) | |
| os.makedirs(dst_dir, exist_ok=True) | |
| os.makedirs(output_dir, exist_ok=True) | |
| if isinstance(dst_indices, str): | |
| dst_indices_list = [int(idx.strip()) for idx in dst_indices.split(",") if idx.strip().isdigit()] | |
| else: | |
| dst_indices_list = [] | |
| 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") | |
| if isinstance(src_img, np.ndarray): cv2.imwrite(src_path, cv2.cvtColor(src_img, cv2.COLOR_RGB2BGR)) | |
| elif isinstance(src_img, str) and os.path.exists(src_img): shutil.copy(src_img, src_path) | |
| 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") | |
| output_path = os.path.join(output_dir, f"output_swapped_{i}_{j}.jpg") | |
| if isinstance(dst_img, np.ndarray): cv2.imwrite(dst_path, cv2.cvtColor(dst_img, cv2.COLOR_RGB2BGR)) | |
| elif isinstance(dst_img, str) and os.path.exists(dst_img): shutil.copy(dst_img, dst_path) | |
| try: | |
| dst_idx = dst_indices_list[j] if j < len(dst_indices_list) else 1 | |
| result = swapper.swap_faces(src_path, 1, dst_path, int(dst_idx)) | |
| cv2.imwrite(output_path, result) | |
| results.append(output_path) | |
| except Exception as e: | |
| log += f"Error {i}-{j}: {e}\n" | |
| count += 1 | |
| progress(count/total) | |
| return results, log | |
| def swap_faces_custom(src_imgs, dst_img, mapping_str, progress=gr.Progress(track_tqdm=True)): | |
| log = "" | |
| dst_path = "CustomSwap/data_dst.jpg" | |
| output_path = "CustomSwap/output_swapped.jpg" | |
| src_dir = "CustomSwap/src" | |
| temp_dir = "CustomSwap/temp" | |
| os.makedirs(src_dir, exist_ok=True) | |
| os.makedirs(temp_dir, exist_ok=True) | |
| os.makedirs(os.path.dirname(dst_path), exist_ok=True) | |
| if dst_img is None: return None, "Error: Dst missing" | |
| cv2.imwrite(dst_path, cv2.cvtColor(dst_img, cv2.COLOR_RGB2BGR)) | |
| src_paths = [] | |
| if not src_imgs: return None, "Error: Src missing" | |
| for i, src_img in enumerate(src_imgs): | |
| src_path = os.path.join(src_dir, f"data_src_{i+1}.jpg") | |
| if isinstance(src_img, tuple): src_img = src_img[0] | |
| if isinstance(src_img, np.ndarray): cv2.imwrite(src_path, cv2.cvtColor(src_img, cv2.COLOR_RGB2BGR)) | |
| elif isinstance(src_img, str): shutil.copy(src_img, src_path) | |
| src_paths.append(src_path) | |
| try: mapping = [int(x.strip()) for x in mapping_str.split(",") if x.strip().isdigit()] | |
| except: return None, "Error: Invalid map" | |
| temp_dst_path = os.path.join(temp_dir, "temp_dst.jpg") | |
| shutil.copy(dst_path, temp_dst_path) | |
| for face_idx, src_idx in enumerate(mapping, start=1): | |
| if src_idx < 1 or src_idx > len(src_paths): continue | |
| try: | |
| swapped_img = swapper.swap_faces(src_paths[src_idx-1], 1, temp_dst_path, face_idx) | |
| cv2.imwrite(temp_dst_path, swapped_img) | |
| log += f"Face {face_idx} <- Src {src_idx}\n" | |
| except Exception as e: | |
| log += f"Fail face {face_idx}: {e}\n" | |
| shutil.copy(temp_dst_path, output_path) | |
| return output_path, log | |
| # --- VIDEO FUNCTIONS --- | |
| def add_audio_to_video(original_video_path, video_no_audio_path, output_path): | |
| 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() | |
| def swap_video(src_img, src_idx, video, dst_idx, delete_frames, add_audio, copy_drive, progress=gr.Progress()): | |
| log = "" | |
| src_path = "VideoSwapping/data_src.jpg" | |
| dst_video_path = "VideoSwapping/data_dst.mp4" | |
| frames_dir = "VideoSwapping/video_frames" | |
| swapped_dir = "VideoSwapping/swapped_frames" | |
| output_video_path = "VideoSwapping/output_tmp_output_video.mp4" | |
| final_output_path = "VideoSwapping/output_with_audio.mp4" | |
| os.makedirs(os.path.dirname(src_path), exist_ok=True) | |
| os.makedirs(os.path.dirname(dst_video_path), exist_ok=True) | |
| os.makedirs(frames_dir, exist_ok=True) | |
| os.makedirs(swapped_dir, exist_ok=True) | |
| if src_img is None: return None, "Error: No source image" | |
| cv2.imwrite(src_path, cv2.cvtColor(src_img, cv2.COLOR_RGB2BGR)) | |
| if isinstance(video, str) and os.path.exists(video): shutil.copy(video, dst_video_path) | |
| else: dst_video_path = video | |
| def extract_frames(video_path, output_folder): | |
| cap = cv2.VideoCapture(video_path) | |
| count = 0 | |
| paths = [] | |
| while cap.isOpened(): | |
| ret, frame = cap.read() | |
| if not ret: break | |
| fname = os.path.join(output_folder, f"frame_{count:05d}.jpg") | |
| cv2.imwrite(fname, frame) | |
| paths.append(fname) | |
| count += 1 | |
| cap.release() | |
| return paths | |
| def frames_to_video(input_folder, output_video, fps): | |
| images = sorted([img for img in os.listdir(input_folder) if img.endswith(".jpg")]) | |
| if not images: return | |
| frame0 = cv2.imread(os.path.join(input_folder, images[0])) | |
| h, w, l = frame0.shape | |
| out = cv2.VideoWriter(output_video, cv2.VideoWriter_fourcc(*'mp4v'), fps, (w, h)) | |
| for image in images: | |
| out.write(cv2.imread(os.path.join(input_folder, image))) | |
| out.release() | |
| progress(0.1, desc="Extracting frames...") | |
| frame_paths = extract_frames(dst_video_path, frames_dir) | |
| cap = cv2.VideoCapture(dst_video_path) | |
| fps = cap.get(cv2.CAP_PROP_FPS) | |
| cap.release() | |
| if fps == 0: fps = 30 | |
| 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: | |
| swapped = swapper.swap_faces(src_path, int(src_idx), frame_path, int(dst_idx)) | |
| cv2.imwrite(out_path, swapped) | |
| except: | |
| shutil.copy(frame_path, out_path) | |
| progress((idx+1)/len(frame_paths), desc=f"Frame {idx+1}/{len(frame_paths)}") | |
| progress(0.9, desc="Encoding video...") | |
| frames_to_video(swapped_dir, output_video_path, fps) | |
| if add_audio: add_audio_to_video(dst_video_path, output_video_path, final_output_path) | |
| else: final_output_path = output_video_path | |
| if delete_frames: | |
| shutil.rmtree(frames_dir, ignore_errors=True) | |
| shutil.rmtree(swapped_dir, ignore_errors=True) | |
| return final_output_path, "Done." | |
| def swap_video_all_faces(src_img, video, num_faces, delete_frames, add_audio, copy_drive, progress=gr.Progress()): | |
| return swap_video(src_img, 1, video, 1, delete_frames, add_audio, copy_drive, progress) | |
| # --- UI LAYOUT --- | |
| with gr.Blocks(theme=shinyy_theme, css=custom_css, title="Shinyy's Face Swapper") as demo: | |
| gr.Markdown("""<h1 class="title-header">s h i n y y ' s Face Swapping Tool</h1>""") | |
| with gr.Tabs(): | |
| # TAB 1: Single Photo | |
| with gr.TabItem("📸 Single Photo"): | |
| with gr.Row(): | |
| with gr.Column(variant="panel"): | |
| gr.Markdown("### Source") | |
| sp_src = gr.Image(label="Source Face", type="numpy", height=300) | |
| sp_src_idx = gr.Number(value=1, label="Face Index", precision=0) | |
| with gr.Column(variant="panel"): | |
| gr.Markdown("### Destination") | |
| sp_dst = gr.Image(label="Target Image", type="numpy", height=300) | |
| sp_dst_idx = gr.Number(value=1, label="Face Index", precision=0) | |
| with gr.Row(): | |
| sp_hair = gr.Checkbox(label="Swap Hair (Beta)", value=False, | |
| info="Warps source hair onto target.") | |
| sp_btn = gr.Button("Swap Face", variant="primary", size="lg") | |
| with gr.Row(): | |
| sp_out = gr.Image(label="Result", interactive=False) | |
| sp_log = gr.Textbox(label="Logs", lines=5, interactive=False) | |
| sp_btn.click(swap_single_photo, inputs=[sp_src, sp_src_idx, sp_dst, sp_dst_idx, sp_hair], outputs=[sp_out, sp_log]) | |
| # TAB 2: Batch | |
| with gr.TabItem("📂 One -> Many"): | |
| with gr.Row(): | |
| ssmd_src = gr.Image(label="Source", type="numpy", height=250) | |
| with gr.Column(): | |
| ssmd_dsts = gr.Gallery(label="Destinations", type="numpy", columns=3, height=250) | |
| ssmd_indices = gr.Textbox(label="Dst Indices", placeholder="1, 1, 1") | |
| ssmd_btn = gr.Button("Batch Swap", variant="primary") | |
| with gr.Row(): | |
| ssmd_out = gr.Gallery(label="Results", columns=3) | |
| ssmd_log = gr.Textbox(label="Logs", lines=5) | |
| ssmd_btn.click(swap_single_src_multi_dst, inputs=[ssmd_src, ssmd_dsts, ssmd_indices], outputs=[ssmd_out, ssmd_log]) | |
| # TAB 3: Multi -> Single | |
| with gr.TabItem("👥 Many -> One"): | |
| with gr.Row(): | |
| mssd_srcs = gr.Gallery(label="Sources", type="numpy", columns=3, height=250) | |
| with gr.Column(): | |
| mssd_dst = gr.Image(label="Target", type="numpy", height=250) | |
| mssd_idx = gr.Number(value=1, label="Target Index", precision=0) | |
| mssd_btn = gr.Button("Swap All", variant="primary") | |
| with gr.Row(): | |
| mssd_out = gr.Gallery(label="Results", columns=3) | |
| mssd_log = gr.Textbox(label="Logs", lines=5) | |
| mssd_btn.click(swap_multi_src_single_dst, inputs=[mssd_srcs, mssd_dst, mssd_idx], outputs=[mssd_out, mssd_log]) | |
| # TAB 4: Many -> Many | |
| with gr.TabItem("🔄 Many -> Many"): | |
| with gr.Row(): | |
| msmd_srcs = gr.Gallery(label="Sources", type="numpy", columns=3) | |
| msmd_dsts = gr.Gallery(label="Targets", type="numpy", columns=3) | |
| msmd_indices = gr.Textbox(label="Dst Indices", placeholder="1,1,1") | |
| msmd_btn = gr.Button("Mass Swap", variant="primary") | |
| with gr.Row(): | |
| msmd_out = gr.Gallery(label="Results") | |
| msmd_log = gr.Textbox(label="Logs", lines=5) | |
| msmd_btn.click(swap_multi_src_multi_dst, inputs=[msmd_srcs, msmd_dsts, msmd_indices], outputs=[msmd_out, msmd_log]) | |
| # TAB 5: Custom | |
| with gr.TabItem("🎛️ Custom Map"): | |
| with gr.Row(): | |
| custom_srcs = gr.Gallery(label="Sources", type="numpy", columns=4, height=200) | |
| custom_dst = gr.Image(label="Destination", type="numpy", height=300) | |
| custom_map = gr.Textbox(label="Map (e.g. 2,1,3)", placeholder="2,1,3") | |
| custom_btn = gr.Button("Execute", variant="primary") | |
| with gr.Row(): | |
| custom_out = gr.Image(label="Result") | |
| custom_log = gr.Textbox(label="Logs") | |
| custom_btn.click(swap_faces_custom, inputs=[custom_srcs, custom_dst, custom_map], outputs=[custom_out, custom_log]) | |
| # TAB 6: Video | |
| with gr.TabItem("🎥 Video"): | |
| with gr.Tabs(): | |
| with gr.TabItem("Single"): | |
| with gr.Row(): | |
| vid_src = gr.Image(label="Source", type="numpy") | |
| vid_target = gr.Video(label="Target Video") | |
| with gr.Row(): | |
| vid_src_idx = gr.Number(value=1, label="Src Idx", precision=0) | |
| vid_dst_idx = gr.Number(value=1, label="Dst Idx", precision=0) | |
| with gr.Accordion("Advanced", open=False): | |
| vid_del = gr.Checkbox(label="Delete Frames", value=True) | |
| vid_audio = gr.Checkbox(label="Audio", value=True) | |
| vid_drive = gr.Checkbox(label="Drive", value=False) | |
| vid_btn = gr.Button("Render", variant="primary") | |
| vid_out = gr.Video(label="Result") | |
| vid_log = gr.Textbox(label="Logs", lines=3) | |
| vid_btn.click(swap_video, inputs=[vid_src, vid_src_idx, vid_target, vid_dst_idx, vid_del, vid_audio, vid_drive], outputs=[vid_out, vid_log]) | |
| with gr.TabItem("All Faces"): | |
| with gr.Row(): | |
| va_src = gr.Image(label="Source", type="numpy") | |
| va_target = gr.Video(label="Target") | |
| va_count = gr.Number(value=1, label="Count", precision=0) | |
| va_btn = gr.Button("Render All", variant="primary") | |
| va_out = gr.Video(label="Result") | |
| va_btn.click(swap_video_all_faces, inputs=[va_src, va_target, va_count, vid_del, vid_audio, vid_drive], outputs=[va_out, vid_log]) | |
| # TAB 7: DOWNLOAD | |
| with gr.TabItem("💾 Download Source Code"): | |
| gr.Markdown("### Backup & Export") | |
| gr.Markdown("Click the button below to zip the entire project folder (excluding models and cache) so you can run it on your local PC.") | |
| dl_btn = gr.Button("📦 Zip & Download Project", variant="primary") | |
| dl_file = gr.File(label="Download Zip") | |
| dl_btn.click(zip_project_directory, inputs=[], outputs=[dl_file]) | |
| gr.Markdown("---") | |
| gr.Markdown("### s h i n y y ' s S u i t e v 2 . 1") | |
| if __name__ == "__main__": | |
| parser = argparse.ArgumentParser() | |
| parser.add_argument("--share", action="store_true", help="Launch with share link") | |
| args = parser.parse_args() | |
| demo.queue().launch(share=args.share) |