Update train.py

train.py

@@ -30,7 +30,7 @@ from so_vits_svc_fork.modules.losses import discriminator_loss, feature_loss, ge
 from so_vits_svc_fork.modules.mel_processing import mel_spectrogram_torch
 from so_vits_svc_fork.modules.synthesizers import SynthesizerTrn
 
-from so_vits_svc_fork.train import VitsLightning
+from so_vits_svc_fork.train import VitsLightning, VCDataModule
 
 LOG = getLogger(__name__)
 torch.set_float32_matmul_precision("high")
@@ -68,37 +68,15 @@ class HuggingFacePushCallback(pl.Callback):
             commit_message="🍻 cheers",
             ignore_patterns=["*.git*", "*README.md*", "*__pycache__*"],
         )
-
-class VCDataModule(pl.LightningDataModule):
-    batch_size: int
-
-    def __init__(self, hparams: Any):
-        super().__init__()
-        self.__hparams = hparams
-        self.batch_size = hparams.train.batch_size
-        if not isinstance(self.batch_size, int):
-            self.batch_size = 1
-        self.collate_fn = TextAudioCollate()
-
-        # these should be called in setup(), but we need to calculate check_val_every_n_epoch
-        self.train_dataset = TextAudioDataset(self.__hparams, is_validation=False)
-        self.val_dataset = TextAudioDataset(self.__hparams, is_validation=True)
-
-    def train_dataloader(self):
-        return DataLoader(
-            self.train_dataset,
-            num_workers=min(cpu_count(), self.__hparams.train.get("num_workers", 8)),
-            batch_size=self.batch_size,
-            collate_fn=self.collate_fn,
-            persistent_workers=self.__hparams.train.get("persistent_workers", True),
-        )
-
-    def val_dataloader(self):
-        return DataLoader(
-            self.val_dataset,
-            batch_size=1,
-            collate_fn=self.collate_fn,
-        )
+        ckpt_pattern = r'^(D_|G_)\d+\.pth$'
+        todelete = []
+        repo_ckpts = [x for x in list_repo_files(self.repo_id) if re.match(ckpt_pattern, x) and x not in ["G_0.pth", "D_0.pth"]]
+        local_ckpts = [x.name for x in Path(model_dir).glob("*.pth") if re.match(ckpt_pattern, x.name)]
+        to_delete = set(repo_ckpts) - set(local_ckpts)
+
+        for fname in to_delete:
+            print("🗑 Deleting {fname} from repo")
+            delete_file(fname, repo_id)
 
 
 def train(