From d57156114909dd7b55e4749c45f97ee46b4a3372 Mon Sep 17 00:00:00 2001
From: Roberts Slisans <roberts.slisans@gmail.com>
Date: Thu, 8 May 2025 19:05:32 +0300
Subject: [PATCH] add the cpu_offload module

---
 acestep/cpu_offload.py | 41 +++++++++++++++++++++++++++++++++++++++++
 1 file changed, 41 insertions(+)
 create mode 100644 acestep/cpu_offload.py

diff --git a/acestep/cpu_offload.py b/acestep/cpu_offload.py
new file mode 100644
index 0000000..efe6044
--- /dev/null
+++ b/acestep/cpu_offload.py
@@ -0,0 +1,41 @@
+import torch
+import functools
+from typing import Callable, TypeVar
+
+
+class CpuOffloader:
+    def __init__(self, model, device="cpu"):
+        self.model = model
+        self.original_device = device
+        self.original_dtype = model.dtype
+    
+    def __enter__(self):
+        self.model.to(self.original_device, dtype=self.original_dtype)
+        return self.model
+    
+    def __exit__(self, *args):
+        self.model.to("cpu")
+        if torch.cuda.is_available():
+            torch.cuda.empty_cache()
+            torch.cuda.synchronize()
+
+
+T = TypeVar('T')
+
+def cpu_offload(model_attr: str):
+    def decorator(func: Callable[..., T]) -> Callable[..., T]:
+        @functools.wraps(func)
+        def wrapper(self, *args, **kwargs):
+            if not self.cpu_offload:
+                return func(self, *args, **kwargs)
+
+            # Get the device from the class
+            device = self.device
+            # Get the model from the class attribute
+            model = getattr(self, model_attr)
+            
+            with CpuOffloader(model, device):
+                return func(self, *args, **kwargs)
+                        
+        return wrapper
+    return decorator