Make transpose_to_landscape picklable for ZeroGPU (closures -> __call__ classes)

dust3r/dust3r/utils/misc.py

@@ -51,19 +51,24 @@ def interleave(tensor1, tensor2):
     return res1, res2
 
 
-def transpose_to_landscape(head, activate=True):
-    """ Predict in the correct aspect-ratio,
-        then transpose the result in landscape 
-        and stack everything back together.
-    """
-    def wrapper_no(decout, true_shape):
-        B = len(true_shape)
+class _TransposeToLandscapeNo:
+    """Picklable replacement for the `wrapper_no` closure — needed so the
+    model can be pickled across processes (e.g. ZeroGPU's worker fork)."""
+    def __init__(self, head):
+        self.head = head
+
+    def __call__(self, decout, true_shape):
         assert true_shape[0:1].allclose(true_shape), 'true_shape must be all identical'
         H, W = true_shape[0].cpu().tolist()
-        res = head(decout, (H, W))
-        return res
+        return self.head(decout, (H, W))
+
 
-    def wrapper_yes(decout, true_shape):
+class _TransposeToLandscapeYes:
+    """Picklable replacement for the `wrapper_yes` closure."""
+    def __init__(self, head):
+        self.head = head
+
+    def __call__(self, decout, true_shape):
         B = len(true_shape)
         # by definition, the batch is in landscape mode so W >= H
         H, W = int(true_shape.min()), int(true_shape.max())
@@ -72,16 +77,15 @@ def transpose_to_landscape(head, activate=True):
         is_landscape = (width >= height)
         is_portrait = ~is_landscape
 
-        # true_shape = true_shape.cpu()
         if is_landscape.all():
-            return head(decout, (H, W))
+            return self.head(decout, (H, W))
         if is_portrait.all():
-            return transposed(head(decout, (W, H)))
+            return transposed(self.head(decout, (W, H)))
 
-        # batch is a mix of both portraint & landscape
+        # batch is a mix of both portrait & landscape
         def selout(ar): return [d[ar] for d in decout]
-        l_result = head(selout(is_landscape), (H, W))
-        p_result = transposed(head(selout(is_portrait), (W, H)))
+        l_result = self.head(selout(is_landscape), (H, W))
+        p_result = transposed(self.head(selout(is_portrait), (W, H)))
 
         # allocate full result
         result = {}
@@ -93,7 +97,14 @@ def transpose_to_landscape(head, activate=True):
 
         return result
 
-    return wrapper_yes if activate else wrapper_no
+
+def transpose_to_landscape(head, activate=True):
+    """Predict in the correct aspect-ratio, then transpose the result in
+    landscape and stack everything back together.
+    Returns a picklable callable (important for ZeroGPU, which sends the model
+    across a multiprocessing boundary and can't pickle local-function closures).
+    """
+    return _TransposeToLandscapeYes(head) if activate else _TransposeToLandscapeNo(head)
 
 
 def transposed(dic):