Fix/wan2 2 vae encode api (#244)

* bugfix:adapt to  5B dit model, derive attention_head_dim from config[dim]

* [Fix] Wan2.2 Vae Encode refactor: drop args parameter and use self.cpu_offload

lightx2v/models/runners/wan/wan_audio_runner.py

@@ -512,7 +512,7 @@ class WanAudioRunner(WanRunner):  # type:ignore
 
     def load_audio_adapter(self):
         audio_adapter = AudioAdapter(
-            attention_head_dim=5120 // self.config["num_heads"],
+            attention_head_dim=self.config["dim"] // self.config["num_heads"],
             num_attention_heads=self.config["num_heads"],
             base_num_layers=self.config["num_layers"],
             interval=1,

lightx2v/models/video_encoders/hf/wan/vae_2_2.py

@@ -985,11 +985,11 @@ class Wan2_2_VAE:
         self.inv_std = self.inv_std.cuda()
         self.scale = [self.mean, self.inv_std]
 
-    def encode(self, videos, args):
-        if hasattr(args, "cpu_offload") and args.cpu_offload:
+    def encode(self, videos):
+        if self.cpu_offload:
             self.to_cuda()
         out = self.model.encode(videos.unsqueeze(0), self.scale).float().squeeze(0)
-        if hasattr(args, "cpu_offload") and args.cpu_offload:
+        if self.cpu_offload:
             self.to_cpu()
         return out