Function main

examples/WanVSR/infer_flashvsr_tiny.py:196–230

()

    return pipe

def main():
    RESULT_ROOT = "./results"
    os.makedirs(RESULT_ROOT, exist_ok=True)
    inputs = [
        "./inputs/example0.mp4",
        "./inputs/example1.mp4",
        "./inputs/example2.mp4",
        "./inputs/example3.mp4",
    ]
    seed, scale, dtype, device = 0, 4.0, torch.bfloat16, 'cuda'
    sparse_ratio = 2.0      # Recommended: 1.5 or 2.0. 1.5 → faster; 2.0 → more stable.
    pipe = init_pipeline()

    for p in inputs:
        torch.cuda.empty_cache(); torch.cuda.ipc_collect()
        name = os.path.basename(p.rstrip('/'))
        if name.startswith('.'):
            continue
        try:
            LQ, th, tw, F, fps = prepare_input_tensor(p, scale=scale, dtype=dtype, device=device)
        except Exception as e:
            print(f"[Error] {name}: {e}"); continue

        video = pipe(
            prompt="", negative_prompt="", cfg_scale=1.0, num_inference_steps=1, seed=seed,
            LQ_video=LQ, num_frames=F, height=th, width=tw, is_full_block=False, if_buffer=True,
            topk_ratio=sparse_ratio*768*1280/(th*tw), 
            kv_ratio=3.0,
            local_range=11,  # Recommended: 9 or 11. local_range=9 → sharper details; 11 → more stable results.
            color_fix = True,
        )
        video = tensor2video(video)
        save_video(video, os.path.join(RESULT_ROOT, f"FlashVSR_Tiny_{name.split('.')[0]}_seed{seed}.mp4"), fps=fps, quality=6)

    print("Done.")

if __name__ == "__main__":
    main()