From 203e571aba6b93b58801857df04e90a785ce365c Mon Sep 17 00:00:00 2001
From: wangzhibo <wangzhibo@shsening.com>
Date: 星期三, 30 七月 2025 22:11:26 +0800
Subject: [PATCH] add face swapper
---
auto_media_publisher/modifiers/face_swap/face_swaper_ffmpeg.py | 79 ++++++++++++++++++---------------------
1 files changed, 36 insertions(+), 43 deletions(-)
diff --git a/auto_media_publisher/utils/face_swap/face_swaper_ffmpeg.py b/auto_media_publisher/modifiers/face_swap/face_swaper_ffmpeg.py
similarity index 80%
rename from auto_media_publisher/utils/face_swap/face_swaper_ffmpeg.py
rename to auto_media_publisher/modifiers/face_swap/face_swaper_ffmpeg.py
index 7b28497..12f4d11 100644
--- a/auto_media_publisher/utils/face_swap/face_swaper_ffmpeg.py
+++ b/auto_media_publisher/modifiers/face_swap/face_swaper_ffmpeg.py
@@ -15,7 +15,9 @@
from .face_analyser import FaceAnalyser
from auto_media_publisher.utils.logger import get_logger
from auto_media_publisher.config.conf_base import FACE_SWAPPER_MODEL_PATH
-
+from auto_media_publisher.utils.media_input_output import VideoReader,VideoWriter
+from torch.utils.data import DataLoader
+from torchvision import transforms
class FaceSwapWorker:
def __init__(self, enable_enhancer:bool = False,enable_nsfw_predictor:bool=False):
@@ -25,6 +27,7 @@
self.main_analyser = FaceAnalyser()
self.enable_enhancer = enable_enhancer
self.enable_nsfw_predictor = enable_nsfw_predictor
+ self.batch_size = 4
if self.enable_nsfw_predictor:
self.nsfw_predictor = NSFWPredictor()
self._init()
@@ -48,56 +51,46 @@
return 'mps'
else:
return 'cpu'
+
+ def auto_downsample_ratio(self,h, w):
+ """
+ Automatically find a downsample ratio so that the largest side of the resolution be 512px.
+ """
+ return min(512 / max(h, w), 1)
def _decode_video(self):
- # === Step 1: 鐢� ffprobe 鑾峰彇瑙嗛淇℃伅 ===
- probe_cmd = [
- "ffprobe",
- "-v", "error",
- "-select_streams", "v:0",
- "-show_entries", "stream=width,height,r_frame_rate",
- "-of", "json",
- self.input_video_path
- ]
- probe_result = subprocess.run(probe_cmd, capture_output=True, text=True, check=True)
- probe_data = json.loads(probe_result.stdout)
- video_stream = probe_data['streams'][0]
- self.width = int(video_stream['width'])
- self.height = int(video_stream['height'])
- self.fps = eval(video_stream['r_frame_rate']) # e.g. "30/1"
+ transform = transforms.ToTensor()
+ source = VideoReader(self.input_video_path, transform)
+
+ reader = DataLoader(source, batch_size=self.batch_size, pin_memory=True, num_workers=self.num_processors)
- print(f"Video info: {self.width}x{self.height} @ {self.fps}fps")
+ writer = VideoWriter(
+ path=self.output_video_path,
+ frame_rate=30,
+ bit_rate=int(1 * 1000000))
+
+ for src in reader:
- # === Step 2: 鐢� ffmpeg 瑙ｇ爜杈撳嚭 rawvideo 鍒� stdout pipe ===
- decode_cmd = [
- "ffmpeg",
- "-loglevel", "info",
- "-i", self.input_video_path,
- "-f", "rawvideo",
- "-pix_fmt", "rgb24",
- "pipe:1"
- ]
- process = subprocess.Popen(decode_cmd, stdout=subprocess.PIPE, bufsize=10**8)
+ if downsample_ratio is None:
+ downsample_ratio = auto_downsample_ratio(*src.shape[2:])
- frame_size = self.width * self.height * 3
- idx = 0
- while True:
- in_bytes = process.stdout.read(frame_size)
- if not in_bytes:
- break
- frame = np.frombuffer(in_bytes, dtype=np.uint8).reshape((self.height, self.width, 3))
- self.frame_queue.put((idx, frame)) # 鏀逛负 tuple 褰㈠紡
- idx += 1
+ src = src.to(device, dtype, non_blocking=True).unsqueeze(0) # [B, T, C, H, W]
+ fgr, pha, *rec = model(src, *rec, downsample_ratio)
- self.total_frames = idx
- print(f"_decode_video Total frames: {self.total_frames}")
+ if output_foreground is not None:
+ writer_fgr.write(fgr[0])
+ if output_alpha is not None:
+ writer_pha.write(pha[0])
+ if output_composition is not None:
+ if output_type == 'video':
+ com = fgr * pha + bgr * (1 - pha)
+ else:
+ fgr = fgr * pha.gt(0)
+ com = torch.cat([fgr, pha], dim=-3)
+ writer_com.write(com[0])
+
- process.stdout.close()
- process.wait()
- # 閫氱煡澶勭悊绾跨▼缁撴潫
- for _ in range(self.num_processors):
- self.frame_queue.put(None)
def _process_frame(self):
thread_face_analyser = FaceAnalyser()
--
Gitblit v1.9.1