SeedVR2-3B-test

Running on Zero

Add resumable ZeroGPU video segments
7e0b9c2 verified 2 months ago
39.9 kB
# // Copyright (c) 2025 Bytedance Ltd. and/or its affiliates
# //
# // Licensed under the Apache License, Version 2.0 (the "License");
# // you may not use this file except in compliance with the License.
# // You may obtain a copy of the License at
# //
# //     http://www.apache.org/licenses/LICENSE-2.0
# //
# // Unless required by applicable law or agreed to in writing, software
# // distributed under the License is distributed on an "AS IS" BASIS,
# // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# // See the License for the specific language governing permissions and
# // limitations under the License.
import spaces
import subprocess

import json
import os
import torch
import mediapy
from einops import rearrange
from omegaconf import OmegaConf
print(os.getcwd())
import datetime
from tqdm import tqdm
import gc
import shutil
import tempfile

from data.image.transforms.divisible_crop import DivisibleCrop
from data.image.transforms.na_resize import NaResize
from data.video.transforms.rearrange import Rearrange
if os.path.exists("./projects/video_diffusion_sr/color_fix.py"):
    from projects.video_diffusion_sr.color_fix import wavelet_reconstruction
    use_colorfix=True
else:
    use_colorfix = False
    print('Note!!!!!! Color fix is not avaliable!')
from torchvision.transforms import Compose, Lambda, Normalize
from torchvision.io.video import read_video
import argparse
from PIL import Image

from common.distributed import (
    get_device,
    init_torch,
)

from common.distributed.advanced import (
    get_data_parallel_rank,
    get_data_parallel_world_size,
    get_sequence_parallel_rank,
    get_sequence_parallel_world_size,
    init_sequence_parallel,
)

from projects.video_diffusion_sr.infer import VideoDiffusionInfer
from common.config import load_config
from common.distributed.ops import sync_data
from common.seed import set_seed

import gradio as gr
from fractions import Fraction
from huggingface_hub import hf_hub_download
from pathlib import Path
import uuid
import mimetypes
import torchvision.transforms as T

os.environ["MASTER_ADDR"] = "127.0.0.1"
os.environ["MASTER_PORT"] = "12355"
os.environ["RANK"] = str(0)
os.environ["WORLD_SIZE"] = str(1)

subprocess.run(
    "pip install flash-attn --no-build-isolation",
    env={"FLASH_ATTENTION_SKIP_CUDA_BUILD": "TRUE"},
    shell=True,
)

MODEL_REPO_ID = "ByteDance-Seed/SeedVR2-3B"
MODEL_REVISION = "37255ff8cccfb01071b87f635a5948ca8d53117c"
MODEL_FILENAMES = {
    "vae": "ema_vae.pth",
    "dit": "seedvr2_ema_3b.pth",
    "pos_emb": "pos_emb.pt",
    "neg_emb": "neg_emb.pt",
}

# README preload_from_hub places these files in this cache during the Space build.
MODEL_PATHS = {
    name: hf_hub_download(
        repo_id=MODEL_REPO_ID,
        filename=filename,
        revision=MODEL_REVISION,
    )
    for name, filename in MODEL_FILENAMES.items()
}


def download_demo_video(url, destination):
    if not os.path.exists(destination):
        torch.hub.download_url_to_file(url, destination)


download_demo_video(
    "https://huggingface.co/datasets/Iceclear/SeedVR_VideoDemos/resolve/main/seedvr_videos_crf23/aigc1k/23_1_lq.mp4",
    "01.mp4",
)
download_demo_video(
    "https://huggingface.co/datasets/Iceclear/SeedVR_VideoDemos/resolve/main/seedvr_videos_crf23/aigc1k/28_1_lq.mp4",
    "02.mp4",
)
download_demo_video(
    "https://huggingface.co/datasets/Iceclear/SeedVR_VideoDemos/resolve/main/seedvr_videos_crf23/aigc1k/2_1_lq.mp4",
    "03.mp4",
)

CHUNK_FRAME_LIMIT = 121
CHUNK_OVERLAP_FRAMES = 2
MODEL_SIZE_MULTIPLE = 16
VIDEO_OUTPUT_RESOLUTIONS = {
    "720p": (1280, 720),
    "1080p": (1920, 1080),
}
SEGMENT_DURATION_CHOICES = {
    "About 5 seconds (ZeroGPU-safe)": 5,
    "About 5 minutes (Dedicated GPU)": 300,
}
SEGMENT_JOB_ROOT = Path("output") / "segment_jobs"
# ZeroGPU validates the requested duration before a function starts. Keep one
# realistic reservation around the public request; nested GPU decorators would
# otherwise reserve additional quota while a worker is already allocated.
GPU_LOOP_DURATION = 180

def configure_sequence_parallel(sp_size):
    if sp_size > 1:
        init_sequence_parallel(sp_size)

