Spaces:
Sleeping
Sleeping
| import os | |
| # ========================= | |
| # FORCE CPU MODE | |
| # ========================= | |
| os.environ["CUDA_VISIBLE_DEVICES"] = "" | |
| os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "" | |
| import torch | |
| import sys | |
| import asyncio | |
| import imageio | |
| import tempfile | |
| import numpy as np | |
| import gradio as gr | |
| from typing import Sequence, Mapping, Any, Union | |
| from PIL import Image | |
| from huggingface_hub import hf_hub_download | |
| # ========================= | |
| # DOWNLOAD MODELS (ONLY IF NOT EXISTS) | |
| # ========================= | |
| def download_if_not_exists(repo, filename, local_dir): | |
| path = os.path.join(local_dir, filename) | |
| if not os.path.exists(path): | |
| os.makedirs(local_dir, exist_ok=True) | |
| hf_hub_download(repo_id=repo, filename=filename, local_dir=local_dir) | |
| download_if_not_exists("ezioruan/inswapper_128.onnx", "inswapper_128.onnx", "models/insightface") | |
| download_if_not_exists("martintomov/comfy", "facerestore_models/GPEN-BFR-512.onnx", "models") | |
| download_if_not_exists("facefusion/models-3.3.0", "hyperswap_1a_256.onnx", "models/hyperswap") | |
| download_if_not_exists("facefusion/models-3.3.0", "hyperswap_1b_256.onnx", "models/hyperswap") | |
| download_if_not_exists("facefusion/models-3.3.0", "hyperswap_1c_256.onnx", "models/hyperswap") | |
| # ========================= | |
| # COMFY INIT (GIỮ NGUYÊN) | |
| # ========================= | |
| from comfy.model_management import CPUState | |
| import comfy.model_management | |
| comfy.model_management.cpu_state = CPUState.CPU | |
| def get_value_at_index(obj: Union[Sequence, Mapping], index: int) -> Any: | |
| try: | |
| return obj[index] | |
| except Exception: | |
| return obj["result"][index] | |
| def find_path(name: str, path: str = None) -> str: | |
| if path is None: | |
| path = os.getcwd() | |
| if name in os.listdir(path): | |
| return os.path.join(path, name) | |
| parent = os.path.dirname(path) | |
| if parent == path: | |
| return None | |
| return find_path(name, parent) | |
| def add_comfyui_directory_to_sys_path(): | |
| comfyui_path = find_path("ComfyUI") | |
| if comfyui_path and os.path.isdir(comfyui_path): | |
| sys.path.append(comfyui_path) | |
| add_comfyui_directory_to_sys_path() | |
| def add_extra_model_paths(): | |
| try: | |
| from main import load_extra_path_config | |
| except ImportError: | |
| from utils.extra_config import load_extra_path_config | |
| extra_model_paths = find_path("extra_model_paths.yaml") | |
| if extra_model_paths: | |
| load_extra_path_config(extra_model_paths) | |
| add_extra_model_paths() | |
| def import_custom_nodes(): | |
| import execution | |
| from nodes import init_extra_nodes | |
| import server | |
| loop = asyncio.new_event_loop() | |
| asyncio.set_event_loop(loop) | |
| server_instance = server.PromptServer(loop) | |
| execution.PromptQueue(server_instance) | |
| loop.run_until_complete(init_extra_nodes()) | |
| import_custom_nodes() | |
| from nodes import NODE_CLASS_MAPPINGS | |
| loadimage = NODE_CLASS_MAPPINGS["LoadImage"]() | |
| reactorfaceswap = NODE_CLASS_MAPPINGS["ReActorFaceSwap"]() | |
| # ========================= | |
| # MAIN FUNCTION | |
| # ========================= | |
| def generate_image(source_files, target_files, target_index, | |
| swap_model, face_restore_model, restore_strength): | |
| os.makedirs("output", exist_ok=True) | |
| output_paths = [] | |
| if not source_files or not target_files: | |
| return [] | |
| with torch.inference_mode(): | |
| for s in source_files: | |
| source_path = s.name | |
| loadimage_source = loadimage.load_image(image=source_path) | |
| source_tensor = get_value_at_index(loadimage_source, 0) | |
| source_base = os.path.splitext(os.path.basename(source_path))[0] | |
| for t in target_files: | |
| target_path = t.name | |
| target_base = os.path.splitext(os.path.basename(target_path))[0] | |
| # ================= GIF ================= | |
| if target_path.lower().endswith(".gif"): | |
| reader = imageio.get_reader(target_path) | |
| frames = [] | |
| durations = [] | |
| for i, frame in enumerate(reader): | |
| frame_rgb = Image.fromarray(frame).convert("RGB") | |
| frames.append(np.array(frame_rgb)) | |
| meta = reader.get_meta_data(index=i) | |
| durations.append(meta.get("duration", 100)) | |
| reader.close() | |
| output_frames = [] | |
| for frame in frames: | |
| with tempfile.NamedTemporaryFile(suffix=".png", delete=False) as tmp: | |
| Image.fromarray(frame).save(tmp.name) | |
| temp_path = tmp.name | |
| loadimage_target = loadimage.load_image(image=temp_path) | |
| target_tensor = get_value_at_index(loadimage_target, 0) | |
| result = reactorfaceswap.execute( | |
| enabled=True, | |
| swap_model=swap_model, | |
| facedetection="YOLOv5l", | |
| face_restore_model=face_restore_model, | |
| face_restore_visibility=restore_strength, | |
| codeformer_weight=0.5, | |
| detect_gender_input="no", | |
| detect_gender_source="no", | |
| input_faces_index=str(target_index), | |
| source_faces_index="0", | |
| console_log_level=1, | |
| input_image=target_tensor, | |
| source_image=source_tensor, | |
| ) | |
| swapped = get_value_at_index(result, 0)[0] | |
| if isinstance(swapped, torch.Tensor): | |
| swapped = swapped.detach().cpu().float().numpy() | |
| if swapped.max() <= 1.0: | |
| swapped *= 255.0 | |
| swapped = np.clip(swapped, 0, 255).astype(np.uint8) | |
| output_frames.append(Image.fromarray(swapped).convert("RGB")) | |
| os.remove(temp_path) | |
| output_path = f"output/{source_base}_to_{target_base}.webp" | |
| output_frames[0].save( | |
| output_path, | |
| save_all=True, | |
| append_images=output_frames[1:], | |
| duration=durations, | |
| loop=0, | |
| format="WEBP", | |
| quality=90, | |
| method=6 | |
| ) | |
| # ================= IMAGE ================= | |
| else: | |
| loadimage_target = loadimage.load_image(image=target_path) | |
| target_tensor = get_value_at_index(loadimage_target, 0) | |
| result = reactorfaceswap.execute( | |
| enabled=True, | |
| swap_model=swap_model, | |
| facedetection="YOLOv5l", | |
| face_restore_model=face_restore_model, | |
| face_restore_visibility=restore_strength, | |
| codeformer_weight=0.5, | |
| detect_gender_input="no", | |
| detect_gender_source="no", | |
| input_faces_index=str(target_index), | |
| source_faces_index="0", | |
| console_log_level=1, | |
| input_image=target_tensor, | |
| source_image=source_tensor, | |
| ) | |
| swapped = get_value_at_index(result, 0)[0] | |
| if isinstance(swapped, torch.Tensor): | |
| swapped = swapped.detach().cpu().float().numpy() | |
| if swapped.max() <= 1.0: | |
| swapped *= 255.0 | |
| swapped = np.clip(swapped, 0, 255).astype(np.uint8) | |
| output_path = f"output/{source_base}_to_{target_base}.webp" | |
| Image.fromarray(swapped).save( | |
| output_path, | |
| format="WEBP", | |
| quality=90, | |
| method=6 | |
| ) | |
| output_paths.append(output_path) | |
| return output_paths | |
| # ========================= | |
| # GRADIO UI | |
| # ========================= | |
| with gr.Blocks() as app: | |
| source_files = gr.File(label="Source Faces", file_count="multiple", interactive=True) | |
| target_files = gr.File(label="Target Images / GIFs", file_count="multiple", interactive=True) | |
| swap_model = gr.Dropdown( | |
| choices=["inswapper_128.onnx", | |
| "hyperswap_1a_256.onnx", | |
| "hyperswap_1b_256.onnx", | |
| "hyperswap_1c_256.onnx"], | |
| value="hyperswap_1b_256.onnx", | |
| label="Swap Model" | |
| ) | |
| face_restore_model = gr.Dropdown( | |
| choices=["none", "GPEN-BFR-512.onnx"], | |
| value="none", | |
| label="Face Restore Model" | |
| ) | |
| restore_strength = gr.Slider(0, 1, 0.7, step=0.05) | |
| target_index = gr.Dropdown([0,1,2,3,4], value=0) | |
| generate_btn = gr.Button("Generate") | |
| output_files = gr.Files(label="Output WebPs") | |
| generate_btn.click( | |
| fn=generate_image, | |
| inputs=[source_files, target_files, target_index, | |
| swap_model, face_restore_model, restore_strength], | |
| outputs=output_files | |
| ) | |
| app.launch(share=True) |