Spaces:
Running
Running
File size: 2,307 Bytes
060fbda | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 | from contextlib import contextmanager
from pathlib import Path
from time import perf_counter
from sorawm.core import SoraWM
from sorawm.schemas import CleanerType
@contextmanager
def timer(name: str):
start = perf_counter()
yield
elapsed = perf_counter() - start
print(f"[{name}] Time elapsed: {elapsed:.2f}s")
if __name__ == "__main__":
input_video_path = Path("resources/dog_vs_sam.mp4")
output_video_path = Path("outputs/sora_watermark_removed")
# # 1. LAMA is fast and good quality, but not time consistent.
# sora_wm = SoraWM(cleaner_type=CleanerType.LAMA)
# with timer("LAMA"):
# sora_wm.run(input_video_path, Path(f"{output_video_path}_lama.mp4"))
# # 2. E2FGVI_HQ ensures time consistency, but will be very slow on no-cuda device.
# sora_wm = SoraWM(cleaner_type=CleanerType.E2FGVI_HQ, enable_torch_compile=False)
# with timer("E2FGVI_HQ"):
# sora_wm.run(input_video_path, Path(f"{output_video_path}_e2fgvi_hq.mp4"))
# 3. E2FGVI_HQ with torch compile is fast and good quality, but not time consistent.
sora_wm = SoraWM(cleaner_type=CleanerType.E2FGVI_HQ, enable_torch_compile=True)
with timer("E2FGVI_HQ + torch.compile"):
sora_wm.run(
input_video_path, Path(f"{output_video_path}_e2fgvi_hq_torch_compile.mp4")
)
# 4. Enable batch detection
batch_size = 4
sora_wm = SoraWM(
cleaner_type=CleanerType.E2FGVI_HQ,
enable_torch_compile=True,
detect_batch_size=4,
)
with timer("E2FGVI_HQ + torch.compile + batch"):
sora_wm.run(
input_video_path,
Path(f"{output_video_path}_e2fgvi_hq_torch_compile_batch.mp4"),
)
# 5. Enable bf16 inference
sora_wm = SoraWM(
cleaner_type=CleanerType.E2FGVI_HQ,
enable_torch_compile=True,
detect_batch_size=4,
use_bf16=True,
)
with timer("E2FGVI_HQ + torch.compile + batch + bf16"):
sora_wm.run(
input_video_path,
Path(f"{output_video_path}_e2fgvi_hq_torch_compile_batch_bf16.mp4"),
)
with timer("E2FGVI_HQ + torch.compile + batch + bf16"):
sora_wm.run(
input_video_path,
Path(f"{output_video_path}_e2fgvi_hq_torch_compile_batch_bf16.mp4"),
)
|