#!/usr/bin/env python3 """Unified InstructAV2AV inference: source AV + instruction -> edited AV.""" from __future__ import annotations import argparse import logging import sys from pathlib import Path from typing import Any import pandas as pd import torch from omegaconf import OmegaConf from tqdm import tqdm from ovi.distributed_comms.parallel_states import ( get_sequence_parallel_state, initialize_sequence_parallel_state, nccl_info, ) from ovi.distributed_comms.util import get_global_rank, get_local_rank, get_world_size from ovi.ovi_fusion_engine import OviFusionEngine from ovi.utils.av_edit_data import ( get_instruction, get_video_info, load_audio_array, load_manifest, load_video_array, resolve_media_path, save_audio, snap_num_frames, to_audio_tensor, to_video_tensor, ) from ovi.utils.io_utils import save_video DEFAULT_CONFIG = "ovi/configs/inference/inference_av_edit.yaml" def parse_args() -> argparse.Namespace: parser = argparse.ArgumentParser(description=__doc__) parser.add_argument("--config-file", default=str(DEFAULT_CONFIG)) inputs = parser.add_mutually_exclusive_group(required=True) inputs.add_argument("--source-video", help="Source video for a single edit.") inputs.add_argument("--manifest", help="CSV/JSON/JSONL batch manifest.") parser.add_argument("--source-audio", help="Optional source audio for a single edit.") parser.add_argument("--instruction", help="Editing instruction for a single edit.") parser.add_argument( "--instruction-column", help="Instruction column for batch inference. Default: instruction.", ) parser.add_argument("--output", help="Single-edit output MP4 path.") parser.add_argument("--output-dir", help="Batch output directory or single-edit default directory.") parser.add_argument("--finetune-path", help="Editing checkpoint override.") parser.add_argument("--ckpt-dir", help="Base Ovi checkpoint directory override.") parser.add_argument("--seed", type=int) parser.add_argument("--sample-steps", type=int) parser.add_argument("--num-frames", type=int) parser.add_argument("--video-guidance-scale", type=float) parser.add_argument("--audio-guidance-scale", type=float) return parser.parse_args() def configure_logging(rank: int) -> None: logging.basicConfig( level=logging.INFO if rank == 0 else logging.ERROR, format="[%(asctime)s] %(levelname)s: %(message)s", handlers=[logging.StreamHandler(stream=sys.stdout)], ) def apply_overrides(config, args: argparse.Namespace) -> None: overrides = { "finetune_path": args.finetune_path, "ckpt_dir": args.ckpt_dir, "output_dir": args.output_dir, "seed": args.seed, "sample_steps": args.sample_steps, "num_frames": args.num_frames, "video_guidance_scale": args.video_guidance_scale, "audio_guidance_scale": args.audio_guidance_scale, } for key, value in overrides.items(): if value is not None: config[key] = value config.av2av_edit = True config.has_video = True config.has_audio = True config.mode = "t2v" def build_rows(args: argparse.Namespace, config) -> tuple[list[dict[str, Any]], Path | None]: if args.source_video: if not args.instruction or not args.instruction.strip(): raise ValueError("--instruction is required with --source-video.") row = { "source_video": args.source_video, "source_audio": args.source_audio, "instruction": args.instruction.strip(), } return [row], None manifest_path = Path(args.manifest).expanduser().resolve() rows = load_manifest(manifest_path) instruction_column = args.instruction_column or config.get( "instruction_column", "instruction" ) for row in rows: row["instruction"] = get_instruction(row, instruction_column) return rows, manifest_path def distributed_layout(world_size: int, global_rank: int) -> tuple[int, int, int]: if get_sequence_parallel_state(): sp_size = nccl_info.sp_size sp_rank = nccl_info.rank_within_group group_id = global_rank // sp_size group_count = world_size // sp_size return sp_rank, group_id, group_count return 0, global_rank, world_size def output_paths( args: argparse.Namespace, config, row: dict[str, Any], row_index: int, ) -> tuple[Path, Path]: if args.source_video and args.output: video_path = Path(args.output).expanduser().resolve() else: output_dir = Path(config.get("output_dir", "./outputs/av_edits")).expanduser().resolve() source_stem = Path(str(row["source_video"])).stem video_path = output_dir / f"{row_index:06d}_{source_stem}_edited.mp4" audio_path = video_path.with_suffix(".wav") video_path.parent.mkdir(parents=True, exist_ok=True) return video_path, audio_path def prepare_inputs(row: dict[str, Any], manifest_path: Path | None, config, engine): pseudo_manifest = manifest_path or (Path.cwd() / "single_input.csv") source_video = resolve_media_path( row.get("source_video"), pseudo_manifest, required=True ) source_audio = resolve_media_path( row.get("source_audio"), pseudo_manifest, required=False ) fps, total_frames = get_video_info(source_video) configured_frames = config.get("num_frames", None) available_frames = total_frames if configured_frames is not None: available_frames = min(available_frames, int(configured_frames)) num_frames = snap_num_frames(available_frames) frame_size = config.get("video_frame_height_width", [704, 1280]) video, _ = load_video_array( source_video, num_frames=num_frames, height=int(frame_size[0]), width=int(frame_size[1]), max_pixels=int(frame_size[0]) * int(frame_size[1]), ) sample_rate = int(config.get("audio_sample_rate", 16000)) num_audio_samples = max(1, round(num_frames / fps * sample_rate)) audio = load_audio_array( source_audio or source_video, sample_rate=sample_rate, num_samples=num_audio_samples, ) return { "source_video": source_video, "source_audio": source_audio or source_video, "video_array": video, "audio_array": audio, "video_tensor": to_video_tensor(video, engine.device, engine.target_dtype), "audio_tensor": to_audio_tensor(audio, engine.device), "fps": fps, "num_frames": num_frames, "sample_rate": sample_rate, } def main() -> None: args = parse_args() config = OmegaConf.load(args.config_file) apply_overrides(config, args) world_size = get_world_size() global_rank = get_global_rank() local_rank = get_local_rank() configure_logging(global_rank) if not torch.cuda.is_available(): raise RuntimeError("Inference requires a CUDA device.") torch.cuda.set_device(local_rank) sp_size = int(config.get("sp_size", 1)) if sp_size > world_size or world_size % sp_size: raise ValueError("sp_size must divide the torchrun world size.") if world_size > 1: torch.distributed.init_process_group(backend="nccl", init_method="env://") elif sp_size != 1: raise ValueError("sp_size must be 1 for a single process.") initialize_sequence_parallel_state(sp_size) rows, manifest_path = build_rows(args, config) if not rows: raise ValueError("No inference rows were found.") sp_rank, group_id, group_count = distributed_layout(world_size, global_rank) assigned_indices = list(range(len(rows)))[group_id::group_count] logging.info("Loading the AV editing checkpoint: %s", config.get("finetune_path")) engine = OviFusionEngine(config=config, device=local_rank, target_dtype=torch.bfloat16) engine.eval() result_rows = [] base_seed = int(config.get("seed", 103)) for row_index in tqdm(assigned_indices, disable=sp_rank != 0, desc="Editing"): row = rows[row_index] prepared = prepare_inputs(row, manifest_path, config, engine) sample_seed = base_seed + row_index torch.manual_seed(sample_seed) torch.cuda.manual_seed_all(sample_seed) generated = engine.generate( text_prompt=row["instruction"], image_path=None, video_frame_height_width=list(config.get("video_frame_height_width", [704, 1280])), seed=sample_seed, solver_name=str(config.get("solver_name", "unipc")), sample_steps=int(config.get("sample_steps", 50)), shift=float(config.get("shift", 5.0)), video_guidance_scale=float(config.get("video_guidance_scale", 4.0)), audio_guidance_scale=float(config.get("audio_guidance_scale", 3.0)), slg_layer=int(config.get("slg_layer", 11)), video_negative_prompt=str(config.get("video_negative_prompt", "")), audio_negative_prompt=str(config.get("audio_negative_prompt", "")), input_video=prepared["video_tensor"], input_audio=prepared["audio_tensor"], ) if generated is None: raise RuntimeError(f"Generation failed for row {row_index}.") generated_video, generated_audio, _ = generated if generated_video is None or generated_audio is None: raise RuntimeError(f"Generation returned empty AV for row {row_index}.") if sp_rank == 0: edited_video_path, edited_audio_path = output_paths(args, config, row, row_index) save_video( str(edited_video_path), generated_video, generated_audio, sample_rate=prepared["sample_rate"], fps=prepared["fps"], ) save_audio(edited_audio_path, generated_audio, prepared["sample_rate"]) result_rows.append( { "row_index": row_index, "source_video": str(prepared["source_video"]), "source_audio": str(prepared["source_audio"]), "instruction": row["instruction"], "seed": sample_seed, "num_frames": prepared["num_frames"], "fps": prepared["fps"], "edited_video": str(edited_video_path), "edited_audio": str(edited_audio_path), } ) if world_size > 1: gathered = [None for _ in range(world_size)] torch.distributed.all_gather_object(gathered, result_rows) else: gathered = [result_rows] if global_rank == 0: merged = [row for rows_from_rank in gathered for row in (rows_from_rank or [])] merged.sort(key=lambda item: item["row_index"]) output_dir = Path(config.get("output_dir", "./outputs/av_edits")).expanduser().resolve() output_dir.mkdir(parents=True, exist_ok=True) result_manifest = output_dir / "results.csv" pd.DataFrame(merged).to_csv(result_manifest, index=False) OmegaConf.save(config, output_dir / "config_resolved.yaml", resolve=True) logging.info("Saved %d result(s). Manifest: %s", len(merged), result_manifest) if world_size > 1: torch.distributed.barrier() if __name__ == "__main__": main()