Force-patch GPT data module for worker 0

app.py

@@ -66,7 +66,7 @@ MODEL_PATTERNS = [
     "*.safetensors",
     "*.model",
 ]
-UPSTREAM_PATCH_VERSION = "2026-05-25-worker0"
+UPSTREAM_PATCH_VERSION = "2026-05-25-worker0-v2"
 
 logging.basicConfig(level=logging.INFO, format="%(asctime)s [%(levelname)s] %(message)s")
 log = logging.getLogger(__name__)
@@ -651,18 +651,60 @@ def patch_upstream_repo():
         sv_content = sv_content.replace(old_load, new_load, 1)
         sv_script.write_text(sv_content, encoding="utf-8")
     data_module = GPT_SOVITS_DIR / "GPT_SoVITS" / "AR" / "data" / "data_module.py"
-    data_module_content = data_module.read_text(encoding="utf-8")
-    old_train_loader = """        return DataLoader(
-            self._train_dataset,
-            batch_size=batch_size,
-            sampler=sampler,
-            collate_fn=self._train_dataset.collate,
-            num_workers=self.num_workers,
-            persistent_workers=True,
-            prefetch_factor=16,
+    data_module.write_text(
+        """# modified from https://github.com/yangdongchao/SoundStorm/blob/master/soundstorm/s1/AR/data/data_module.py
+# reference: https://github.com/lifeiteng/vall-e
+from pytorch_lightning import LightningDataModule
+from torch.utils.data import DataLoader
+
+from AR.data.bucket_sampler import DistributedBucketSampler
+from AR.data.dataset import Text2SemanticDataset
+
+
+class Text2SemanticDataModule(LightningDataModule):
+    def __init__(
+        self,
+        config,
+        train_semantic_path,
+        train_phoneme_path,
+        dev_semantic_path=None,
+        dev_phoneme_path=None,
+    ):
+        super().__init__()
+        self.config = config
+        self.train_semantic_path = train_semantic_path
+        self.train_phoneme_path = train_phoneme_path
+        self.dev_semantic_path = dev_semantic_path
+        self.dev_phoneme_path = dev_phoneme_path
+        self.num_workers = self.config["data"]["num_workers"]
+
+    def prepare_data(self):
+        pass
+
+    def setup(self, stage=None, output_logs=False):
+        self._train_dataset = Text2SemanticDataset(
+            phoneme_path=self.train_phoneme_path,
+            semantic_path=self.train_semantic_path,
+            max_sec=self.config["data"]["max_sec"],
+            pad_val=self.config["data"]["pad_val"],
         )
-"""
-    new_train_loader = """        loader_kwargs = dict(
+        self._dev_dataset = self._train_dataset
+        # self._dev_dataset = Text2SemanticDataset(
+        #     phoneme_path=self.dev_phoneme_path,
+        #     semantic_path=self.dev_semantic_path,
+        #     max_sample=self.config['data']['max_eval_sample'],
+        #     max_sec=self.config['data']['max_sec'],
+        #     pad_val=self.config['data']['pad_val'])
+
+    def train_dataloader(self):
+        batch_size = (
+            self.config["train"]["batch_size"] // 2
+            if self.config["train"].get("if_dpo", False) is True
+            else self.config["train"]["batch_size"]
+        )
+        batch_size = max(min(batch_size, len(self._train_dataset) // 4), 1)  # 防止不保存
+        sampler = DistributedBucketSampler(self._train_dataset, batch_size=batch_size)
+        loader_kwargs = dict(
             batch_size=batch_size,
             sampler=sampler,
             collate_fn=self._train_dataset.collate,
@@ -672,20 +714,9 @@ def patch_upstream_repo():
             loader_kwargs["persistent_workers"] = True
             loader_kwargs["prefetch_factor"] = 16
         return DataLoader(self._train_dataset, **loader_kwargs)
-"""
-    if old_train_loader in data_module_content:
-        data_module_content = data_module_content.replace(old_train_loader, new_train_loader, 1)
-    old_val_loader = """        return DataLoader(
-            self._dev_dataset,
-            batch_size=1,
-            shuffle=False,
-            collate_fn=self._train_dataset.collate,
-            num_workers=max(self.num_workers, 12),
-            persistent_workers=True,
-            prefetch_factor=16,
-        )
-"""
-    new_val_loader = """        num_workers = self.num_workers
+
+    def val_dataloader(self):
+        num_workers = self.num_workers
         loader_kwargs = dict(
             batch_size=1,
             shuffle=False,
@@ -696,10 +727,18 @@ def patch_upstream_repo():
             loader_kwargs["persistent_workers"] = True
             loader_kwargs["prefetch_factor"] = 16
         return DataLoader(self._dev_dataset, **loader_kwargs)
-"""
-    if old_val_loader in data_module_content:
-        data_module_content = data_module_content.replace(old_val_loader, new_val_loader, 1)
-    data_module.write_text(data_module_content, encoding="utf-8")
+
+    # 这个会使用到嘛？
+    def test_dataloader(self):
+        return DataLoader(
+            self._dev_dataset,
+            batch_size=1,
+            shuffle=False,
+            collate_fn=self._train_dataset.collate,
+        )
+""",
+        encoding="utf-8",
+    )
     patch_marker.write_text(f"{UPSTREAM_PATCH_VERSION}\n", encoding="utf-8")