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

* update parallel

* update pab

* update test

* empty cache

* update tests

Author: oahzxl

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,

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."

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")

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)

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)

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")

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")

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