def run_media_command(command, error_message):
    result = subprocess.run(command, capture_output=True, text=True)
    if result.returncode != 0:
        details = result.stderr.strip() or result.stdout.strip() or "unknown error"
        raise RuntimeError(f"{error_message}\n{details}")
    return result


def safe_remove(path):
    if path and os.path.exists(path):
        os.remove(path)


def build_video_transform(target_pixels):
    return Compose(
        [
            NaResize(
                resolution=target_pixels ** 0.5,
                mode="area",
                downsample_only=False,
            ),
            Lambda(lambda x: torch.clamp(x, 0.0, 1.0)),
            DivisibleCrop((MODEL_SIZE_MULTIPLE, MODEL_SIZE_MULTIPLE)),
            Normalize(0.5, 0.5),
            Rearrange("t c h w -> c t h w"),
        ]
    )


def build_image_transform():
    return Compose(
        [
            NaResize(
                resolution=(2560 * 1440) ** 0.5,
                mode="area",
                downsample_only=False,
            ),
            Lambda(lambda x: torch.clamp(x, 0.0, 1.0)),
            DivisibleCrop((MODEL_SIZE_MULTIPLE, MODEL_SIZE_MULTIPLE)),
            Normalize(0.5, 0.5),
            Rearrange("t c h w -> c t h w"),
        ]
    )


def get_text_embeds():
    return {
        "texts_pos": [torch.load(MODEL_PATHS["pos_emb"], map_location="cpu")],
        "texts_neg": [torch.load(MODEL_PATHS["neg_emb"], map_location="cpu")],
    }


def probe_video_metadata(video_path):
    def _safe_float(value):
        try:
            return float(value)
        except (TypeError, ValueError):
            return 0.0

    command = [
        "ffprobe",
        "-v",
        "error",
        "-show_entries",
        "stream=index,codec_type,width,height,avg_frame_rate,r_frame_rate,nb_frames,duration",
        "-show_entries",
        "format=duration",
        "-of",
        "json",
        video_path,
    ]
    result = run_media_command(command, f"Could not read video metadata from {video_path}.")
    data = json.loads(result.stdout)
    streams = data.get("streams", [])
    video_stream = next((stream for stream in streams if stream.get("codec_type") == "video"), None)
    if video_stream is None:
        raise RuntimeError("The uploaded file does not contain a readable video stream.")

    fps_value = video_stream.get("avg_frame_rate") or video_stream.get("r_frame_rate") or "0/1"
    try:
        fps_fraction = Fraction(fps_value)
    except (ValueError, ZeroDivisionError):
        fps_fraction = Fraction(0, 1)
    fps = float(fps_fraction) if fps_fraction.numerator else 0.0
    duration = _safe_float(
        video_stream.get("duration")
        or data.get("format", {}).get("duration")
        or 0.0
    )
    nb_frames = video_stream.get("nb_frames")
    if nb_frames and str(nb_frames).isdigit():
        total_frames = int(nb_frames)
    elif fps > 0 and duration > 0:
        total_frames = max(1, int(round(duration * fps)))
    else:
        raise RuntimeError("Could not determine the source video length automatically.")

    return {
        "duration": duration if duration > 0 else (total_frames / fps if fps > 0 else 0.0),
        "fps": fps,
        "fps_display": str(fps_fraction) if fps_fraction.numerator else str(fps),
        "total_frames": total_frames,
        "width": int(video_stream.get("width") or 0),
        "height": int(video_stream.get("height") or 0),
        "has_audio": any(stream.get("codec_type") == "audio" for stream in streams),
    }


def build_chunk_plan(total_frames, chunk_frames=CHUNK_FRAME_LIMIT, overlap_frames=CHUNK_OVERLAP_FRAMES):
    if total_frames <= 0:
        raise RuntimeError("The source video does not contain any frames to process.")
    if chunk_frames <= overlap_frames:
        raise ValueError("Chunk size must be larger than the overlap.")

    stride = chunk_frames - overlap_frames
    chunks = []
    start_frame = 0
    while start_frame < total_frames:
        end_frame = min(start_frame + chunk_frames, total_frames)
        chunks.append(
            {
                "start_frame": start_frame,
                "end_frame": end_frame,
                "trim_leading_frames": 0 if not chunks else overlap_frames,
            }
        )
        if end_frame >= total_frames:
            break
        start_frame += stride
    return chunks


def build_output_segment_plan(chunk_plan, fps, segment_seconds):
    """Group existing model chunks without changing their overlap behavior."""
    if fps <= 0:
        raise RuntimeError("Could not determine a valid FPS for video segmentation.")

    segment_frames = max(1, int(round(fps * segment_seconds)))
    segments = []
    current_segment = []
    current_frames = 0
    for chunk in chunk_plan:
        output_frames = chunk["end_frame"] - chunk["start_frame"] - chunk["trim_leading_frames"]
        if current_segment and current_frames + output_frames > segment_frames:
            segments.append(current_segment)
            current_segment = []
            current_frames = 0
        current_segment.append(chunk)
        current_frames += output_frames
    if current_segment:
        segments.append(current_segment)
    return segments


