Update infer/lib/train/process_ckpt.py

infer/lib/train/process_ckpt.py

@@ -10,10 +10,53 @@ from i18n.i18n import I18nAuto
 i18n = I18nAuto()
 
 
-def savee(ckpt, sr, if_f0, name, epoch, version, hps):
+# def savee(ckpt, sr, if_f0, name, epoch, version, hps):
+#     try:
+#         opt = OrderedDict()
+#         opt["weight"] = {}
+#         for key in ckpt.keys():
+#             if "enc_q" in key:
+#                 continue
+#             opt["weight"][key] = ckpt[key].half()
+#         opt["config"] = [
+#             hps.data.filter_length // 2 + 1,
+#             32,
+#             hps.model.inter_channels,
+#             hps.model.hidden_channels,
+#             hps.model.filter_channels,
+#             hps.model.n_heads,
+#             hps.model.n_layers,
+#             hps.model.kernel_size,
+#             hps.model.p_dropout,
+#             hps.model.resblock,
+#             hps.model.resblock_kernel_sizes,
+#             hps.model.resblock_dilation_sizes,
+#             hps.model.upsample_rates,
+#             hps.model.upsample_initial_channel,
+#             hps.model.upsample_kernel_sizes,
+#             hps.model.spk_embed_dim,
+#             hps.model.gin_channels,
+#             hps.data.sampling_rate,
+#         ]
+#         opt["info"] = "%sepoch" % epoch
+#         opt["sr"] = sr
+#         opt["f0"] = if_f0
+#         opt["version"] = version
+#         torch.save(opt, "assets/weights/%s.pth" % name)
+          
+#         return "Success."
+#     except:
+#         return traceback.format_exc()
+
+import os
+import torch
+from google.oauth2 import service_account
+from googleapiclient.discovery import build
+from googleapiclient.http import MediaFileUpload
+
+def savee(ckpt, sr, if_f0, name, epoch, version, hps, credentials_path):
     try:
-        print("ckpt, sr, if_f0, name, epoch, version, hps", ckpt, sr, if_f0, name, epoch, version, hps)
-        opt = OrderedDict()
+        opt = {}
         opt["weight"] = {}
         for key in ckpt.keys():
             if "enc_q" in key:
@@ -43,11 +86,39 @@ def savee(ckpt, sr, if_f0, name, epoch, version, hps):
         opt["sr"] = sr
         opt["f0"] = if_f0
         opt["version"] = version
-        torch.save(opt, "assets/weights/%s.pth" % name)
 
+        # Save the checkpoint to a local file
+        checkpoint_file = "assets/weights/%s.pth" % name
+        torch.save(opt, checkpoint_file)
+
+        # Set up Google Drive API and upload the file
+        SCOPES = ['https://www.googleapis.com/auth/drive.file']
+        creds = service_account.Credentials.from_service_account_file(credentials_path, scopes=SCOPES)
+
+        drive_service = build('drive', 'v3', credentials=creds)
+
+        # Specify the file name for Google Drive
+        drive_file_name = 'your_drive_file_name_here'
+
+        # Create a MediaFileUpload object to upload the file
+        media = MediaFileUpload(checkpoint_file, resumable=True)
+
+        # Create a file on Google Drive
+        file_metadata = {'name': drive_file_name}
+        file = drive_service.files().create(body=file_metadata, media_body=media, fields='id').execute()
+
+        # Make the uploaded file publicly accessible (optional)
+        drive_service.permissions().create(
+            fileId=file['id'],
+            body={'type': 'anyone', 'role': 'reader'}
+        ).execute()
+
+        print(f'File ID: {file.get("id")}')
+        print(file)
         return "Success."
-    except:
-        return traceback.format_exc()
+    except Exception as e:
+        print("Error:", str(e))
+
 
 
 def show_info(path):