fix parallel compile

lightx2v/models/networks/wan/infer/post_infer.py

@@ -14,7 +14,6 @@ class WanPostInfer:
     def set_scheduler(self, scheduler):
         self.scheduler = scheduler
 
-    @torch.compile(disable=not CHECK_ENABLE_GRAPH_MODE())
     def infer(self, x, pre_infer_out):
         x = self.unpatchify(x, pre_infer_out.grid_sizes)
 

lightx2v/models/networks/wan/infer/pre_infer.py

@@ -32,7 +32,6 @@ class WanPreInfer:
     def set_scheduler(self, scheduler):
         self.scheduler = scheduler
 
-    @torch.compile(disable=not CHECK_ENABLE_GRAPH_MODE())
     def infer(self, weights, inputs, kv_start=0, kv_end=0):
         x = self.scheduler.latents
 

lightx2v/models/networks/wan/infer/transformer_infer.py

@@ -97,7 +97,6 @@ class WanTransformerInfer(BaseTransformerInfer):
                 freqs_i = compute_freqs(q.size(2) // 2, grid_sizes, freqs)
         return freqs_i
 
-    @torch.compile(disable=not CHECK_ENABLE_GRAPH_MODE())
     def infer(self, weights, pre_infer_out):
         x = self.infer_main_blocks(weights, pre_infer_out)
         return self.infer_post_blocks(weights, x, pre_infer_out.embed)

lightx2v/models/networks/wan/model.py

@@ -354,6 +354,7 @@ class WanModel:
                 self.pre_weight.to_cpu()
                 self.transformer_weights.post_weights_to_cpu()
 
+    @torch.compile(disable=not CHECK_ENABLE_GRAPH_MODE())
     @torch.no_grad()
     def _infer_cond_uncond(self, inputs, infer_condition=True):
         self.scheduler.infer_condition = infer_condition