From 787b44807a41c8eebd685fd71001acc64149d9db Mon Sep 17 00:00:00 2001
From: mrq <barry.quiggles@protonmail.com>
Date: Tue, 28 Feb 2023 15:36:06 +0000
Subject: [PATCH] added to embedded metadata: datetime, model path, model hash

---
 src/utils.py | 64 ++++++++++++++++++++++++++++++++++++++++++++--------
 src/webui.py |  7 +++---
 2 files changed, 59 insertions(+), 12 deletions(-)

diff --git a/src/utils.py b/src/utils.py
index 14aa533..7896f1d 100755
--- a/src/utils.py
+++ b/src/utils.py
@@ -307,6 +307,10 @@ def generate(
 		'cond_free_k': cond_free_k,
 		'experimentals': experimental_checkboxes,
 		'time': time.time()-full_start_time,
+
+		'datetime': datetime.now().isoformat(),
+		'model': tts.autoregressive_model_path,
+		'model_hash': tts.autoregressive_model_hash if hasattr(tts, 'autoregressive_model_hash') else None,
 	}
 
 	"""
@@ -324,6 +328,7 @@ def generate(
 
 	if args.voice_fixer:
 		if not voicefixer:
+			progress(0, "Loading voicefix...")
 			load_voicefixer()
 
 		fixed_cache = {}
@@ -1006,7 +1011,33 @@ def get_voice_list(dir=get_voice_dir(), append_defaults=False):
 		res = res + ["random", "microphone"]
 	return res
 
-def get_autoregressive_models(dir="./models/finetunes/"):
+def hash_file(path, algo="md5", buffer_size=0):
+	import hashlib
+
+	hash = None
+	if algo == "md5":
+		hash = hashlib.md5()
+	elif algo == "sha1":
+		hash = hashlib.sha1()
+	else:
+		raise Exception(f'Unknown hash algorithm specified: {algo}')
+
+	if not os.path.exists(path):
+		raise Exception(f'Path not found: {path}')
+
+	with open(path, 'rb') as f:
+		if buffer_size > 0:
+			while True:
+				data = f.read(buffer_size)
+				if not data:
+					break
+				hash.update(data)
+		else:
+			hash.update(f.read())
+
+	return "{0}".format(hash.hexdigest())
+
+def get_autoregressive_models(dir="./models/finetunes/", prefixed=False):
 	os.makedirs(dir, exist_ok=True)
 	base = [get_model_path('autoregressive.pth')]
 	halfp = get_halfp_model_path()
@@ -1018,12 +1049,20 @@ def get_autoregressive_models(dir="./models/finetunes/"):
 	for training in os.listdir(f'./training/'):
 		if not os.path.isdir(f'./training/{training}/') or not os.path.isdir(f'./training/{training}/models/'):
 			continue
-		#found = found + sorted([ f'./training/{training}/models/{d}' for d in os.listdir(f'./training/{training}/models/') if d[-8:] == "_gpt.pth" ])
 		models = sorted([ int(d[:-8]) for d in os.listdir(f'./training/{training}/models/') if d[-8:] == "_gpt.pth" ])
 		found = found + [ f'./training/{training}/models/{d}_gpt.pth' for d in models ]
-		#found.append(f'./training/{training}/models/{models[-1]}_gpt.pth')
 
-	return base + additionals + found
+	res = base + additionals + found
+	
+	if prefixed:
+		for i in range(len(res)):
+			path = res[i]
+			hash = hash_file(path)
+			shorthash = hash[:8]
+
+			res[i] = f'[{shorthash}] {path}'
+
+	return res
 
 def get_dataset_list(dir="./training/"):
 	return sorted([d for d in os.listdir(dir) if os.path.isdir(os.path.join(dir, d)) and len(os.listdir(os.path.join(dir, d))) > 0 and "train.txt" in os.listdir(os.path.join(dir, d)) ])
@@ -1256,8 +1295,6 @@ def save_args_settings():
 		'training-default-bnb': args.training_default_bnb,
 	}
 
-	print(settings)
-
 	os.makedirs('./config/', exist_ok=True)
 	with open(f'./config/exec.json', 'w', encoding="utf-8") as f:
 		f.write(json.dumps(settings, indent='\t') )
@@ -1322,9 +1359,7 @@ def read_generate_settings(file, read_latents=True):
 	except Exception as e:
 		pass
 
-	if j is None:
-		print("No metadata found in audio file to read")
-	else:
+	if j is not None:
 		if 'latents' in j:
 			if read_latents:
 				latents = base64.b64decode(j['latents'])
@@ -1360,6 +1395,10 @@ def load_tts( restart=False, model=None ):
 	except Exception as e:
 		tts = TextToSpeech(minor_optimizations=not args.low_vram)
 		load_autoregressive_model(args.autoregressive_model)
+
+	if not hasattr(tts, 'autoregressive_model_hash'):
+		tts.autoregressive_model_hash = hash_file(tts.autoregressive_model_path)
+
 	tts_loading = False
 
 	get_model_path('dvae.pth')
@@ -1381,6 +1420,10 @@ def reload_tts( model=None ):
 	load_tts( restart=True, model=model )
 
 def update_autoregressive_model(autoregressive_model_path):
+	match = re.findall(r'^\[[a-fA-F0-9]{8}\] (.+?)$', autoregressive_model_path)
+	if match:
+		autoregressive_model_path = match[0]
+
 	if not autoregressive_model_path or not os.path.exists(autoregressive_model_path):
 		print(f"Invalid model: {autoregressive_model_path}")
 		return
@@ -1416,6 +1459,9 @@ def update_autoregressive_model(autoregressive_model_path):
 		if tts.preloaded_tensors:
 			tts.autoregressive = tts.autoregressive.to(tts.device)
 
+	if not hasattr(tts, 'autoregressive_model_hash'):
+		tts.autoregressive_model_hash = hash_file(autoregressive_model_path)
+
 	print(f"Loaded model: {tts.autoregressive_model_path}")
 
 	do_gc()
diff --git a/src/webui.py b/src/webui.py
index b4c548f..aed451d 100755
--- a/src/webui.py
+++ b/src/webui.py
@@ -129,6 +129,9 @@ history_headers = {
 	"Rep Pen": "repetition_penalty",
 	"Cond-Free K": "cond_free_k",
 	"Time": "time",
+	"Datetime": "datetime",
+	"Model": "model",
+	"Model Hash": "model_hash",
 }
 
 def history_view_results( voice ):
@@ -147,7 +150,7 @@ def history_view_results( voice ):
 		for k in history_headers:
 			v = file
 			if k != "Name":
-				v = metadata[history_headers[k]]
+				v = metadata[history_headers[k]] if history_headers[k] in metadata else '?'
 			values.append(v)
 
 
@@ -174,8 +177,6 @@ def read_generate_settings_proxy(file, saveAs='.temp'):
 		
 		latents = f'{outdir}/cond_latents.pth'
 
-	print(j, latents)
-
 	return (
 		gr.update(value=j, visible=j is not None),
 		gr.update(visible=j is not None),