More cleanup

optimization.py

@@ -9,8 +9,8 @@ import spaces
 import torch
 from torch.utils._pytree import tree_map_only
 
-from pipeline_utils import capture_component_call
-from zerogpu import aoti_compile
+from optimization_utils import capture_component_call
+from optimization_utils import aoti_compile
 
 
 P = ParamSpec('P')

zerogpu.py → optimization_utils.py

@@ -1,9 +1,11 @@
 """
 """
+import contextlib
 from contextvars import ContextVar
 from io import BytesIO
 from typing import Any
 from typing import cast
+from unittest.mock import patch
 
 import torch
 from torch._inductor.package.package import package_aoti
@@ -60,3 +62,35 @@ def aoti_compile(
     package_aoti(archive_file, files)
     weights, = (artifact for artifact in artifacts if isinstance(artifact, Weights))
     return ZeroGPUCompiledModel(archive_file, weights)
+
+
+@contextlib.contextmanager
+def capture_component_call(
+    pipeline: Any,
+    component_name: str,
+    component_method='forward',
+):
+
+    class CapturedCallException(Exception):
+        def __init__(self, *args, **kwargs):
+            super().__init__()
+            self.args = args
+            self.kwargs = kwargs
+
+    class CapturedCall:
+        def __init__(self):
+            self.args: tuple[Any, ...] = ()
+            self.kwargs: dict[str, Any] = {}
+
+    component = getattr(pipeline, component_name)
+    captured_call = CapturedCall()
+
+    def capture_call(*args, **kwargs):
+        raise CapturedCallException(*args, **kwargs)
+
+    with patch.object(component, component_method, new=capture_call):
+        try:
+            yield captured_call
+        except CapturedCallException as e:
+            captured_call.args = e.args
+            captured_call.kwargs = e.kwargs

pipeline_utils.py

@@ -1,40 +0,0 @@
-"""
-"""
-
-import contextlib
-from unittest.mock import patch
-
-from typing import Any
-
-
-class CapturedCallException(Exception):
-    def __init__(self, *args, **kwargs):
-        super().__init__()
-        self.args = args
-        self.kwargs = kwargs
-
-
-class CapturedCall:
-    def __init__(self):
-        self.args: tuple[Any, ...] = ()
-        self.kwargs: dict[str, Any] = {}
-
-
-@contextlib.contextmanager
-def capture_component_call(
-    pipeline: Any,
-    component_name: str,
-    component_method='forward',
-):
-    component = getattr(pipeline, component_name)
-    captured_call = CapturedCall()
-
-    def capture_call(*args, **kwargs):
-        raise CapturedCallException(*args, **kwargs)
-
-    with patch.object(component, component_method, new=capture_call):
-        try:
-            yield captured_call
-        except CapturedCallException as e:
-            captured_call.args = e.args
-            captured_call.kwargs = e.kwargs