auto_media_publisher/modifiers/face_swap/__init__.py

 auto_media_publisher/modifiers/face_swap/face_analyser.py

File was renamed from auto_media_publisher/utils/face_swap/face_analyser.py
@@ -18,7 +18,6 @@
        self._analyser.prepare(ctx_id=0, det_size=(640, 640))

    def _select_device(self):
        return 'cpu'
        if torch.cuda.is_available():
            return 'cuda'
        elif torch.backends.mps.is_available():

 auto_media_publisher/modifiers/face_swap/face_enhancer.py

File was renamed from auto_media_publisher/utils/face_swap/face_enhancer.py
@@ -16,7 +16,6 @@
        self._enhancer = GFPGANer(model_path=FACE_ENHANCER_MODEL_PATH, upscale=1, device=self.device)

    def _select_device(self):
        return 'cpu'
        if torch.cuda.is_available():
            return 'cuda'
        elif torch.backends.mps.is_available():

 auto_media_publisher/modifiers/face_swap/face_nsfw_predictor.py

 auto_media_publisher/modifiers/face_swap/face_swaper_avpy.py

New file
@@ -0,0 +1,166 @@
import av
import torch
import threading, queue
from threading import current_thread
import numpy as np
import insightface
import cv2
import time
from insightface.app.common import Face
from gfpgan.utils import GFPGANer
from .typing import Frame # 你已有定义
from auto_media_publisher.config.conf_base import FACE_ENHANCER_MODEL_PATH,FACE_SWAPPER_MODEL_PATH
from insightface.app import FaceAnalysis

SIMILAR_FACE_DISTANCE: float = 0.45     # < 0.4 同一人 0.4 -0.6 大概是同一人 >0.6 不是同一个

def _select_device():
    if torch.cuda.is_available():
        return 'cuda:0'
    elif torch.backends.mps.is_available():
        return 'mps'
    else:
        return 'cpu'
    
NUM_WORKERS = 1

#face_swapper = insightface.model_zoo.get_model(FACE_SWAPPER_MODEL_PATH, providers=['CUDAExecutionProvider'])
#face_enhancer = GFPGANer(model_path=FACE_ENHANCER_MODEL_PATH, upscale=1, device=_select_device())

frame_q = queue.Queue(maxsize=10)
processed_q = queue.Queue(maxsize=10)

THREAD_SEMAPHORE =threading.Semaphore(4)  # 不限制了。

def enhance_face(thread_face_enhancer,target_face, temp_frame:np.ndarray) -> np.ndarray:
    start_x, start_y, end_x, end_y = map(int, target_face['bbox'])
    padding_x = int((end_x - start_x) * 0.5)
    padding_y = int((end_y - start_y) * 0.5)
    start_x = max(0, start_x - padding_x)
    start_y = max(0, start_y - padding_y)
    end_x = max(0, end_x + padding_x)
    end_y = max(0, end_y + padding_y)
    temp_face = temp_frame[start_y:end_y, start_x:end_x]
    #print(temp_face.size)
    if temp_face.size:
        _, _, temp_face = thread_face_enhancer.enhance(
                temp_face,
                paste_back=True
        )
        temp_frame[start_y:end_y, start_x:end_x] = temp_face
    return temp_frame
    

def decode_video(video_path):
    container = av.open(video_path)
    for idx, frame in enumerate(container.decode(video=0)):
        img = frame.to_ndarray(format='rgb24')
        frame_q.put((idx, img))
    for _ in range(NUM_WORKERS):
        frame_q.put((None, None))

def process_frame_loop(new_face):
    start_time = time.time()
    thread_face_analyser = FaceAnalysis(name='buffalo_l', root='./', providers=['CUDAExecutionProvider'])
    thread_face_analyser.prepare(ctx_id=3, det_size=(640, 640))
    #thread_face_enhancer = GFPGANer(model_path=FACE_ENHANCER_MODEL_PATH, upscale=1, device=_select_device())
    thread_face_enhancer = GFPGANer(model_path=FACE_ENHANCER_MODEL_PATH, upscale=1, device=_select_device())
    thread_face_swapper = insightface.model_zoo.get_model(FACE_SWAPPER_MODEL_PATH, providers=['CUDAExecutionProvider'])

    while True:
        idx, frame = frame_q.get()
        if idx is None:
            processed_q.put((None, None))
            break
        try:
            many_faces = thread_face_analyser.get(frame)
            if many_faces:
                old_face = many_faces[0]
                frame = thread_face_swapper.get(frame, old_face, new_face, paste_back=True)

                for sinlge_face in many_faces:
                    frame_ = enhance_face(thread_face_enhancer,sinlge_face, frame)                

            processed_q.put((idx, frame_))
            print(f"{current_thread().name} has proceesed frame :{idx}")           

        except Exception as e:
            print(f"[ERROR] Frame {idx} failed: {e}")
    
    print(f"totol time is : {time.time()- start_time}")

def encode_video(output_path, fps, width, height):
    out = av.open(output_path, mode='w')

    #stream = out.add_stream('h264_nvenc', rate=fps)
    stream = out.add_stream('libx264', rate=fps)
    stream.width = width
    stream.height = height
    stream.pix_fmt = 'yuv420p'

    stream.options = {'crf': '18', 'preset': 'veryfast'}
    #stream.options = {'rc': 'constqp', 'qp': '18','preset': 'p7','profile': 'high'}

    next_index = 0
    buffer = {}
    end_count = 0

    while True:
        idx, frame = processed_q.get()
        if idx is None:
            end_count += 1
            if end_count >= NUM_WORKERS:
                break
            continue
        buffer[idx] = frame

        # 只要有当前要写入的帧，就输出
        while next_index in buffer:
            frame = buffer.pop(next_index)
            video_frame = av.VideoFrame.from_ndarray(frame, format='rgb24')
            for packet in stream.encode(video_frame):
                out.mux(packet)
            next_index += 1

    # 清空编码缓存
    for packet in stream.encode():
        out.mux(packet)
    out.close()
# 启动任务
video_path = './temp/buhui.mp4'
output_path = './temp/buhui-wangrong-enhancer.mp4'
new_face_image_path = './temp/wangrong.jpg'
face_analyser = FaceAnalysis(name='buffalo_l', root='./', providers=['CUDAExecutionProvider'])
face_analyser.prepare(ctx_id=0, det_size=(640, 640))

faces = face_analyser.get(cv2.imread(new_face_image_path))
if faces:
    new_face = sorted(faces, key=lambda x: x.bbox[0])[0]
        
    #threading.Thread(target=decode_video, args=(video_path,)).start()
    #threading.Thread(target=process_frame_loop, args=(new_face,)).start()
    #threading.Thread(target=encode_video, args=(output_path, 25,576,1080)).start()
    threads = []
        
    # 解码线程
    t_decode = threading.Thread(target=decode_video,args=(video_path,))
    threads.append(t_decode)
    t_decode.start()
    
    # 多处理线程
    for i in range(NUM_WORKERS):
        t = threading.Thread(target=process_frame_loop, args=(new_face,),name=f"process_worker_{i}")
        threads.append(t)
        t.start()
    
    # 编码线程
    t_encode = threading.Thread(target=encode_video,args=(output_path, 25,576,1080))
    threads.append(t_encode)
    t_encode.start()
    
    for t in threads:
        t.join()
    
    print("Processing complete.")



 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()

 auto_media_publisher/modifiers/face_swap/face_swaper_ffmpeg_stream.py

File was renamed from auto_media_publisher/utils/face_swap/face_swaper_ffmpeg_stream.py
@@ -8,7 +8,6 @@
import torch
import json
import insightface
from insightface.app.common import Face
from .face_nsfw_predictor import NSFWPredictor
from .face_enhancer import FaceEnhancer
from .face_analyser import FaceAnalyser
@@ -41,7 +40,6 @@
        self.stop_event = threading.Event()

    def _select_device(self):
        return 'cpu'
        if torch.cuda.is_available():
            return 'cuda'
        elif torch.backends.mps.is_available():

 auto_media_publisher/modifiers/face_swap/typing.py

New file
@@ -0,0 +1,7 @@

from typing import Any
from insightface.app.common import Face
import numpy

Face = Face
Frame = numpy.ndarray[Any, Any]

 auto_media_publisher/utils/face_swap/face_swaper_cv2.py

File was deleted

 auto_media_publisher/utils/face_swap/typing.py

File was deleted

 auto_media_publisher/utils/media_input_output.py

New file
@@ -0,0 +1,93 @@
import av
import os
import pims
import numpy as np
from torch.utils.data import Dataset
from torchvision.transforms.functional import to_pil_image
from PIL import Image


class VideoReader(Dataset):
    def __init__(self, path, transform=None):
        self.video = pims.PyAVVideoReader(path)
        self.rate = self.video.frame_rate
        self.transform = transform
        
    @property
    def frame_rate(self):
        return self.rate
        
    def __len__(self):
        return len(self.video)
        
    def __getitem__(self, idx):
        frame = self.video[idx]
        frame = Image.fromarray(np.asarray(frame))
        if self.transform is not None:
            frame = self.transform(frame)
        return frame


class VideoWriter:
    def __init__(self, path, frame_rate, bit_rate=1000000):
        self.container = av.open(path, mode='w')
        self.stream = self.container.add_stream('h264', rate=round(frame_rate))
        self.stream.pix_fmt = 'yuv420p'
        self.stream.bit_rate = bit_rate
        self.stream.codec_context.options = {
            'crf': '18',  # 控制画质，越小越清晰（18是高质量）
            'preset': 'veryfast',     # 影响编码速度，medium 是1x，fast是2x faster是3x veryfast是4x superfast是6x ultrafast是10x，文件体积大约会增加(以medium为基准) 10%，15%，20% 30% 50%
            'tune': 'film',
        }
    
    def write(self, frames):
        # frames: [T, C, H, W]
        self.stream.width = frames.size(3)
        self.stream.height = frames.size(2)
        if frames.size(1) == 1:
            frames = frames.repeat(1, 3, 1, 1) # convert grayscale to RGB
        frames = frames.mul(255).byte().cpu().permute(0, 2, 3, 1).numpy()
        for t in range(frames.shape[0]):
            frame = frames[t]
            frame = av.VideoFrame.from_ndarray(frame, format='rgb24')
            self.container.mux(self.stream.encode(frame))
                
    def close(self):
        self.container.mux(self.stream.encode())
        self.container.close()


class ImageSequenceReader(Dataset):
    def __init__(self, path, transform=None):
        self.path = path
        self.files = sorted(os.listdir(path))
        self.transform = transform
        
    def __len__(self):
        return len(self.files)
    
    def __getitem__(self, idx):
        with Image.open(os.path.join(self.path, self.files[idx])) as img:
            img.load()
        if self.transform is not None:
            return self.transform(img)
        return img


class ImageSequenceWriter:
    def __init__(self, path, extension='jpg'):
        self.path = path
        self.extension = extension
        self.counter = 0
        os.makedirs(path, exist_ok=True)
    
    def write(self, frames):
        # frames: [T, C, H, W]
        for t in range(frames.shape[0]):
            to_pil_image(frames[t]).save(os.path.join(
                self.path, str(self.counter).zfill(4) + '.' + self.extension))
            self.counter += 1
            
    def close(self):
        pass
        

 data/auto_media_publisher.sqlite3

Binary files differ