Implemented cpu offload (#197)

* Implemented cpu offload for other frameworks

* polish example code

* Added test for low mem settings.

* polish

* format

* fix arg

* update pipeline and test

---------

Co-authored-by: ExtremeViscent <[email protected]>

Author: oahzxl

examples/latte/sample.py

@@ -16,6 +16,15 @@ def run_base():
     engine.save_video(video, f"./outputs/{prompt}.mp4")
 
 
+def run_low_mem():
+    config = LatteConfig("maxin-cn/Latte-1", cpu_offload=True)
+    engine = VideoSysEngine(config)
+
+    prompt = "Sunset over the sea."
+    video = engine.generate(prompt).video[0]
+    engine.save_video(video, f"./outputs/{prompt}.mp4")
+
+
 def run_pab():
     config = LatteConfig("maxin-cn/Latte-1", enable_pab=True)
     engine = VideoSysEngine(config)
@@ -27,4 +36,5 @@ def run_pab():
 
 if __name__ == "__main__":
     run_base()
+    # run_low_mem()
     # run_pab()

examples/open_sora/sample.py

@@ -20,6 +20,15 @@ def run_base():
     engine.save_video(video, f"./outputs/{prompt}.mp4")
 
 
+def run_low_mem():
+    config = OpenSoraConfig(cpu_offload=True, tiling_size=1)
+    engine = VideoSysEngine(config)
+
+    prompt = "Sunset over the sea."
+    video = engine.generate(prompt).video[0]
+    engine.save_video(video, f"./outputs/{prompt}.mp4")
+
+
 def run_pab():
     config = OpenSoraConfig(enable_pab=True)
     engine = VideoSysEngine(config)
@@ -31,4 +40,5 @@ def run_pab():
 
 if __name__ == "__main__":
     run_base()
+    # run_low_mem()
     # run_pab()

examples/open_sora_plan/sample.py

@@ -16,8 +16,17 @@ def run_base():
     engine.save_video(video, f"./outputs/{prompt}.mp4")
 
 
+def run_low_mem():
+    config = OpenSoraPlanConfig(cpu_offload=True, enable_tiling=True)
+    engine = VideoSysEngine(config)
+
+    prompt = "Sunset over the sea."
+    video = engine.generate(prompt).video[0]
+    engine.save_video(video, f"./outputs/{prompt}.mp4")
+
+
 def run_pab():
-    config = OpenSoraPlanConfig(num_gpus=1, enable_pab=True)
+    config = OpenSoraPlanConfig(enable_pab=True)
     engine = VideoSysEngine(config)
 
     prompt = "Sunset over the sea."
@@ -27,4 +36,5 @@ def run_pab():
 
 if __name__ == "__main__":
     run_base()
+    # run_low_mem()
     # run_pab()

tests/examples/test_sample.py

@@ -12,12 +12,19 @@
 import examples.open_sora.sample as open_sora
 import examples.open_sora_plan.sample as open_sora_plan
 
+files = [cogvideox, latte, open_sora, open_sora_plan]
+members = []
 
-@pytest.mark.parametrize("file", [cogvideox, latte, open_sora, open_sora_plan])
-def test_examples(file):
-    funcs = inspect.getmembers(file, inspect.isfunction)
-    for name, func in funcs:
-        try:
-            func()
-        except Exception as e:
-            raise Exception(f"Failed to run {name} in {file.__file__}") from e
+for file in files:
+    for m in inspect.getmembers(file, inspect.isfunction):
+        members.append(m)
+print(members)
+
+
+@pytest.mark.parametrize("members", members)
+def test_examples(members):
+    name, func = members
+    try:
+        func()
+    except Exception as e:
+        raise Exception(f"Failed to run {name} in {file.__file__}") from e

tests/pipelines/cogvideox/test_cogvideox.py

@@ -21,3 +21,13 @@ def test_pab(num_gpus):
     prompt = "Sunset over the sea."
     video = engine.generate(prompt).video[0]
     engine.save_video(video, f"./test_outputs/{prompt}_cogvideo_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)
+    engine = VideoSysEngine(config)
+
+    prompt = "Sunset over the sea."
+    video = engine.generate(prompt).video[0]
+    engine.save_video(video, f"./test_outputs/{prompt}_cogvideo_low_mem_{num_gpus}.mp4")

tests/pipelines/latte/test_latte.py

@@ -21,3 +21,13 @@ def test_pab(num_gpus):
     prompt = "Sunset over the sea."
     video = engine.generate(prompt).video[0]
     engine.save_video(video, f"./test_outputs/{prompt}_latte_pab_{num_gpus}.mp4")
+
+
+@pytest.mark.parametrize("num_gpus", [1])
+def test_low_mem(num_gpus):
+    config = LatteConfig(num_gpus=num_gpus, cpu_offload=True)
+    engine = VideoSysEngine(config)
+
+    prompt = "Sunset over the sea."
+    video = engine.generate(prompt).video[0]
+    engine.save_video(video, f"./test_outputs/{prompt}_latte_low_mem_{num_gpus}.mp4")

tests/pipelines/open_sora/test_open_sora.py

@@ -21,3 +21,13 @@ def test_pab(num_gpus):
     prompt = "Sunset over the sea."
     video = engine.generate(prompt).video[0]
     engine.save_video(video, f"./test_outputs/{prompt}_open_sora_pab_{num_gpus}.mp4")
+
+
+@pytest.mark.parametrize("num_gpus", [1])
+def test_low_mem(num_gpus):
+    config = OpenSoraConfig(num_gpus=num_gpus, cpu_offload=True, tiling_size=1)
+    engine = VideoSysEngine(config)
+
+    prompt = "Sunset over the sea."
+    video = engine.generate(prompt).video[0]
+    engine.save_video(video, f"./test_outputs/{prompt}_open_sora_low_mem_{num_gpus}.mp4")

tests/pipelines/open_sora_plan/test_open_sora_plan.py

@@ -21,3 +21,13 @@ def test_pab(num_gpus):
     prompt = "Sunset over the sea."
     video = engine.generate(prompt).video[0]
     engine.save_video(video, f"./test_outputs/{prompt}_open_sora_plan_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)
+    engine = VideoSysEngine(config)
+
+    prompt = "Sunset over the sea."
+    video = engine.generate(prompt).video[0]
+    engine.save_video(video, f"./test_outputs/{prompt}_open_sora_plan_low_mem_{num_gpus}.mp4")

videosys/core/engine.py

@@ -3,6 +3,7 @@
 from typing import Any, Optional
 
 import torch
+import torch.distributed as dist
 
 import videosys
 
@@ -22,9 +23,6 @@ def __init__(self, config):
     def _init_worker(self, pipeline_cls):
         world_size = self.config.num_gpus
 
-        if "CUDA_VISIBLE_DEVICES" not in os.environ:
-            os.environ["CUDA_VISIBLE_DEVICES"] = ",".join(str(i) for i in range(world_size))
-
         # Disable torch async compiling which won't work with daemonic processes
         os.environ["TORCHINDUCTOR_COMPILE_THREADS"] = "1"
 
@@ -124,7 +122,7 @@ def save_video(self, video, output_path):
     def shutdown(self):
         if (worker_monitor := getattr(self, "worker_monitor", None)) is not None:
             worker_monitor.close()
-        torch.distributed.destroy_process_group()
+        dist.destroy_process_group()
 
     def __del__(self):
         self.shutdown()

videosys/models/autoencoders/autoencoder_kl_open_sora.py

@@ -670,6 +670,9 @@ def encode(self, x):
             return (z - self.shift) / self.scale
 
     def decode(self, z, num_frames=None):
+        device = z.device
+        self.scale = self.scale.to(device)
+        self.shift = self.shift.to(device)
         if not self.cal_loss:
             z = z * self.scale.to(z.dtype) + self.shift.to(z.dtype)