def get_job_directory(job_state):
    if not isinstance(job_state, dict) or not job_state.get("job_dir"):
        return None

    job_dir = Path(job_state["job_dir"]).resolve()
    job_root = SEGMENT_JOB_ROOT.resolve()
    if job_dir.parent != job_root:
        raise RuntimeError("The saved video job is not valid anymore. Please start a new upload.")
    return job_dir


def create_segment_job(video_path, metadata, output_resolution, segment_duration_label):
    segment_seconds = SEGMENT_DURATION_CHOICES.get(
        segment_duration_label,
        SEGMENT_DURATION_CHOICES["About 5 seconds (ZeroGPU-safe)"],
    )
    extension = Path(video_path).suffix or ".mp4"
    job_dir = (SEGMENT_JOB_ROOT / uuid.uuid4().hex).resolve()
    job_dir.mkdir(parents=True, exist_ok=False)
    source_path = job_dir / f"source{extension}"
    shutil.copy2(video_path, source_path)

    return {
        "job_dir": str(job_dir),
        "source_name": source_path.name,
        "metadata": metadata,
        "output_resolution": output_resolution,
        "segment_seconds": segment_seconds,
        "next_segment": 0,
        "outputs": [],
    }


def remove_segment_job(job_state):
    job_dir = get_job_directory(job_state)
    if job_dir and job_dir.exists():
        shutil.rmtree(job_dir, ignore_errors=True)


def extract_video_chunk(source_path, chunk_path, start_frame, end_frame):
    filter_graph = f"trim=start_frame={start_frame}:end_frame={end_frame},setpts=PTS-STARTPTS"
    command = [
        "ffmpeg",
        "-y",
        "-v",
        "error",
        "-i",
        source_path,
        "-an",
        "-vf",
        filter_graph,
        "-c:v",
        "ffv1",
        chunk_path,
    ]
    run_media_command(
        command,
        f"Could not extract frames {start_frame} to {end_frame} from the source video.",
    )


def write_concat_manifest(manifest_path, file_paths):
    with open(manifest_path, "w", encoding="utf-8") as manifest:
        for file_path in file_paths:
            escaped_path = Path(file_path).as_posix().replace("'", r"'\''")
            manifest.write(f"file '{escaped_path}'\n")


def concatenate_video_chunks(chunk_paths, output_path, temp_dir, fps_out):
    if not chunk_paths:
        raise RuntimeError("No processed video chunks were generated.")
    if len(chunk_paths) == 1:
        shutil.copyfile(chunk_paths[0], output_path)
        return

    manifest_path = os.path.join(temp_dir, "concat_manifest.txt")
    write_concat_manifest(manifest_path, chunk_paths)

    try:
        run_media_command(
            [
                "ffmpeg",
                "-y",
                "-v",
                "error",
                "-f",
                "concat",
                "-safe",
                "0",
                "-i",
                manifest_path,
                "-c",
                "copy",
                output_path,
            ],
            "Could not concatenate processed video chunks losslessly.",
        )
    except RuntimeError:
        run_media_command(
            [
                "ffmpeg",
                "-y",
                "-v",
                "error",
                "-f",
                "concat",
                "-safe",
                "0",
                "-i",
                manifest_path,
                "-r",
                str(fps_out),
                "-c:v",
                "libx264",
                "-preset",
                "medium",
                "-crf",
                "14",
                "-pix_fmt",
                "yuv420p",
                output_path,
            ],
            "Could not concatenate processed video chunks.",
        )


def finalize_video_output(silent_video_path, source_video_path, output_path, target_width, target_height, has_audio):
    video_filter = (
        f"scale={target_width}:{target_height}:force_original_aspect_ratio=decrease,"
        f"pad={target_width}:{target_height}:(ow-iw)/2:(oh-ih)/2:black,"
        "setsar=1,format=yuv420p"
    )

    base_command = [
        "ffmpeg",
        "-y",
        "-v",
        "error",
        "-i",
        silent_video_path,
    ]
    if has_audio:
        base_command.extend(["-i", source_video_path])

    copy_audio_command = base_command + [
        "-map",
        "0:v:0",
    ]
    if has_audio:
        copy_audio_command.extend(["-map", "1:a:0"])
    copy_audio_command.extend(
        [
            "-vf",
            video_filter,
            "-c:v",
            "libx264",
            "-preset",
            "medium",
            "-crf",
            "14",
            "-movflags",
            "+faststart",
        ]
    )
    if has_audio:
        copy_audio_command.extend(["-c:a", "copy", "-shortest"])
    else:
        copy_audio_command.append("-an")
    copy_audio_command.append(output_path)

    try:
        run_media_command(copy_audio_command, "Could not finalize the output video with copied audio.")
    except RuntimeError:
        if not has_audio:
            raise
        safe_remove(output_path)
        fallback_command = base_command + [
            "-map",
            "0:v:0",
            "-map",
            "1:a:0",
            "-vf",
            video_filter,
            "-c:v",
            "libx264",
            "-preset",
            "medium",
            "-crf",
            "14",
            "-c:a",
            "aac",
            "-b:a",
            "192k",
            "-movflags",
            "+faststart",
            "-shortest",
            output_path,
        ]
        run_media_command(fallback_command, "Could not finalize the output video.")


