code refactorized

.gitignore

@@ -0,0 +1,10 @@
+Deep3DFaceRecon_pytorch/models/arcface_torch/backbones/__pycache__/
+Deep3DFaceRecon_pytorch/models/__pycache__/
+Deep3DFaceRecon_pytorch/util/__pycache__/
+arcface_torch/backbones/__pycache__/
+benchmark/__pycache__/
+HRNet/__pycache__/
+models/__pycache__/
+configs/__pycache__/
+.idea
+playground.py

Dockerfile

@@ -1,16 +0,0 @@
-FROM xuehy93/hififace:1.0
-
-
-RUN apt update && apt install -y wget
-
-
-WORKDIR /
-
-RUN wget https://public.ph.files.1drv.com/y4m_El1_AyFLmGuZaPWOqkytzM4qYtDc3BvNNL99JV1OLCEkmD4RTQjtHEXZ0SAWb7UPLV1IPB0KO2rFlyGJaV_kITLbuAHzJ73GwR_cgvXpkIGywaTnKsKVV1jJe1LoFcl7XsxatyGpaC8-Gupq6jjBnaqSBH4dgfYAmzUk8Wqiiuj_ml2duU7No0M1T426y3RqOJsqVHXEMVfV0B6HjzQFKCCZIgfHjjHvLIB3B3xP8Q?AVOverride=1 -O checkpoints.tar.gz
-
-RUN tar xfz checkpoints.tar.gz
-
-WORKDIR /app
-ADD ./ /app
-RUN chmod +x ./server.sh
-CMD ["./server.sh"]

README.md

@@ -1,11 +1,14 @@
 ---
-title: HiFiFace Inference
-emoji: 📉
-colorFrom: indigo
-colorTo: indigo
-sdk: docker
+title: HiFiFace image swap
+emoji: 👁
+colorFrom: green
+colorTo: green
+sdk: gradio
+sdk_version: 5.47.0
+app_file: app.py
 pinned: false
 license: mit
+short_description: Swap faces from photos
 ---
 
 Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference

app.py

@@ -1,134 +1,110 @@
-import argparse
+#######################################################################################
+#
+# MIT License
+#
+# Copyright (c) [2025] [leonelhs@gmail.com]
+#
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+#
+# The above copyright notice and this permission notice shall be included in all
+# copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+# SOFTWARE.
+#
+#######################################################################################
+#
+# Source code is based on or inspired by several projects.
+# For more details and proper attribution, please refer to the following resources:
+#
+# - [hyxue] - [https://huggingface.co/spaces/hyxue/HiFiFace-inference-demo]
+# - [maum-ai] [https://github.com/maum-ai/hififace]
+#
 
 import gradio as gr
+import torch
+from huggingface_hub import hf_hub_download
 
 from benchmark.app_image import ImageSwap
-from benchmark.app_video import VideoSwap
-from configs.train_config import TrainConfig
-from models.model import HifiFace
+from models.model import HifiFaceST, HifiFaceWGM
 
+REPO_ID = "leonelhs/HiFiFace"
+
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+gen_st_path = hf_hub_download(repo_id=REPO_ID,
+                              filename="hififace_pretrained/standard_model/generator_320000.pth")
+
+gen_wgm_path = hf_hub_download(repo_id=REPO_ID,
+                               filename="hififace_pretrained/with_gaze_and_mouth/generator_190000.pth")
+
+fade_detector_path = hf_hub_download(repo_id=REPO_ID,
+                                     filename="face_detector/face_detector_scrfd_10g_bnkps.onnx")
+
+identity_extractor_config = {
+    "f_3d_checkpoint_path": hf_hub_download(repo_id=REPO_ID, filename="Deep3DFaceRecon/epoch_20.pth"),
+    "f_id_checkpoint_path": hf_hub_download(repo_id=REPO_ID, filename="arcface/ms1mv3_arcface_r100_fp16_backbone.pth")
+}
 
 class ConfigPath:
-    face_detector_weights = "/checkpoints/face_detector/face_detector_scrfd_10g_bnkps.onnx"
+    face_detector_weights = fade_detector_path
     model_path = ""
     model_idx = 80000
-    ffmpeg_device = "cuda"
-    device = "cuda"
+    ffmpeg_device = device
+    device = device
 
