♻️ Refactor: Move audio inference files to 'infer/audio' subdirectory

lightx2v/models/networks/wan/audio_model.py

@@ -12,8 +12,8 @@ from lightx2v.models.networks.wan.infer.pre_infer import WanPreInfer
 from lightx2v.models.networks.wan.infer.post_infer import WanPostInfer
 
 from lightx2v.models.networks.wan.infer.pre_infer import WanPreInfer
-from lightx2v.models.networks.wan.infer.pre_wan_audio_infer import WanAudioPreInfer
-from lightx2v.models.networks.wan.infer.post_wan_audio_infer import WanAudioPostInfer
+from lightx2v.models.networks.wan.infer.audio.pre_wan_audio_infer import WanAudioPreInfer
+from lightx2v.models.networks.wan.infer.audio.post_wan_audio_infer import WanAudioPostInfer
 from lightx2v.models.networks.wan.infer.feature_caching.transformer_infer import WanTransformerInferTeaCaching
 from safetensors import safe_open
 import lightx2v.attentions.distributed.ulysses.wrap as ulysses_dist_wrap
@@ -57,9 +57,6 @@ class WanAudioModel(WanModel):
             _, c, h, w = self.scheduler.latents.shape
             num_frame = c + 1  # for r2v
             video_token_num = num_frame * (h // 2) * (w // 2)
-            from loguru import logger
-
-            logger.info(f"video_token_num: {video_token_num}, num_frame: {num_frame}")
             self.transformer_infer.mask_map = MaskMap(video_token_num, num_frame)
 
         embed, grid_sizes, pre_infer_out, valid_patch_length = self.pre_infer.infer(self.pre_weight, inputs, positive=True)

lightx2v/models/networks/wan/infer/pre_wan_audio_infer.py → lightx2v/models/networks/wan/infer/audio/pre_wan_audio_infer.py

 import torch
 import math
-from .utils import rope_params, sinusoidal_embedding_1d
+from ..utils import rope_params, sinusoidal_embedding_1d
 from lightx2v.models.networks.wan.infer.pre_infer import WanPreInfer
 from loguru import logger
 

scripts/wan/run_wan_i2v_audio.sh

@@ -29,6 +29,10 @@ export ENABLE_PROFILING_DEBUG=true
 export ENABLE_GRAPH_MODE=false
 export DTYPE=BF16
 
+
+export TORCH_CUDA_ARCH_LIST="9.0"
+export PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True
+
 python -m lightx2v.infer \
 --model_cls wan2.1_audio \
 --task i2v \