def finalize_video_segment(
    silent_video_path,
    source_video_path,
    output_path,
    target_width,
    target_height,
    has_audio,
    audio_start_seconds,
    audio_duration_seconds,
):
    """Encode one downloadable segment with the matching range of source audio."""
    video_filter = (
        f"scale={target_width}:{target_height}:force_original_aspect_ratio=decrease,"
        f"pad={target_width}:{target_height}:(ow-iw)/2:(oh-ih)/2:black,"
        "setsar=1,format=yuv420p"
    )
    command = [
        "ffmpeg",
        "-y",
        "-v",
        "error",
        "-i",
        silent_video_path,
    ]
    if has_audio:
        command.extend(["-i", source_video_path])
    command.extend(
        [
        "-map",
        "0:v:0",
        "-vf",
        video_filter,
        "-c:v",
        "libx264",
        "-preset",
        "medium",
        "-crf",
        "14",
        "-movflags",
        "+faststart",
        ]
    )
    if has_audio:
        audio_filter = (
            f"[1:a:0]atrim=start={audio_start_seconds:.9f}:duration={audio_duration_seconds:.9f},"
            "asetpts=PTS-STARTPTS[audio]"
        )
        command.extend(
            [
                "-filter_complex",
                audio_filter,
                "-map",
                "[audio]",
                "-c:a",
                "aac",
                "-b:a",
                "192k",
                "-shortest",
            ]
        )
    else:
        command.append("-an")
    command.append(output_path)
    run_media_command(command, "Could not finalize this video segment with its source audio.")


def restore_media_file(
    runner,
    media_path,
    text_embeds,
    fps_out,
    video_target_pixels,
    trim_leading_frames=0,
    output_path=None,
    sp_size=1,
):
    def cut_videos(videos, sp_size):
        t = videos.size(1)
        if t <= 4 * sp_size:
            print(f"Cut input video size: {videos.size()}")
            padding = [videos[:, -1].unsqueeze(1)] * (4 * sp_size - t + 1)
            padding = torch.cat(padding, dim=1)
            videos = torch.cat([videos, padding], dim=1)
            return videos
        if (t - 1) % (4 * sp_size) == 0:
            return videos
        padding = [videos[:, -1].unsqueeze(1)] * (
            4 * sp_size - ((t - 1) % (4 * sp_size))
        )
        padding = torch.cat(padding, dim=1)
        videos = torch.cat([videos, padding], dim=1)
        assert (videos.size(1) - 1) % (4 * sp_size) == 0
        return videos

    media_type, _ = mimetypes.guess_type(media_path)
    is_image = bool(media_type and media_type.startswith("image"))
    is_video = bool(media_type and media_type.startswith("video"))
    if not is_image and not is_video:
        raise RuntimeError("Please upload a supported image or video file.")

    output_path = output_path or os.path.join("output", f"{uuid.uuid4()}{'.png' if is_image else '.mp4'}")
    transform = build_image_transform() if is_image else build_video_transform(video_target_pixels)

    if is_video:
        media_tensor = read_video(media_path, output_format="TCHW")[0] / 255.0
        print(f"Read video size: {media_tensor.size()}")
    else:
        image = Image.open(media_path).convert("RGB")
        media_tensor = T.ToTensor()(image).unsqueeze(0)
        print(f"Read Image size: {media_tensor.size()}")

    cond_latent = transform(media_tensor.to(torch.device("cuda")))
    ori_length = cond_latent.size(1)
    input_video = cond_latent
    cond_latents = [cut_videos(cond_latent, sp_size)] if is_video else [cond_latent]

    print(f"Encoding videos: {list(map(lambda x: x.size(), cond_latents))}")
    cond_latents = runner.vae_encode(cond_latents)

    text_embeds_cuda = {
        "texts_pos": [embed.to(torch.device("cuda")) for embed in text_embeds["texts_pos"]],
        "texts_neg": [embed.to(torch.device("cuda")) for embed in text_embeds["texts_neg"]],
    }

    samples = generation_step(runner, text_embeds_cuda, cond_latents=cond_latents)
    del cond_latents

    sample = samples[0]
    if ori_length < sample.shape[0]:
        sample = sample[:ori_length]

    input_video = (
        rearrange(input_video[:, None], "c t h w -> t c h w")
        if input_video.ndim == 3
        else rearrange(input_video, "c t h w -> t c h w")
    )
    if use_colorfix:
        sample = wavelet_reconstruction(
            sample.to("cpu"), input_video[: sample.size(0)].to("cpu")
        )
    else:
        sample = sample.to("cpu")

    if is_video and trim_leading_frames > 0:
        if trim_leading_frames >= sample.size(0):
            raise RuntimeError("Chunk overlap trimming removed all generated frames.")
        sample = sample[trim_leading_frames:]

    sample = (
        rearrange(sample[:, None], "t c h w -> t h w c")
        if sample.ndim == 3
        else rearrange(sample, "t c h w -> t h w c")
    )
    sample = sample.clip(-1, 1).mul_(0.5).add_(0.5).mul_(255).round()
    sample = sample.to(torch.uint8).numpy()

    if is_image:
        mediapy.write_image(output_path, sample[0])
    else:
        mediapy.write_video(output_path, sample, fps=fps_out)

    gc.collect()
    torch.cuda.empty_cache()
    return output_path, is_image