+cfg = ConfigPath()
 
-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"):
+model_standard = HifiFaceST(identity_extractor_config, device=device, generator_path=gen_st_path)
+
+model_wgm = HifiFaceWGM(identity_extractor_config, device=device, generator_path=gen_wgm_path)
+
+image_infer_standard = ImageSwap(cfg, model_standard)
+image_infer_wgm = ImageSwap(cfg, model_wgm)
+
+MODELS = {
+    "Standard model": "standard",
+    "Eye and mouth hm loss": "eyeandmouth",
+}
+
+def inference_image(source_face, target_face, method="standard", shape_rate=1.0, id_rate=1.0, iterations=1):
+    if method == "standard":
+        return target_face, image_infer_standard.inference(source_face, target_face, shape_rate, id_rate, int(iterations))
+    return target_face, image_infer_wgm.inference(source_face, target_face, shape_rate, id_rate, int(iterations))
+
+
+with gr.Blocks(title="FaceSwap") as app:
+    gr.Markdown("## HiFiFace image swap")
+    with gr.Row():
+        with gr.Column(scale=1):
             with gr.Row():
-                source_image = gr.Image(shape=None, label="source image")
-                target_image = gr.Image(shape=None, label="target image")
+                source_image = gr.Image(type="numpy", label="Face image")
+                target_image = gr.Image(type="numpy", label="Body image")
+            mod = gr.Dropdown(choices=list(MODELS.items()), label="Model generator", value="standard")
+            image_btn = gr.Button("Swap image")
+            with gr.Accordion("Fine tunes", open=False):
+                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")
+        with gr.Column(scale=1):
             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()
+                output_image = gr.ImageSlider(label="Swapped image", type="pil")
+
+    image_btn.click(
+        fn=inference_image,
+        inputs=[source_image, target_image, mod, structure_sim, id_sim, iters],
+        outputs=output_image,
+    )
+
+app.launch(share=False, debug=True, show_error=True, mcp_server=True, pwa=True)
+app.queue()

models/model.py

@@ -1,6 +1,6 @@
 import os
+from abc import abstractmethod
 from typing import Dict
-from typing import Optional
 from typing import Tuple
 
 import kornia
@@ -10,28 +10,30 @@ import torch.nn as nn
 import torch.nn.functional as F
 from loguru import logger
 
-from arcface_torch.backbones.iresnet import iresnet100
-from configs.train_config import TrainConfig
 from Deep3DFaceRecon_pytorch.models.bfm import ParametricFaceModel
 from Deep3DFaceRecon_pytorch.models.networks import ReconNetWrapper
 from HRNet.hrnet import HighResolutionNet
+from arcface_torch.backbones.iresnet import iresnet100
 from models.discriminator import Discriminator
 from models.gan_loss import GANLoss
 from models.generator import Generator
 from models.init_weight import init_net
 
-
 class HifiFace:
     def __init__(
         self,
         identity_extractor_config,
-        is_training=True,
-        device="cpu",
-        load_checkpoint: Optional[Tuple[str, int]] = None,
+        generator_path,
+        is_training=False,
+        device="cpu"
     ):
         super(HifiFace, self).__init__()
+        self.d_optimizer = None
+        self.g_optimizer = None
         self.generator = Generator(identity_extractor_config)
         self.is_training = is_training
+        self.device = device
+        self.generator_path = generator_path
 
         if self.is_training:
             self.lr = TrainConfig().lr
@@ -80,10 +82,9 @@ class HifiFace:
 
             self.dilation_kernel = torch.ones(5, 5)
 
-        if load_checkpoint is not None:
-            self.load(load_checkpoint[0], load_checkpoint[1])
+        self.load_checkpoint()
 
-        self.setup(device)
+        self.setup(self.device)
 
     def save(self, path, idx=None):
         os.makedirs(path, exist_ok=True)
@@ -100,18 +101,9 @@ class HifiFace:
             torch.save(self.generator.state_dict(), g_path)
             torch.save(self.discriminator.state_dict(), d_path)
 
