scripts/edit.py

#!/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()