def configure_runner(sp_size):
    config_path = os.path.join('./configs_3b', 'main.yaml')
    config = load_config(config_path)
    runner = VideoDiffusionInfer(config)
    OmegaConf.set_readonly(runner.config, False)
    
    init_torch(cudnn_benchmark=False, timeout=datetime.timedelta(seconds=3600))
    configure_sequence_parallel(sp_size)
    runner.configure_dit_model(device="cuda", checkpoint=MODEL_PATHS["dit"])
    runner.config.vae.checkpoint = MODEL_PATHS["vae"]
    runner.configure_vae_model()
    # Set memory limit.
    if hasattr(runner.vae, "set_memory_limit"):
        runner.vae.set_memory_limit(**runner.config.vae.memory_limit)
    return runner

def generation_step(runner, text_embeds_dict, cond_latents):
    def _move_to_cuda(x):
        return [i.to(torch.device("cuda")) for i in x]

    noises = [torch.randn_like(latent) for latent in cond_latents]
    aug_noises = [torch.randn_like(latent) for latent in cond_latents]
    print(f"Generating with noise shape: {noises[0].size()}.")
    noises, aug_noises, cond_latents = sync_data((noises, aug_noises, cond_latents), 0)
    noises, aug_noises, cond_latents = list(
        map(lambda x: _move_to_cuda(x), (noises, aug_noises, cond_latents))
    )
    cond_noise_scale = 0.1

    def _add_noise(x, aug_noise):
        t = (
            torch.tensor([1000.0], device=torch.device("cuda"))
            * cond_noise_scale
        )
        shape = torch.tensor(x.shape[1:], device=torch.device("cuda"))[None]
        t = runner.timestep_transform(t, shape)
        print(
            f"Timestep shifting from"
            f" {1000.0 * cond_noise_scale} to {t}."
        )
        x = runner.schedule.forward(x, aug_noise, t)
        return x

    conditions = [
        runner.get_condition(
            noise,
            task="sr",
            latent_blur=_add_noise(latent_blur, aug_noise),
        )
        for noise, aug_noise, latent_blur in zip(noises, aug_noises, cond_latents)
    ]

    with torch.no_grad(), torch.autocast("cuda", torch.bfloat16, enabled=True):
        video_tensors = runner.inference(
            noises=noises,
            conditions=conditions,
            dit_offload=False,
            **text_embeds_dict,
        )

    samples = [
        (
            rearrange(video[:, None], "c t h w -> t c h w")
            if video.ndim == 3
            else rearrange(video, "c t h w -> t c h w")
        )
        for video in video_tensors
    ]
    del video_tensors

    return samples


def prepare_generation_runner(seed, cfg_scale=1.0, cfg_rescale=0.0, sample_steps=1, sp_size=1):
    runner = configure_runner(sp_size)
    runner.config.diffusion.cfg.scale = cfg_scale
    runner.config.diffusion.cfg.rescale = cfg_rescale
    runner.config.diffusion.timesteps.sampling.steps = sample_steps
    runner.configure_diffusion()
    set_seed(int(seed) % (2**32), same_across_ranks=True)
    return runner

