From ff1feb611fc31d30b3cf57017b1bd65f53dc03df Mon Sep 17 00:00:00 2001
From: vincentchen <chenlufang@x2robot.com>
Date: Mon, 8 Sep 2025 21:38:24 +0800
Subject: [PATCH] fix load model

---
 wall_x/data/load_lerobot_dataset.py                   |  2 +-
 wall_x/model/qwen2_5_based/modeling_qwen2_5_vl_act.py | 11 ++++++++---
 wall_x/trainer/qwen_vl_act_trainer.py                 |  2 +-
 workspace/lerobot_example/config_qact.yml             | 10 ++++------
 4 files changed, 14 insertions(+), 11 deletions(-)

diff --git a/wall_x/data/load_lerobot_dataset.py b/wall_x/data/load_lerobot_dataset.py
index d502417..344d567 100644
--- a/wall_x/data/load_lerobot_dataset.py
+++ b/wall_x/data/load_lerobot_dataset.py
@@ -207,7 +207,7 @@ class DataCollator:
         self.load_processor()
 
     def load_processor(self):
-        processor_path = self.config["processor_path"]
+        processor_path = self.config["pretrained_qwen_vl_path"]
         action_tokenizer_path = self.config["action_tokenizer_path"]
 
         # Use cached processors if available
diff --git a/wall_x/model/qwen2_5_based/modeling_qwen2_5_vl_act.py b/wall_x/model/qwen2_5_based/modeling_qwen2_5_vl_act.py
index 1a0d65b..cf0edb8 100644
--- a/wall_x/model/qwen2_5_based/modeling_qwen2_5_vl_act.py
+++ b/wall_x/model/qwen2_5_based/modeling_qwen2_5_vl_act.py
@@ -1,6 +1,7 @@
 import os
 import torch
 import numpy as np
+import glob
 import torch.nn as nn
 from torchdiffeq import odeint
 from dataclasses import dataclass
@@ -685,7 +686,6 @@ class Qwen2_5_VLMoEForAction(Qwen2_5_VLForConditionalGeneration):
         """
         
         # Load model components from pretrained path
-        model_path = os.path.join(pretrained_model_path, "model.safetensors")
         config_path = os.path.join(pretrained_model_path, "config.json")
         config = cls.config_class.from_pretrained(config_path)
         processor = AutoProcessor.from_pretrained(pretrained_model_path, use_fast=True)
@@ -703,8 +703,13 @@ class Qwen2_5_VLMoEForAction(Qwen2_5_VLForConditionalGeneration):
         model.resize_token_embeddings(len(processor.tokenizer))
         
         # Load model state dict from safetensors file
-        state_dict = load_file(model_path, device="cpu")
-        msg = model.load_state_dict(state_dict, strict=False)
+        safetensor_files = glob.glob(os.path.join(pretrained_model_path, "*.safetensors"))
+        state_dict = {}
+        for file in safetensor_files:
+            sd = load_file(file, device="cpu")
+            state_dict.update(sd)
+
+        model.load_state_dict(state_dict, strict=False)
         
         return model
 
diff --git a/wall_x/trainer/qwen_vl_act_trainer.py b/wall_x/trainer/qwen_vl_act_trainer.py
index f3301ea..05dc648 100644
--- a/wall_x/trainer/qwen_vl_act_trainer.py
+++ b/wall_x/trainer/qwen_vl_act_trainer.py
@@ -108,7 +108,7 @@ class QwenVlAct_Trainer:
             ValueError: If required configuration keys are missing
         """
         # Validate required configuration keys
-        required_keys = ["processor_path", "qwen_vl_act_config_path", "learning_rate", "num_epoch"]
+        required_keys = ["learning_rate", "num_epoch"]
         for key in required_keys:
             if key not in config:
                 raise ValueError(f"Missing required configuration key: {key}")
diff --git a/workspace/lerobot_example/config_qact.yml b/workspace/lerobot_example/config_qact.yml
index 4fa62cd..fa683ae 100644
--- a/workspace/lerobot_example/config_qact.yml
+++ b/workspace/lerobot_example/config_qact.yml
@@ -5,9 +5,7 @@
 log_name: "robotic_training"
 log_project: "vla_training"
 model_type: qwen2_5
-processor_path: "/path/to/model/"
 pretrained_qwen_vl_path: "/path/to/qwen_vl_model/"
-qwen_vl_act_config_path: "/path/to/config.json"
 action_tokenizer_path: "/path/to/fast/"
 save_path: "/path/to/workspace/"
 
@@ -52,10 +50,10 @@ agent_pos_config:
   height: 1
   car_pose: 3
 
-# Checkpoint resuming configuration
-resume:
-  ckpt: "/path/to/resume_model/"
-  load_ckpt_only: true
+# # Checkpoint resuming configuration
+# resume:
+#   ckpt: "/path/to/resume_model/"
+#   load_ckpt_only: true
 
 # Data configuration
 data: