import shutil import cv2 import numpy as np import torch from .ffmpeg_utils import mux_video_audio from .file_utils import create_task_dir from .media_utils import ( TARGET_AUDIO_SR, TARGET_FPS, denoise_audio_file, encode_rgb_frames_to_mp4, frames_uint8_rgb_to_image_tensor, images_to_uint8_rgb, resample_frames_to_fps, save_audio_to_wav, trim_or_pad_audio, ) def _log(message): print(f"[HighSync] {message}", flush=True) def _maybe_throw_if_interrupted(): try: import comfy.model_management as model_management except Exception: return model_management.throw_exception_if_processing_interrupted() def _empty_cache(device): if isinstance(device, torch.device) and device.type == "cuda" and torch.cuda.is_available(): torch.cuda.empty_cache() try: import comfy.model_management as model_management model_management.soft_empty_cache() except Exception: return def create_mask(mask_path): raw_mask = cv2.imread(str(mask_path)) if raw_mask is None: raise FileNotFoundError(f"HighSync mask could not be read: {mask_path}") mask = 1 - (raw_mask / 255.0) h, w, c = mask.shape new_mask = np.zeros((h, w, c + 1), dtype=np.float32) new_mask[:, :, :c] = mask new_mask[:, :, c] = mask[:, :, c - 1] return torch.FloatTensor(new_mask).permute(2, 0, 1).unsqueeze(1) def create_mask_alpha(width, height, coordinate): left, top, right, bottom = coordinate mask = np.ones((height, width, 3), dtype=np.float32) mask[top:bottom, left:right] = 0 pad = min(int((right - left) / 15), int((bottom - top) / 15)) if pad <= 0: return mask for idx in range(left, right): mask[bottom - pad:bottom, idx] = np.array([np.linspace(0, 1, pad)] * 3).transpose(1, 0) mask[top:top + pad, idx] = np.array([np.linspace(1, 0, pad)] * 3).transpose(1, 0) for idx in range(top, bottom): mask[idx, right - pad:right] = np.array([np.linspace(0, 1, pad)] * 3).transpose(1, 0) mask[idx, left:left + pad] = np.array([np.linspace(1, 0, pad)] * 3).transpose(1, 0) for idx, start, end in zip(range(top, top + pad), np.linspace(1, 0.7, pad), np.linspace(0.7, 0, pad)): mask[idx, left:left + pad] = np.array([np.linspace(start, end, pad)] * 3).transpose(1, 0) mask[idx, right - pad:right] = np.array([np.linspace(end, start, pad)] * 3).transpose(1, 0) for idx, start, end in zip(range(bottom - pad, bottom), np.linspace(0.7, 1, pad), np.linspace(0, 0.7, pad)): mask[idx, left:left + pad] = np.array([np.linspace(start, end, pad)] * 3).transpose(1, 0) mask[idx, right - pad:right] = np.array([np.linspace(end, start, pad)] * 3).transpose(1, 0) return mask def _prepare_audio_files(task, audio, duration_sec, trim_audio_to_video, denoise_audio, device, model_cache): audio_for_video = trim_or_pad_audio( audio, duration_sec=duration_sec if trim_audio_to_video else None, pad_to_duration=bool(trim_audio_to_video), ) mux_audio_path = save_audio_to_wav( audio_for_video, task["media_dir"] / "mux_audio.wav", ) feature_audio_path = save_audio_to_wav( audio_for_video, task["media_dir"] / "feature_audio_16k.wav", mono=True, sample_rate=TARGET_AUDIO_SR, ) if denoise_audio: feature_audio_path = denoise_audio_file( feature_audio_path, task["media_dir"] / "feature_audio_16k_denoised.wav", device=device, cache=model_cache, ) return audio_for_video, mux_audio_path, feature_audio_path def _ensure_highsync_model(highsync_model): required = ("config", "pipeline", "audio_processor", "image_processor_cls", "device", "weight_dtype") if not isinstance(highsync_model, dict) or any(key not in highsync_model for key in required): raise ValueError("highsync_model must come from the HighSync Model Loader node.") def _append_original_frame(output_frames_rgb, frame_bgr): output_frames_rgb.append(cv2.cvtColor(frame_bgr, cv2.COLOR_BGR2RGB)) def _run_generation_chunk( pipeline, image_processor, config, audio_fea_final, frames_video, output_frames_rgb, count_img, generator, device, weight_dtype, inference_steps, cfg_scale, ): chunk_frame_count = len(frames_video) chunk_global_start = count_img chunk_global_end = count_img + chunk_frame_count - 1 _log( f"Stage 4/6 face tracking: frames {chunk_global_start}-{chunk_global_end} " f"({chunk_frame_count} frame(s))." ) source_image_pixels_all, bbox_all, frame_list_all = image_processor.preprocess_frames_jump(frames_video) img_size = (int(config.data.source_image.width), int(config.data.source_image.height)) sample_frames = int(config.data.n_sample_frames) mask_path = str(config.mask_path) segment_count = len(frame_list_all) for segment_index, (source_image_pixels_list, bbox_list, frame_list) in enumerate( zip(source_image_pixels_all, bbox_all, frame_list_all), start=1, ): _maybe_throw_if_interrupted() if not frame_list: continue segment_start = count_img segment_end = count_img + len(frame_list) - 1 if bbox_list[0] != (0, 0, 0, 0): _log( f"Stage 4/6 inference segment {segment_index}/{segment_count}: " f"global frames {segment_start}-{segment_end} ({len(frame_list)} frame(s)), " f"{inference_steps} diffusion step(s)." ) audio_tensor = audio_fea_final[:, count_img: count_img + len(frame_list)] main_length = len(frame_list) diff = sample_frames - main_length if diff > 0: _log( f"Stage 4/6 segment {segment_index}/{segment_count}: padded " f"{main_length} -> {sample_frames} frame(s) for the HighSync window." ) source_image_pixels_list.extend([source_image_pixels_list[-1]] * diff) audio_tensor = torch.cat((audio_tensor, audio_tensor[:, -1:].repeat(1, diff, 1, 1)), dim=1) source_image_pixels = torch.cat(source_image_pixels_list, dim=0) mask_main = create_mask(mask_path).unsqueeze(0).to(device=device, dtype=weight_dtype) _empty_cache(device) with torch.no_grad(): pipeline_output = pipeline( gt_encode=source_image_pixels, audio_tensor=audio_tensor, mask_main=mask_main, width=img_size[0], height=img_size[1], video_length=len(source_image_pixels_list), num_inference_steps=int(inference_steps), guidance_scale=float(cfg_scale), generator=generator, progress_label=f"HighSync seg {segment_index}/{segment_count} frames {segment_start}-{segment_end}", ) _empty_cache(device) _log( f"Stage 4/6 blend segment {segment_index}/{segment_count}: " f"writing {main_length} generated frame(s) back to original resolution." ) images_out = pipeline_output.videos.squeeze(0).permute(1, 2, 3, 0).cpu().numpy()[:main_length] images_out = np.clip(images_out * 255, 0, 255).astype(np.uint8) for image_out, bbox, frame in zip(images_out, bbox_list, frame_list): frame_h, frame_w, _ = frame.shape if bbox == (0, 0, 0, 0): _append_original_frame(output_frames_rgb, frame) else: image_out_bgr = cv2.cvtColor(image_out, cv2.COLOR_RGB2BGR) frame_in_out = frame.copy() frame_in_out[bbox[1]:bbox[3], bbox[0]:bbox[2]] = cv2.resize( image_out_bgr, (bbox[2] - bbox[0], bbox[3] - bbox[1]), ) mask = create_mask_alpha(frame_w, frame_h, bbox) blended = (frame * mask + frame_in_out * (1 - mask)).astype(np.uint8) _append_original_frame(output_frames_rgb, blended) count_img += 1 else: _log( f"Stage 4/6 segment {segment_index}/{segment_count}: " f"no face box, copying frames {segment_start}-{segment_end} unchanged." ) for frame in frame_list: _append_original_frame(output_frames_rgb, frame) count_img += 1 return count_img def run_highsync( highsync_model, images, audio, input_fps=25, max_frames=0, inference_steps=None, cfg_scale=None, seed=42, preprocess_chunk_frames=1500, denoise_audio=True, trim_audio_to_video=True, keep_intermediate_outputs=True, ): _ensure_highsync_model(highsync_model) config = highsync_model["config"] pipeline = highsync_model["pipeline"] audio_processor = highsync_model["audio_processor"] image_processor = highsync_model["image_processor_cls"]((int(config.data.source_image.width), int(config.data.source_image.height))) device = torch.device(highsync_model["device"]) weight_dtype = highsync_model["weight_dtype"] inference_steps = int(inference_steps if inference_steps is not None else config.inference_steps) cfg_scale = float(cfg_scale if cfg_scale is not None else config.cfg_scale) preprocess_chunk_frames = max(12, int(preprocess_chunk_frames)) task = create_task_dir() _log( f"Stage 1/6 start: input_fps={input_fps}, target_fps={TARGET_FPS}, " f"max_frames={max_frames}, inference_steps={inference_steps}, cfg_scale={cfg_scale}, " f"task_dir={task['task_dir']}" ) frames_rgb = images_to_uint8_rgb(images) input_frame_count = len(frames_rgb) frames_rgb = resample_frames_to_fps(frames_rgb, input_fps=input_fps, output_fps=TARGET_FPS, max_frames=max_frames) frames_bgr = [cv2.cvtColor(frame, cv2.COLOR_RGB2BGR) for frame in frames_rgb] initial_duration_sec = len(frames_bgr) / float(TARGET_FPS) _log( f"Stage 1/6 media prepared: {input_frame_count} input frame(s) -> " f"{len(frames_bgr)} frame(s) at {TARGET_FPS} FPS " f"({initial_duration_sec:.2f}s)." ) _log( f"Stage 2/6 audio features: converting audio to mono {TARGET_AUDIO_SR} Hz" f"{' and denoising' if denoise_audio else ''}." ) _, _, feature_audio_path = _prepare_audio_files( task=task, audio=audio, duration_sec=initial_duration_sec, trim_audio_to_video=trim_audio_to_video, denoise_audio=denoise_audio, device=device, model_cache=highsync_model, ) _log("Stage 2/6 audio features: running Whisper audio feature extraction.") whisper_feature = audio_processor.audio2feat(str(feature_audio_path)) whisper_chunks = audio_processor.feature2chunks(feature_array=whisper_feature, fps=TARGET_FPS) audio_frame_num = int(whisper_chunks.shape[0]) video_length = min(len(frames_bgr), audio_frame_num) if video_length <= 0: raise RuntimeError("HighSync could not align any video frames with the provided audio.") frames_bgr = frames_bgr[:video_length] audio_fea_final = torch.as_tensor(whisper_chunks[:video_length], dtype=pipeline.vae.dtype, device=pipeline.vae.device) audio_fea_final = audio_fea_final.unsqueeze(0) _log( f"Stage 3/6 alignment: video has {len(frames_bgr)} frame(s), " f"audio features cover {audio_frame_num} frame(s), processing {video_length} frame(s)." ) final_duration_sec = video_length / float(TARGET_FPS) output_audio, mux_audio_path, _ = _prepare_audio_files( task=task, audio=audio, duration_sec=final_duration_sec, trim_audio_to_video=trim_audio_to_video, denoise_audio=False, device=device, model_cache=highsync_model, ) torch.manual_seed(int(seed)) generator = torch.manual_seed(int(seed)) output_frames_rgb = [] frames_video = [] count_motion = 0 count_img = 0 while True: _maybe_throw_if_interrupted() if count_motion >= video_length: finish = True else: frames_video.append(frames_bgr[count_motion]) count_motion += 1 finish = False if len(frames_video) == 0: break if (count_motion % preprocess_chunk_frames == 0) or finish: _log(f"Stage 4/6 chunk queued: frames {count_img}-{count_motion - 1} / {video_length}.") count_img = _run_generation_chunk( pipeline=pipeline, image_processor=image_processor, config=config, audio_fea_final=audio_fea_final, frames_video=frames_video, output_frames_rgb=output_frames_rgb, count_img=count_img, generator=generator, device=device, weight_dtype=weight_dtype, inference_steps=inference_steps, cfg_scale=cfg_scale, ) frames_video = [] if finish: break if len(output_frames_rgb) == 0: raise RuntimeError("HighSync finished without producing output frames.") output_frames = np.ascontiguousarray(np.stack(output_frames_rgb)) _log(f"Stage 5/6 encoding: writing {len(output_frames)} frame(s) to MP4.") no_audio_path = encode_rgb_frames_to_mp4(output_frames, task["media_dir"] / "highsync_no_audio.mp4", fps=TARGET_FPS) final_video_path = mux_video_audio(no_audio_path, mux_audio_path, task["task_dir"] / "final.mp4") output_images = frames_uint8_rgb_to_image_tensor(output_frames) if not keep_intermediate_outputs: for path in (task["media_dir"], task["frames_dir"], task["logs_dir"]): shutil.rmtree(path, ignore_errors=True) _log(f"Stage 6/6 done: final video saved to {final_video_path.resolve()}") return output_images, output_audio, TARGET_FPS, str(final_video_path.resolve())