Improved Train code cell of DFL_colab

train_colab.py

#@title Training
Model = "SAEHD" #@param ["SAEHD", "Quick96", "XSeg"]
Archive_name_backup = "workspace.zip" #@param {type:"string"}
Backup_every_specified_time = False #@param {type:"boolean"}
backup_time = 1 #@param {type:"integer"}
Silent_Start = True #@param {type:"boolean"}

%cd "/content"

#Mount Google Drive as folder
from google.colab import drive
drive.mount('/content/drive')

import psutil, os, time

p = psutil.Process(os.getpid())
uptime = time.time() - p.create_time()

if (Backup_every_specified_time):
  if not os.path.exists(Archive_name_backup):
    print("Creating workspace archive ...")
    os.system('zip -r -q '+Archive_name_backup+' workspace')
    print("Archive created!")
  else:
    print("Archive exist!")

if (Backup_every_specified_time):
  print("Time to end session: "+str(round((43200-uptime)/3600))+" hours")
  backup_time = str(backup_time * 3600)
  backup_cmd = " --execute-program -"+backup_time+" \"import os; os.system('zip -r -q "+Archive_name_backup+" workspace/model'); os.system('cp /content/"+Archive_name_backup+" /content/drive/My\ Drive/'); print('Backed up!') \"" 
elif (round(39600-uptime) > 0):
  print("Time to backup: "+str(round((39600-uptime)/3600))+" hours")
  backup_time = str(round(39600-uptime))
  backup_cmd = " --execute-program "+backup_time+" \"import os; os.system('zip -r -q "+Archive_name_backup+" workspace'); os.system('cp /content/"+Archive_name_backup+" /content/drive/My\ Drive/'); print('Backed up!') \"" 
else:
  print("Session expires in less than an hour.")
  backup_cmd = ""
    
cmd = "DeepFaceLab/main.py train --training-data-src-dir workspace/data_src/aligned --training-data-dst-dir workspace/data_dst/aligned --pretraining-data-dir pretrain --model-dir workspace/model --model "+Model

if Model == "Quick96":
  cmd+= " --pretrained-model-dir pretrain_Q96"

if Silent_Start:
  cmd+= " --silent-start"

if (backup_cmd != ""):
  train_cmd = (cmd+backup_cmd)
else:
  train_cmd = (cmd)

!python $train_cmd

Brown sugar