@spaces.GPU(duration=GPU_LOOP_DURATION)
def generation_loop(
    video_path="./test_videos",
    seed=666,
    fps_out=12,
    output_resolution="720p",
    batch_size=1,
    cfg_scale=1.0,
    cfg_rescale=0.0,
    sample_steps=1,
    res_h=1280,
    res_w=720,
    sp_size=1,
    progress=gr.Progress(),
):
    runner = prepare_generation_runner(
        seed,
        cfg_scale=cfg_scale,
        cfg_rescale=cfg_rescale,
        sample_steps=sample_steps,
        sp_size=sp_size,
    )
    os.makedirs('output/', exist_ok=True)
    progress(0.02, desc="Loading text embeddings")
    text_embeds = get_text_embeds()

    media_type, _ = mimetypes.guess_type(video_path)
    is_image = bool(media_type and media_type.startswith("image"))
    is_video = bool(media_type and media_type.startswith("video"))
    if not is_image and not is_video:
        raise RuntimeError("Please upload a supported image or video file.")

    target_width, target_height = VIDEO_OUTPUT_RESOLUTIONS.get(
        output_resolution, VIDEO_OUTPUT_RESOLUTIONS["720p"]
    )
    target_pixels = target_width * target_height

    if is_image:
        progress(0.1, desc="Upscaling image")
        output_path, _ = restore_media_file(
            runner=runner,
            media_path=video_path,
            text_embeds=text_embeds,
            fps_out=fps_out,
            video_target_pixels=target_pixels,
            sp_size=sp_size,
        )
        progress(1.0, desc="Done")
        return output_path, None, output_path

    progress(0.05, desc="Analysing source video")
    metadata = probe_video_metadata(video_path)
    chunk_plan = build_chunk_plan(metadata["total_frames"])
    video_fps = metadata["fps"] if metadata["fps"] > 0 else fps_out
    final_output_path = os.path.join("output", f"{uuid.uuid4()}.mp4")

    try:
        with tempfile.TemporaryDirectory(prefix="seedvr2_chunks_") as temp_dir:
            processed_chunk_paths = []
            for index, chunk in enumerate(chunk_plan, start=1):
                progress(
                    0.05 + 0.85 * ((index - 1) / max(1, len(chunk_plan))),
                    desc=f"Chunk {index}/{len(chunk_plan)}",
                )
                input_chunk_path = os.path.join(temp_dir, f"input_chunk_{index:04d}.mkv")
                output_chunk_path = os.path.join(temp_dir, f"output_chunk_{index:04d}.mp4")

                extract_video_chunk(
                    source_path=video_path,
                    chunk_path=input_chunk_path,
                    start_frame=chunk["start_frame"],
                    end_frame=chunk["end_frame"],
                )
                restore_media_file(
                    runner=runner,
                    media_path=input_chunk_path,
                    text_embeds=text_embeds,
                    fps_out=video_fps,
                    video_target_pixels=target_pixels,
                    trim_leading_frames=chunk["trim_leading_frames"],
                    output_path=output_chunk_path,
                    sp_size=sp_size,
                )
                safe_remove(input_chunk_path)
                processed_chunk_paths.append(output_chunk_path)

            progress(0.93, desc="Combining processed chunks")
            silent_video_path = os.path.join(temp_dir, "merged_silent.mp4")
            concatenate_video_chunks(processed_chunk_paths, silent_video_path, temp_dir, video_fps)

            progress(0.97, desc=f"Applying original audio and final {output_resolution} output")
            finalize_video_output(
                silent_video_path=silent_video_path,
                source_video_path=video_path,
                output_path=final_output_path,
                target_width=target_width,
                target_height=target_height,
                has_audio=metadata["has_audio"],
            )
    except Exception:
        safe_remove(final_output_path)
        raise

    progress(1.0, desc="Done")
    return None, final_output_path, final_output_path