-    def load(self, path, idx=None):
-        if idx is None:
-            g_path = os.path.join(path, "generator.pth")
-            d_path = os.path.join(path, "discriminator.pth")
-        else:
-            g_path = os.path.join(path, f"generator_{idx}.pth")
-            d_path = os.path.join(path, f"discriminator_{idx}.pth")
-        logger.info(f"Loading generator from {g_path}")
-        self.generator.load_state_dict(torch.load(g_path, map_location="cpu"))
-        if self.is_training:
-            logger.info(f"Loading discriminator from {d_path}")
-            self.discriminator.load_state_dict(torch.load(d_path, map_location="cpu"))
+    @abstractmethod
+    def load_checkpoint(self):
+        pass
 
     def setup(self, device):
         self.generator.to(device)
@@ -399,37 +391,18 @@ class HifiFace:
         }
 
 
-if __name__ == "__main__":
-    import torch
-    import cv2
-    from configs.train_config import TrainConfig
-
-    identity_extractor_config = TrainConfig().identity_extractor_config
-
-    model = HifiFace(identity_extractor_config, is_training=True)
-
-    # src = cv2.imread("/home/xuehongyang/data/test1.jpg")
-    # tgt = cv2.imread("/home/xuehongyang/data/test2.jpg")
-    # src = cv2.cvtColor(src, cv2.COLOR_BGR2RGB)
-    # tgt = cv2.cvtColor(tgt, cv2.COLOR_BGR2RGB)
-    # src = cv2.resize(src, (256, 256))
-    # tgt = cv2.resize(tgt, (256, 256))
-    # src = src.transpose(2, 0, 1)[None, ...]
-    # tgt = tgt.transpose(2, 0, 1)[None, ...]
-    # source_img = torch.from_numpy(src).float() / 255.0
-    # target_img = torch.from_numpy(tgt).float() / 255.0
-    # same_id_mask = torch.Tensor([1]).unsqueeze(0)
-    # tgt_mask = target_img[:, 0, :, :].unsqueeze(1)
-    # if torch.cuda.is_available():
-    #     model.to("cuda:3")
-    #     source_img = source_img.to("cuda:3")
-    #     target_img = target_img.to("cuda:3")
-    #     tgt_mask = tgt_mask.to("cuda:3")
-    #     same_id_mask = same_id_mask.to("cuda:3")
-    #     source_img = source_img.repeat(16, 1, 1, 1)
-    #     target_img = target_img.repeat(16, 1, 1, 1)
-    #     tgt_mask = tgt_mask.repeat(16, 1, 1, 1)
-    #     same_id_mask = same_id_mask.repeat(16, 1)
-    # while True:
-    #     x = model.optimize(source_img, target_img, tgt_mask, same_id_mask)
-    #     print(x[0]["loss_generator"])
+class HifiFaceST(HifiFace):
+    def __init__(self, identity_extractor_config, device, generator_path):
+        super().__init__(identity_extractor_config, device=device, generator_path=generator_path)
+
+    def load_checkpoint(self):
+        self.generator.load_state_dict(torch.load(self.generator_path, map_location=self.device))
+        logger.info(f"Loading generator from {self.generator_path}")
+
+class HifiFaceWGM(HifiFace):
+    def __init__(self, identity_extractor_config, device, generator_path):
+        super().__init__(identity_extractor_config, device=device, generator_path=generator_path)
+
+    def load_checkpoint(self):
+        self.generator.load_state_dict(torch.load(self.generator_path, map_location=self.device))
+        logger.info(f"Loading generator from {self.generator_path}")

packages.txt

@@ -1 +0,0 @@
-wget

pyrightconfig.json

@@ -1,15 +0,0 @@
-{
-    "reportMissingImports": true,
-    "reportMissingTypeStubs": true,
-    "useLibraryCodeForTypes": true,
-    "reportUnusedImport": "warning",
-    "reportUnusedVariable": "warning",
-    "reportDuplicateImport": true,
-    "reportPrivateImportUsage": false,
-    "reportWildcardImportFromLibrary": "warning",
-    "reportTypedDictNotRequiredAccess": false,
-    "reportGeneralTypeIssues": false,
-    "venvPath": "/home/xuehongyang/miniconda3/envs/",
-    "venv": "pytorch-2.0",
-    "stubPath": "/home/xuehongyang/dev_configs/typings"
-}

server.sh

@@ -1 +0,0 @@
-python3 app.py