[feat] support pab and multi-gpu in open-sora-plan v1.2.0 (#223)

* update parallel * update pab * update test * empty cache * update tests
NUS-HPC-AI-Lab · Sep 28, 2024 · c8dc8c4 · c8dc8c4
1 parent 868c489
commit c8dc8c4
Show file tree

Hide file tree

Showing 13 changed files with 246 additions and 300 deletions.
diff --git a/eval/pab/experiments/opensora_plan.py b/eval/pab/experiments/opensora_plan.py
@@ -1,6 +1,6 @@
 from utils import generate_func, read_prompt_list
 
-from videosys import OpenSoraPlanConfig, OpenSoraPlanPABConfig, VideoSysEngine
+from videosys import OpenSoraPlanConfig, OpenSoraPlanV110PABConfig, VideoSysEngine
 
 
 def eval_base(prompt_list):
@@ -10,7 +10,7 @@ def eval_base(prompt_list):
 
 
 def eval_pab1(prompt_list):
-    pab_config = OpenSoraPlanPABConfig(
+    pab_config = OpenSoraPlanV110PABConfig(
         spatial_gap=2,
         temporal_gap=4,
         cross_gap=6,
@@ -21,7 +21,7 @@ def eval_pab1(prompt_list):
 
 
 def eval_pab2(prompt_list):
-    pab_config = OpenSoraPlanPABConfig(
+    pab_config = OpenSoraPlanV110PABConfig(
         spatial_gap=3,
         temporal_gap=5,
         cross_gap=7,
@@ -32,7 +32,7 @@ def eval_pab2(prompt_list):
 
 
 def eval_pab3(prompt_list):
-    pab_config = OpenSoraPlanPABConfig(
+    pab_config = OpenSoraPlanV110PABConfig(
         spatial_gap=5,
         temporal_gap=7,
         cross_gap=9,

diff --git a/examples/open_sora_plan/sample.py b/examples/open_sora_plan/sample.py
@@ -5,7 +5,7 @@ def run_base():
     # open-sora-plan v1.2.0
     # transformer_type (len, res): 93x480p 93x720p 29x480p 29x720p
     # change num_gpus for multi-gpu inference
-    config = OpenSoraPlanConfig(version="v120", transformer_type="93x480p", num_gpus=1)
+    config = OpenSoraPlanConfig(version="v120", transformer_type="29x480p", num_gpus=1)
     engine = VideoSysEngine(config)
 
     prompt = "Sunset over the sea."

diff --git a/tests/pipelines/cogvideox/test_cogvideox.py b/tests/pipelines/cogvideox/test_cogvideox.py
@@ -1,33 +1,40 @@
 import pytest
 
 from videosys import CogVideoXConfig, VideoSysEngine
+from videosys.utils.test import empty_cache
 
 
 @pytest.mark.parametrize("num_gpus", [1, 2])
-def test_base(num_gpus):
-    config = CogVideoXConfig(num_gpus=num_gpus)
+@pytest.mark.parametrize("model_path", ["THUDM/CogVideoX-2b", "THUDM/CogVideoX-5b"])
+@empty_cache
+def test_base(num_gpus, model_path):
+    config = CogVideoXConfig(model_path=model_path, num_gpus=num_gpus)
     engine = VideoSysEngine(config)
 
     prompt = "Sunset over the sea."
     video = engine.generate(prompt, seed=0).video[0]
-    engine.save_video(video, f"./test_outputs/{prompt}_cogvideo_{num_gpus}.mp4")
+    engine.save_video(video, f"./test_outputs/{prompt}_cogvideo_{model_path.replace('/', '_')}_base_{num_gpus}.mp4")
 
 
 @pytest.mark.parametrize("num_gpus", [1])
-def test_pab(num_gpus):
-    config = CogVideoXConfig(num_gpus=num_gpus, enable_pab=True)
+@pytest.mark.parametrize("model_path", ["THUDM/CogVideoX-2b"])
+@empty_cache
+def test_pab(num_gpus, model_path):
+    config = CogVideoXConfig(model_path=model_path, num_gpus=num_gpus, enable_pab=True)
     engine = VideoSysEngine(config)
 
     prompt = "Sunset over the sea."
     video = engine.generate(prompt, seed=0).video[0]
-    engine.save_video(video, f"./test_outputs/{prompt}_cogvideo_pab_{num_gpus}.mp4")
+    engine.save_video(video, f"./test_outputs/{prompt}_cogvideo_{model_path.replace('/', '_')}_pab_{num_gpus}.mp4")
 
 
 @pytest.mark.parametrize("num_gpus", [1])
-def test_low_mem(num_gpus):
-    config = CogVideoXConfig(num_gpus=num_gpus, cpu_offload=True, vae_tiling=True)
+@pytest.mark.parametrize("model_path", ["THUDM/CogVideoX-2b"])
+@empty_cache
+def test_low_mem(num_gpus, model_path):
+    config = CogVideoXConfig(model_path=model_path, num_gpus=num_gpus, cpu_offload=True, vae_tiling=True)
     engine = VideoSysEngine(config)
 
     prompt = "Sunset over the sea."
     video = engine.generate(prompt, seed=0).video[0]
-    engine.save_video(video, f"./test_outputs/{prompt}_cogvideo_low_mem_{num_gpus}.mp4")
+    engine.save_video(video, f"./test_outputs/{prompt}_cogvideo_{model_path.replace('/', '_')}_low_mem_{num_gpus}.mp4")
diff --git a/tests/pipelines/latte/test_latte.py b/tests/pipelines/latte/test_latte.py
@@ -1,9 +1,11 @@
 import pytest
 
 from videosys import LatteConfig, VideoSysEngine
+from videosys.utils.test import empty_cache
 
 
 @pytest.mark.parametrize("num_gpus", [1, 2])
+@empty_cache
 def test_base(num_gpus):
     config = LatteConfig(num_gpus=num_gpus)
     engine = VideoSysEngine(config)
@@ -14,6 +16,7 @@ def test_base(num_gpus):
 
 
 @pytest.mark.parametrize("num_gpus", [1])
+@empty_cache
 def test_pab(num_gpus):
     config = LatteConfig(num_gpus=num_gpus, enable_pab=True)
     engine = VideoSysEngine(config)
@@ -24,6 +27,7 @@ def test_pab(num_gpus):
 
 
 @pytest.mark.parametrize("num_gpus", [1])
+@empty_cache
 def test_low_mem(num_gpus):
     config = LatteConfig(num_gpus=num_gpus, cpu_offload=True)
     engine = VideoSysEngine(config)

diff --git a/tests/pipelines/open_sora/test_open_sora.py b/tests/pipelines/open_sora/test_open_sora.py
@@ -1,9 +1,11 @@
 import pytest
 
 from videosys import OpenSoraConfig, VideoSysEngine
+from videosys.utils.test import empty_cache
 
 
 @pytest.mark.parametrize("num_gpus", [1, 2])
+@empty_cache
 def test_base(num_gpus):
     config = OpenSoraConfig(num_gpus=num_gpus)
     engine = VideoSysEngine(config)
@@ -14,6 +16,7 @@ def test_base(num_gpus):
 
 
 @pytest.mark.parametrize("num_gpus", [1])
+@empty_cache
 def test_pab(num_gpus):
     config = OpenSoraConfig(num_gpus=num_gpus, enable_pab=True)
     engine = VideoSysEngine(config)
@@ -24,6 +27,7 @@ def test_pab(num_gpus):
 
 
 @pytest.mark.parametrize("num_gpus", [1])
+@empty_cache
 def test_low_mem(num_gpus):
     config = OpenSoraConfig(num_gpus=num_gpus, cpu_offload=True, tiling_size=1)
     engine = VideoSysEngine(config)

diff --git a/tests/pipelines/open_sora_plan/test_open_sora_plan.py b/tests/pipelines/open_sora_plan/test_open_sora_plan.py
@@ -1,33 +1,42 @@
 import pytest
 
 from videosys import OpenSoraPlanConfig, VideoSysEngine
+from videosys.utils.test import empty_cache
 
 
 @pytest.mark.parametrize("num_gpus", [1, 2])
-def test_base(num_gpus):
-    config = OpenSoraPlanConfig(num_gpus=num_gpus)
+@pytest.mark.parametrize("model", [("v120", "29x480p")])
+@empty_cache
+def test_base(num_gpus, model):
+    config = OpenSoraPlanConfig(version=model[0], transformer_type=model[1], num_gpus=num_gpus)
     engine = VideoSysEngine(config)
 
     prompt = "Sunset over the sea."
     video = engine.generate(prompt, seed=0).video[0]
-    engine.save_video(video, f"./test_outputs/{prompt}_open_sora_plan_{num_gpus}.mp4")
+    engine.save_video(video, f"./test_outputs/{prompt}_open_sora_plan_{model[0]}_{model[1]}_{num_gpus}.mp4")
 
 
 @pytest.mark.parametrize("num_gpus", [1])
-def test_pab(num_gpus):
-    config = OpenSoraPlanConfig(num_gpus=num_gpus, enable_pab=True)
+@pytest.mark.parametrize("model", [("v120", "29x480p")])
+@empty_cache
+def test_pab(num_gpus, model):
+    config = OpenSoraPlanConfig(version=model[0], transformer_type=model[1], num_gpus=num_gpus, enable_pab=True)
     engine = VideoSysEngine(config)
 
     prompt = "Sunset over the sea."
     video = engine.generate(prompt, seed=0).video[0]
-    engine.save_video(video, f"./test_outputs/{prompt}_open_sora_plan_pab_{num_gpus}.mp4")
+    engine.save_video(video, f"./test_outputs/{prompt}_open_sora_plan_{model[0]}_{model[1]}_pab_{num_gpus}.mp4")
 
 
 @pytest.mark.parametrize("num_gpus", [1])
-def test_low_mem(num_gpus):
-    config = OpenSoraPlanConfig(num_gpus=num_gpus, cpu_offload=True, enable_tiling=True)
+@pytest.mark.parametrize("model", [("v120", "29x480p")])
+@empty_cache
+def test_low_mem(num_gpus, model):
+    config = OpenSoraPlanConfig(
+        version=model[0], transformer_type=model[1], num_gpus=num_gpus, cpu_offload=True, enable_tiling=True
+    )
     engine = VideoSysEngine(config)
 
     prompt = "Sunset over the sea."
     video = engine.generate(prompt, seed=0).video[0]
-    engine.save_video(video, f"./test_outputs/{prompt}_open_sora_plan_low_mem_{num_gpus}.mp4")
+    engine.save_video(video, f"./test_outputs/{prompt}_open_sora_plan_{model[0]}_{model[1]}_low_mem_{num_gpus}.mp4")
diff --git a/tests/pipelines/vchitect/test_vchitect.py b/tests/pipelines/vchitect/test_vchitect.py
@@ -1,33 +1,44 @@
 import pytest
 
 from videosys import VchitectConfig, VideoSysEngine
+from videosys.utils.test import empty_cache
 
 
 @pytest.mark.parametrize("num_gpus", [1, 2])
-def test_base(num_gpus):
-    config = VchitectConfig(num_gpus=num_gpus)
+@pytest.mark.parametrize("model_path", ["Vchitect/Vchitect-2.0-2B"])
+@empty_cache
+def test_base(num_gpus, model_path):
+    config = VchitectConfig(model_path=model_path, num_gpus=num_gpus)
     engine = VideoSysEngine(config)
 
     prompt = "Sunset over the sea."
     video = engine.generate(prompt, seed=0).video[0]
-    engine.save_video(video, f"./test_outputs/{prompt}_vchitect_{num_gpus}.mp4")
+    engine.save_video(video, f"./test_outputs/{prompt}_vchitect_{model_path.replace('/', '_')}_base_{num_gpus}.mp4")
 
 
 @pytest.mark.parametrize("num_gpus", [1])
-def test_pab(num_gpus):
-    config = VchitectConfig(num_gpus=num_gpus, enable_pab=True)
+@pytest.mark.parametrize("model_path", ["Vchitect/Vchitect-2.0-2B"])
+@empty_cache
+def test_pab(num_gpus, model_path):
+    config = VchitectConfig(model_path=model_path, num_gpus=num_gpus, enable_pab=True)
     engine = VideoSysEngine(config)
 
     prompt = "Sunset over the sea."
     video = engine.generate(prompt, seed=0).video[0]
-    engine.save_video(video, f"./test_outputs/{prompt}_vchitect_pab_{num_gpus}.mp4")
+    engine.save_video(video, f"./test_outputs/{prompt}_vchitect_{model_path.replace('/', '_')}_pab_{num_gpus}.mp4")
 
 
 @pytest.mark.parametrize("num_gpus", [1])
-def test_low_mem(num_gpus):
-    config = VchitectConfig(num_gpus=num_gpus, cpu_offload=True)
+@pytest.mark.parametrize("model_path", ["Vchitect/Vchitect-2.0-2B"])
+@empty_cache
+def test_low_mem(num_gpus, model_path):
+    config = VchitectConfig(
+        model_path=model_path,
+        num_gpus=num_gpus,
+        cpu_offload=True,
+    )
     engine = VideoSysEngine(config)
 
     prompt = "Sunset over the sea."
     video = engine.generate(prompt, seed=0).video[0]
-    engine.save_video(video, f"./test_outputs/{prompt}_vchitect_low_mem_{num_gpus}.mp4")
+    engine.save_video(video, f"./test_outputs/{prompt}_vchitect_{model_path.replace('/', '_')}_low_mem_{num_gpus}.mp4")
diff --git a/videosys/__init__.py b/videosys/__init__.py
@@ -3,15 +3,20 @@
 from .pipelines.cogvideox import CogVideoXConfig, CogVideoXPABConfig, CogVideoXPipeline
 from .pipelines.latte import LatteConfig, LattePABConfig, LattePipeline
 from .pipelines.open_sora import OpenSoraConfig, OpenSoraPABConfig, OpenSoraPipeline
-from .pipelines.open_sora_plan import OpenSoraPlanConfig, OpenSoraPlanPABConfig, OpenSoraPlanPipeline
-from .pipelines.vchitect import VchitectConfig, VchitectXLPipeline
+from .pipelines.open_sora_plan import (
+    OpenSoraPlanConfig,
+    OpenSoraPlanPipeline,
+    OpenSoraPlanV110PABConfig,
+    OpenSoraPlanV120PABConfig,
+)
+from .pipelines.vchitect import VchitectConfig, VchitectPABConfig, VchitectXLPipeline
 
 __all__ = [
     "initialize",
     "VideoSysEngine",
     "LattePipeline", "LatteConfig", "LattePABConfig",
-    "OpenSoraPlanPipeline", "OpenSoraPlanConfig", "OpenSoraPlanPABConfig",
+    "OpenSoraPlanPipeline", "OpenSoraPlanConfig", "OpenSoraPlanV110PABConfig", "OpenSoraPlanV120PABConfig",
     "OpenSoraPipeline", "OpenSoraConfig", "OpenSoraPABConfig",
     "CogVideoXPipeline", "CogVideoXConfig", "CogVideoXPABConfig",
-    "VchitectXLPipeline", "VchitectConfig",
+    "VchitectXLPipeline", "VchitectConfig", "VchitectPABConfig"
 ]  # fmt: skip