Update patch.py

patch.py

@@ -7,7 +7,7 @@ from einops import rearrange
 import torch
 import torch.nn.functional as F
 
-from . import merge
+from .merge import bipartite_soft_matching_randframe, bipartite_soft_matching_2s, do_nothing
 from .utils import isinstance_str, init_generator, join_frame, split_frame, func_warper, join_warper, split_warper
 
 
@@ -42,7 +42,7 @@ def compute_merge(module: torch.nn.Module, x: torch.Tensor, tome_info: Dict[str,
 
         # Recursive merge multi-frame tokens into one set. Such as 4->1 for 4 frames and 8->2->1 for 8 frames when target stride is 4.
         while curF > 1:
-            m, u, ret_dict = merge.bipartite_soft_matching_randframe(
+            m, u, ret_dict = bipartite_soft_matching_randframe(
                 local_tokens, curF, args["local_merge_ratio"], unm, generator, args["target_stride"], args["align_batch"])
             unm += ret_dict["unm_num"]
             m_ls.append(m)
@@ -70,7 +70,7 @@ def compute_merge(module: torch.nn.Module, x: torch.Tensor, tome_info: Dict[str,
                         [module.global_tokens.to(local_tokens), local_tokens], dim=1)
                     local_chunk = 1
 
-                m, u, _ = merge.bipartite_soft_matching_2s(
+                m, u, _ = bipartite_soft_matching_2s(
                     tokens, src_len, args["global_merge_ratio"], args["align_batch"], unmerge_chunk=local_chunk)
                 merged_tokens = m(tokens)
                 m_ls.append(m)
@@ -84,7 +84,7 @@ def compute_merge(module: torch.nn.Module, x: torch.Tensor, tome_info: Dict[str,
         m = func_warper(m_ls)
         u = func_warper(u_ls[::-1])
     else:
-        m, u = (merge.do_nothing, merge.do_nothing)
+        m, u = (do_nothing, do_nothing)
         merged_tokens = x
 
     # Return merge op, unmerge op, and merged tokens.