File size: 4,195 Bytes
a526c1c
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
import gradio as gr
import cv2
import multiprocessing
import os
import requests
from refacer import Refacer

# Hugging Face URL to download the model
model_url = "https://huggingface.co/ofter/4x-UltraSharp/resolve/main/inswapper_128.onnx"
model_path = "./inswapper_128.onnx"

# Function to download the model
def download_model():
    if not os.path.exists(model_path):
        print("Downloading inswapper_128.onnx...")
        response = requests.get(model_url)
        if response.status_code == 200:
            with open(model_path, 'wb') as f:
                f.write(response.content)
            print("Model downloaded successfully!")
        else:
            raise Exception(f"Failed to download the model. Status code: {response.status_code}")
    else:
        print("Model already exists.")

# Download the model when the script runs
download_model()

# Initialize Refacer class (force CPU mode)
refacer = Refacer(force_cpu=True)

# Dummy function to simulate frame-level processing
def process_frame(frame, origin_face, destination_face, threshold):
    # Simulate face swapping or any processing needed
    result_frame = refacer.reface(frame, [{
        'origin': origin_face,
        'destination': destination_face,
        'threshold': threshold
    }])
    return result_frame

# Function to process the video in parallel using multiprocessing
def process_video(video_path, origins, destinations, thresholds, max_processes=2):
    cap = cv2.VideoCapture(video_path)
    frames = []

    # Read all frames from the video
    while cap.isOpened():
        ret, frame = cap.read()
        if not ret:
            break
        frames.append(frame)

    cap.release()

    # Parallel processing of frames with limited processes (for CPU optimization)
    with multiprocessing.Pool(processes=max_processes) as pool:
        processed_frames = pool.starmap(process_frame, [
            (frame, origins[min(i, len(origins) - 1)], destinations[min(i, len(destinations) - 1)], thresholds[min(i, len(thresholds) - 1)])
            for i, frame in enumerate(frames)
        ])

    # Saving the processed frames back into a video
    output_video_path = "processed_video.mp4"
    fourcc = cv2.VideoWriter_fourcc(*'mp4v')  # Compression using mp4 codec
    out = cv2.VideoWriter(output_video_path, fourcc, 30.0, (640, 360))  # Reduce resolution to speed up processing

    for frame in processed_frames:
        out.write(frame)

    out.release()
    return output_video_path

# Gradio Interface function
def run(video_path, *vars):
    # Split the inputs into origins, destinations, and thresholds based on num_faces
    num_faces = 5  # You can adjust this based on your UI
    origins = vars[:num_faces]
    destinations = vars[num_faces:2*num_faces]
    thresholds = vars[2*num_faces:]

    # Ensure there are no index errors by limiting the number of inputs
    if len(origins) != num_faces or len(destinations) != num_faces or len(thresholds) != num_faces:
        return "Please provide input for all faces."

    refaced_video_path = process_video(video_path, origins, destinations, thresholds)
    print(f"Refaced video can be found at {refaced_video_path}")

    return refaced_video_path

# Prepare Gradio components
origin = []
destination = []
thresholds = []

with gr.Blocks() as demo:
    with gr.Row():
        gr.Markdown("# Refacer")
    with gr.Row():
        video_input = gr.Video(label="Original video", format="mp4")
        video_output = gr.Video(label="Refaced video", interactive=False, format="mp4")

    for i in range(5):  # Set max faces to 5
        with gr.Tab(f"Face #{i+1}"):
            with gr.Row():
                origin.append(gr.Image(label="Face to replace"))
                destination.append(gr.Image(label="Destination face"))
            with gr.Row():
                thresholds.append(gr.Slider(label="Threshold", minimum=0.0, maximum=1.0, value=0.2))
    
    with gr.Row():
        button = gr.Button("Reface", variant="primary")

    button.click(fn=run, inputs=[video_input] + origin + destination + thresholds, outputs=[video_output])

# Launch the Gradio app
demo.launch(show_error=True, server_name="0.0.0.0", server_port=7860)