Spaces:
Runtime error
Runtime error
File size: 5,165 Bytes
83d8d3c 54a5078 83d8d3c 54a5078 83d8d3c 54a5078 83d8d3c 54a5078 83d8d3c 54a5078 83d8d3c 54a5078 83d8d3c 54a5078 83d8d3c 54a5078 83d8d3c 54a5078 83d8d3c 54a5078 83d8d3c 54a5078 83d8d3c 54a5078 83d8d3c 54a5078 83d8d3c |
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 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 |
import argparse
import gradio as gr
from benchmark.app_image import ImageSwap
from benchmark.app_video import VideoSwap
from configs.train_config import TrainConfig
from models.model import HifiFace
class ConfigPath:
face_detector_weights = "/checkpoints/face_detector/face_detector_scrfd_10g_bnkps.onnx"
model_path = ""
model_idx = 80000
ffmpeg_device = "cuda"
device = "cuda"
def main():
cfg = ConfigPath()
parser = argparse.ArgumentParser(
prog="benchmark", description="What the program does", epilog="Text at the bottom of help"
)
parser.add_argument("-m", "--model_path", default="/checkpoints/hififace_pretrained/standard_model")
parser.add_argument("-i", "--model_idx", default="320000")
parser.add_argument("-f", "--ffmpeg_device", default="cpu")
parser.add_argument("-d", "--device", default="cpu")
args = parser.parse_args()
cfg.model_path = args.model_path
cfg.model_idx = int(args.model_idx)
cfg.ffmpeg_device = args.ffmpeg_device
cfg.device = args.device
opt = TrainConfig()
checkpoint = (cfg.model_path, cfg.model_idx)
model_path_1 = "/checkpoints/hififace_pretrained/with_gaze_and_mouth"
checkpoint1 = ("/checkpoints/hififace_pretrained/with_gaze_and_mouth", "190000")
model = HifiFace(opt.identity_extractor_config, is_training=False, device=cfg.device, load_checkpoint=checkpoint)
model1 = HifiFace(opt.identity_extractor_config, is_training=False, device=cfg.device, load_checkpoint=checkpoint1)
image_infer = ImageSwap(cfg, model)
image_infer1 = ImageSwap(cfg, model1)
def inference_image(source_face, target_face, shape_rate, id_rate, iterations):
return image_infer.inference(source_face, target_face, shape_rate, id_rate, int(iterations))
def inference_image1(source_face, target_face, shape_rate, id_rate, iterations):
return image_infer1.inference(source_face, target_face, shape_rate, id_rate, int(iterations))
model_name = cfg.model_path.split("/")[-1] + ":" + f"{cfg.model_idx}"
model_name1 = model_path_1.split("/")[-1] + ":" + "190000"
with gr.Blocks(title="FaceSwap") as demo:
gr.Markdown(
f"""
### standard model: {model_name} \n
### model with eye and mouth hm loss: {model_name1}
"""
)
with gr.Tab("Image swap with standard model"):
with gr.Row():
source_image = gr.Image(shape=None, label="source image")
target_image = gr.Image(shape=None, label="target image")
with gr.Row():
with gr.Column():
structure_sim = gr.Slider(
minimum=0.0,
maximum=1.0,
value=1.0,
step=0.1,
label="3d similarity",
)
id_sim = gr.Slider(
minimum=0.0,
maximum=1.0,
value=1.0,
step=0.1,
label="id similarity",
)
iters = gr.Slider(
minimum=1,
maximum=10,
value=1,
step=1,
label="iters",
)
image_btn = gr.Button("image swap")
output_image = gr.Image(shape=None, label="Result")
image_btn.click(
fn=inference_image,
inputs=[source_image, target_image, structure_sim, id_sim, iters],
outputs=output_image,
)
with gr.Tab("Image swap with eye&mouth hm loss model"):
with gr.Row():
source_image = gr.Image(shape=None, label="source image")
target_image = gr.Image(shape=None, label="target image")
with gr.Row():
with gr.Column():
structure_sim = gr.Slider(
minimum=0.0,
maximum=1.0,
value=1.0,
step=0.1,
label="3d similarity",
)
id_sim = gr.Slider(
minimum=0.0,
maximum=1.0,
value=1.0,
step=0.1,
label="id similarity",
)
iters = gr.Slider(
minimum=1,
maximum=10,
value=1,
step=1,
label="iters",
)
image_btn = gr.Button("image swap")
output_image = gr.Image(shape=None, label="Result")
image_btn.click(
fn=inference_image1,
inputs=[source_image, target_image, structure_sim, id_sim, iters],
outputs=output_image,
)
demo.launch(server_name="0.0.0.0", server_port=7860)
if __name__ == "__main__":
main()
|