Fix audio model compile and offload bugs

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

@@ -155,7 +155,6 @@ class WanTransformerInfer(BaseTransformerInfer):
 
             self.weights_stream_mgr._async_prefetch_block(weights.blocks)
 
-
         if self.clean_cuda_cache:
             del grid_sizes, embed, embed0, seq_lens, freqs, context
             torch.cuda.empty_cache()
@@ -295,9 +294,7 @@ class WanTransformerInfer(BaseTransformerInfer):
         for ipa_out in audio_dit_blocks:
             if block_idx in ipa_out:
                 cur_modify = ipa_out[block_idx]
-                x = cur_modify["modify_func"](x,
-                                            grid_sizes,
-                                            **cur_modify["kwargs"])
+                x = cur_modify["modify_func"](x, grid_sizes, **cur_modify["kwargs"])
         return x
 
     def _infer_without_offload(self, weights, grid_sizes, embed, x, embed0, seq_lens, freqs, context, audio_dit_blocks=None):

lightx2v/models/runners/wan/wan_audio_runner.py

@@ -3,7 +3,7 @@ import os
 import subprocess
 from contextlib import contextmanager
 from dataclasses import dataclass
-from typing import Any, Dict, List, Optional, Tuple
+from typing import Dict, List, Optional, Tuple
 
 import numpy as np
 import torch
@@ -738,7 +738,7 @@ class WanAudioRunner(WanRunner):  # type:ignore
                 prev_video=prev_video,
                 prev_frame_length=5,
                 segment_idx=0,
-                total_steps=1
+                total_steps=1,
             )
             # Final cleanup
             self.end_run()