Update generate.py

generate.py

@@ -12,8 +12,8 @@ from .utils import prepare_control, load_latent, load_video, prepare_depth, save
 from .pnp_utils import register_time, register_attention_control, register_conv_control
 
 # will cause an issue
-from . import vidtome
-
+# from . import vidtome
+from .vidtome import update_patch, update_patch
 # suppress partial model loading warning
 logging.set_verbosity_error()
 
@@ -95,7 +95,7 @@ class Generator(nn.Module):
             self.pipe.load_lora_weights(**gene_config.lora)
     
     def activate_vidtome(self):
-        vidtome.apply_patch(self.pipe, self.local_merge_ratio, self.merge_global, self.global_merge_ratio, 
+        apply_patch(self.pipe, self.local_merge_ratio, self.merge_global, self.global_merge_ratio, 
             seed = self.seed, batch_size = self.batch_size, align_batch = self.use_pnp or self.align_batch, global_rand = self.global_rand)        
 
     @torch.no_grad()
@@ -234,7 +234,7 @@ class Generator(nn.Module):
     def post_iter(self, x, t):
         if self.merge_global:
             # Reset global tokens
-            vidtome.update_patch(self.pipe, global_tokens = None)
+            update_patch(self.pipe, global_tokens = None)
 
     @torch.no_grad()
     def pred_noise(self, x, cond, t, batch_idx=None):