add face swapper

wangzhibo

2025-07-30 203e571aba6b93b58801857df04e90a785ce365c

 auto_media_publisher/modifiers/face_swap/face_swaper_ffmpeg.py

File was renamed from auto_media_publisher/utils/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()