@spaces.GPU(duration=GPU_LOOP_DURATION)
def process_next_video_segment(
    video_path,
    job_state,
    seed,
    fps_out,
    output_resolution,
    segment_duration_label,
    progress=gr.Progress(),
):
    """Process one persisted output segment per ZeroGPU allocation."""
    os.makedirs("output", exist_ok=True)
    job_dir = get_job_directory(job_state)

    if job_dir:
        source_path = job_dir / job_state["source_name"]
        metadata = job_state["metadata"]
        if not source_path.exists():
            raise RuntimeError("The saved source file is gone. Please upload the video again.")
    else:
        media_type, _ = mimetypes.guess_type(video_path)
        is_image = bool(media_type and media_type.startswith("image"))
        is_video = bool(media_type and media_type.startswith("video"))
        if not is_image and not is_video:
            raise RuntimeError("Please upload a supported image or video file.")

        target_width, target_height = VIDEO_OUTPUT_RESOLUTIONS.get(
            output_resolution, VIDEO_OUTPUT_RESOLUTIONS["720p"]
        )
        if is_image:
            progress(0.05, desc="Loading model")
            runner = prepare_generation_runner(seed)
            output_path, _ = restore_media_file(
                runner=runner,
                media_path=video_path,
                text_embeds=get_text_embeds(),
                fps_out=fps_out,
                video_target_pixels=target_width * target_height,
            )
            progress(1.0, desc="Done")
            return output_path, None, [output_path], None, "Image processing complete."

        progress(0.02, desc="Preparing resumable video job")
        metadata = probe_video_metadata(video_path)
        job_state = create_segment_job(
            video_path,
            metadata,
            output_resolution,
            segment_duration_label,
        )
        job_dir = get_job_directory(job_state)
        source_path = job_dir / job_state["source_name"]

    target_width, target_height = VIDEO_OUTPUT_RESOLUTIONS.get(
        job_state["output_resolution"], VIDEO_OUTPUT_RESOLUTIONS["720p"]
    )
    video_fps = metadata["fps"] if metadata["fps"] > 0 else float(fps_out)
    chunk_plan = build_chunk_plan(metadata["total_frames"])
    segment_plan = build_output_segment_plan(
        chunk_plan,
        video_fps,
        job_state["segment_seconds"],
    )
    segment_index = int(job_state.get("next_segment", 0))
    completed_outputs = list(job_state.get("outputs", []))
    if segment_index >= len(segment_plan):
        last_output = completed_outputs[-1] if completed_outputs else None
        return (
            None,
            last_output,
            completed_outputs,
            job_state,
            f"All {len(segment_plan)} output parts are already complete.",
        )

    segment_chunks = segment_plan[segment_index]
    first_chunk = segment_chunks[0]
    audio_start_seconds = (
        first_chunk["start_frame"] + first_chunk["trim_leading_frames"]
    ) / video_fps
    output_frame_count = sum(
        chunk["end_frame"] - chunk["start_frame"] - chunk["trim_leading_frames"]
        for chunk in segment_chunks
    )
    audio_duration_seconds = output_frame_count / video_fps
    output_path = job_dir / f"part_{segment_index + 1:04d}.mp4"

    try:
        progress(0.05, desc=f"Part {segment_index + 1}/{len(segment_plan)}: loading model")
        runner = prepare_generation_runner(seed)
        text_embeds = get_text_embeds()
        with tempfile.TemporaryDirectory(prefix="seedvr2_segment_") as temp_dir:
            processed_chunk_paths = []
            for local_index, chunk in enumerate(segment_chunks, start=1):
                global_index = chunk_plan.index(chunk) + 1
                progress(
                    0.08 + 0.78 * ((local_index - 1) / max(1, len(segment_chunks))),
                    desc=(
                        f"Part {segment_index + 1}/{len(segment_plan)} - "
                        f"Chunk {global_index}/{len(chunk_plan)}"
                    ),
                )
                input_chunk_path = os.path.join(temp_dir, f"input_{local_index:04d}.mkv")
                output_chunk_path = os.path.join(temp_dir, f"output_{local_index:04d}.mp4")
                extract_video_chunk(
                    source_path=str(source_path),
                    chunk_path=input_chunk_path,
                    start_frame=chunk["start_frame"],
                    end_frame=chunk["end_frame"],
                )
                restore_media_file(
                    runner=runner,
                    media_path=input_chunk_path,
                    text_embeds=text_embeds,
                    fps_out=video_fps,
                    video_target_pixels=target_width * target_height,
                    trim_leading_frames=chunk["trim_leading_frames"],
                    output_path=output_chunk_path,
                )
                safe_remove(input_chunk_path)
                processed_chunk_paths.append(output_chunk_path)

            progress(0.9, desc=f"Part {segment_index + 1}/{len(segment_plan)}: joining chunks")
            silent_video_path = os.path.join(temp_dir, "merged_silent.mp4")
            concatenate_video_chunks(processed_chunk_paths, silent_video_path, temp_dir, video_fps)
            progress(0.96, desc="Adding the matching original audio")
            finalize_video_segment(
                silent_video_path=silent_video_path,
                source_video_path=str(source_path),
                output_path=str(output_path),
                target_width=target_width,
                target_height=target_height,
                has_audio=metadata["has_audio"],
                audio_start_seconds=audio_start_seconds,
                audio_duration_seconds=audio_duration_seconds,
            )
    except Exception:
        safe_remove(str(output_path))
        raise

    completed_outputs.append(str(output_path))
    job_state["outputs"] = completed_outputs
    job_state["next_segment"] = segment_index + 1
    if job_state["next_segment"] == len(segment_plan):
        safe_remove(str(source_path))
        status = f"All {len(segment_plan)} output parts are complete. The temporary source copy was removed."
    else:
        status = (
            f"Part {segment_index + 1}/{len(segment_plan)} is ready. "
            "Click Process next part to continue without uploading again."
        )
    progress(1.0, desc="Part complete")
    return None, str(output_path), completed_outputs, job_state, status


def discard_video_job(job_state):
    remove_segment_job(job_state)
    return None, None, [], None, "The active job, its source copy, and its completed parts were removed."


with gr.Blocks(title="SeedVR2: One-Step Video Restoration via Diffusion Adversarial Post-Training") as demo:
    # Top logo and title
    gr.HTML("""
        <div style='text-align:center; margin-bottom: 10px;'>
            <img src='https://huggingface.co/ByteDance-Seed/SeedVR2-3B/resolve/main/assets/seedvr_logo.png' style='height:40px;' alt='SeedVR logo'/>
        </div>
        <p><b>Official Gradio demo</b> for 
        <a href='https://github.com/ByteDance-Seed/SeedVR' target='_blank'>
        <b>SeedVR2: One-Step Video Restoration via Diffusion Adversarial Post-Training</b></a>.<br>
        🔥 <b>SeedVR2</b> is a one-step image and video restoration algorithm for real-world and AIGC content.
        </p>
    """)

    # Interface
    with gr.Row():
        input_video = gr.File(
            label="Upload image or video",
            file_types=["image", "video"],
            type="filepath",
        )
        seed = gr.Number(label="Seeds", value=666)
        fps = gr.Number(label="Fallback FPS", value=24, info="Used only if the source video FPS cannot be detected.")
        output_resolution = gr.Radio(
            choices=["720p", "1080p"],
            value="720p",
            label="Video output resolution",
        )
        segment_duration = gr.Radio(
            choices=list(SEGMENT_DURATION_CHOICES),
            value="About 5 seconds (ZeroGPU-safe)",
            label="Output part length",
            info="Use 5 minutes only on a dedicated GPU. ZeroGPU processes one short part per request.",
        )

    with gr.Row():
        output_video = gr.Video(label="Output_Video")
        output_image = gr.Image(label="Output_Image")
        download_link = gr.File(label="Completed downloads", file_count="multiple")

    job_state = gr.State(value=None)
    job_status = gr.Markdown("Upload a video, then process the first part. Completed parts stay available here.")
    with gr.Row():
        run_button = gr.Button("Process next part", variant="primary")
        discard_button = gr.Button("Discard job and parts")
    run_button.click(
        fn=process_next_video_segment,
        inputs=[input_video, job_state, seed, fps, output_resolution, segment_duration],
        outputs=[output_image, output_video, download_link, job_state, job_status],
    )
    discard_button.click(
        fn=discard_video_job,
        inputs=[job_state],
        outputs=[output_image, output_video, download_link, job_state, job_status],
    )

    # Examples
    gr.Examples(
        examples=[
            ["./01.mp4", 4, 24, "720p", "About 5 seconds (ZeroGPU-safe)"],
            ["./02.mp4", 4, 24, "720p", "About 5 seconds (ZeroGPU-safe)"],
            ["./03.mp4", 4, 24, "720p", "About 5 seconds (ZeroGPU-safe)"],
        ],
        inputs=[input_video, seed, fps, output_resolution, segment_duration]
    )

    # Article/Footer
    gr.HTML("""
        <hr>
        <p>If you find SeedVR helpful, please ⭐ the 
        <a href='https://github.com/ByteDance-Seed/SeedVR' target='_blank'>GitHub repository</a>:</p>

        <a href="https://github.com/ByteDance-Seed/SeedVR" target="_blank">
            <img src="https://img.shields.io/github/stars/ByteDance-Seed/SeedVR?style=social" alt="GitHub Stars">
        </a>
        <h4>Notice</h4>
        <p>Uploaded videos are processed automatically from their detected source length. Each processed download keeps the matching portion of the original audio.
        The Space offers exact <b>720p</b> and <b>1080p</b> exports, while keeping the model on a compatible internal resolution and resizing only at the end.</p>
        <h4>How it works</h4>
        <p>Upload an image or a video, choose 720p or 1080p, then click <b>Process next part</b>. The upload is stored only while its resumable job is active. Videos are restored in memory-safe chunks of up to 121 frames. Adjacent chunks share a small frame overlap to reduce visible transitions, then each completed part is joined losslessly before its matching original audio is added back. The progress display shows the current part and chunk, for example <b>Part 3/150 - Chunk 37/150</b>.</p>
        <h4>ZeroGPU</h4>
        <p>This Space runs on Hugging Face <b>ZeroGPU</b>. The default five-second part size fits short GPU reservations, so long videos can be continued with repeated clicks without uploading again. Five-minute parts are available for dedicated GPU hardware, but may exceed a ZeroGPU allocation. Use <b>Discard job and parts</b> to remove an unfinished upload and all files created for it.</p>
        <h4>Limitations</h4>
        <p>May fail on heavy degradations or small-motion AIGC clips, causing oversharpening or poor restoration.</p>

        <h4>Citation</h4>
        <pre style="font-size: 12px;">
        @article{wang2025seedvr2,
            title={SeedVR2: One-Step Video Restoration via Diffusion Adversarial Post-Training},
            author={Wang, Jianyi and Lin, Shanchuan and Lin, Zhijie and Ren, Yuxi and Wei, Meng and Yue, Zongsheng and Zhou, Shangchen and Chen, Hao and Zhao, Yang and Yang, Ceyuan and Xiao, Xuefeng and Loy, Chen Change and Jiang, Lu},
            booktitle={arXiv preprint arXiv:2506.05301},
            year={2025}
        }

        @inproceedings{wang2025seedvr,
            title={SeedVR: Seeding Infinity in Diffusion Transformer Towards Generic Video Restoration},
            author={Wang, Jianyi and Lin, Zhijie and Wei, Meng and Zhao, Yang and Yang, Ceyuan and Loy, Chen Change and Jiang, Lu},
            booktitle={CVPR},
            year={2025}
        }
        </pre>

        <h4>License</h4>
        <p>Licensed under the 
        <a href="http://www.apache.org/licenses/LICENSE-2.0" target="_blank">Apache 2.0 License</a>.</p>

        
        </p>
    """)

demo.queue()
demo.launch()