Update app.py

app.py

@@ -87,21 +87,49 @@ def gradio_wrapper(inputs):
 	return decorated
 
 # returns a list of models, assuming the models are placed under ./training/ or ./models/ or ./data/models/
-def get_model_paths( paths=[Path("./training/"), Path("./models/"), Path("./data/models/")] ):
+def get_model_paths(paths=["./training/", "./models/", "./data/models/"] ):
 	configs = []
 
 	for path in paths:
+		if not isinstance( path, Path ):
+			path = Path(path)
+
 		if not path.exists():
 			continue
 
 		for yaml in path.glob("**/*.yaml"):
 			if "/logs/" in str(yaml):
 				continue
+			if "lora" in str(yaml):
+				continue
 			configs.append( yaml )
 		
 		for sft in path.glob("**/*.sft"):
 			if "/logs/" in str(sft):
 				continue
+			if "lora" in str(sft):
+				continue
+			configs.append( sft )
+
+	configs = [ str(p) for p in configs ]
+
+	return configs
+
+def get_lora_paths(paths=["./training/", "./models/", "./data/models/"] ):
+	configs = []
+
+	for path in paths:
+		if not isinstance( path, Path ):
+			path = Path(path)
+
+		if not path.exists():
+			continue
+
+		for sft in path.glob("**/*.sft"):
+			if "/logs/" in str(sft):
+				continue
+			if "lora" not in str(sft):
+				continue
 			configs.append( sft )
 
 	configs = [ str(p) for p in configs ]
@@ -115,10 +143,10 @@ def get_attentions():
 	return AVAILABLE_ATTENTIONS + ["auto"]
 
 #@gradio_wrapper(inputs=layout["settings"]["inputs"].keys())
-def load_model( config, device, dtype, attention ):
+def load_model( config, lora, device, dtype, attention ):
 	gr.Info(f"Loading: {config}")
 	try:
-		init_tts( config=Path(config), restart=True, device=device, dtype=dtype, attention=attention )
+		init_tts( config=Path(config), lora=Path(lora) if lora is not None else None, restart=True, device=device, dtype=dtype, attention=attention )
 	except Exception as e:
 		raise gr.Error(e)
 	gr.Info(f"Loaded model")
@@ -130,7 +158,7 @@ def get_languages():
 	return list(get_lang_symmap().keys()) + ["auto"]
 
 def get_tasks():
-	return ["tts", "sr", "nr", "vc"]
+	return ["tts", "sr", "ns", "vc"]
 
 #@gradio_wrapper(inputs=layout["dataset"]["inputs"].keys())
 def load_sample( speaker ):
@@ -219,18 +247,20 @@ def do_inference_tts( progress=gr.Progress(track_tqdm=True), *args, **kwargs ):
 	parser.add_argument("--voice-convert", type=str, default=kwargs["voice-convert"])
 	parser.add_argument("--language", type=str, default=kwargs["language"])
 	parser.add_argument("--text-language", type=str, default=kwargs["text-language"])
+	parser.add_argument("--no-phonemize", action="store_true")
+	parser.add_argument("--play", action="store_true")
 	parser.add_argument("--split-text-by", type=str, default=kwargs["split-text-by"])
 	parser.add_argument("--context-history", type=int, default=kwargs["context-history"])
 	parser.add_argument("--input-prompt-length", type=float, default=kwargs["input-prompt-length"])
-	parser.add_argument("--input-prompt-prefix", action='store_true', default=kwargs["input-prompt-prefix"])
+	#parser.add_argument("--input-prompt-prefix", action='store_true', default=kwargs["input-prompt-prefix"])
 	parser.add_argument("--max-duration", type=int, default=int(kwargs["max-duration"]*cfg.dataset.frames_per_second))
-	parser.add_argument("--max-levels", type=int, default=kwargs["max-levels"])
+	#parser.add_argument("--max-levels", type=int, default=kwargs["max-levels"])
 	parser.add_argument("--max-steps", type=int, default=kwargs["max-steps"])
 	parser.add_argument("--ar-temperature", type=float, default=kwargs["ar-temperature"])
 	parser.add_argument("--nar-temperature", type=float, default=kwargs["nar-temperature"])
 	parser.add_argument("--min-ar-temperature", type=float, default=kwargs["min-ar-temperature"])
 	parser.add_argument("--min-nar-temperature", type=float, default=kwargs["min-nar-temperature"])
-	parser.add_argument("--prefix-silence", type=float, default=kwargs["prefix-silence"])
+	#parser.add_argument("--prefix-silence", type=float, default=kwargs["prefix-silence"])
 	parser.add_argument("--top-p", type=float, default=kwargs["top-p"])
 	parser.add_argument("--top-k", type=int, default=kwargs["top-k"])
 	parser.add_argument("--top-no", type=float, default=kwargs["top-no"])
@@ -238,6 +268,7 @@ def do_inference_tts( progress=gr.Progress(track_tqdm=True), *args, **kwargs ):
 	parser.add_argument("--repetition-penalty", type=float, default=kwargs["repetition-penalty"])
 	parser.add_argument("--repetition-penalty-decay", type=float, default=kwargs["repetition-penalty-decay"])
 	parser.add_argument("--length-penalty", type=float, default=kwargs["length-penalty"])
+	"""
 	parser.add_argument("--beam-width", type=int, default=kwargs["beam-width"])
 	parser.add_argument("--mirostat-tau", type=float, default=kwargs["mirostat-tau"])
 	parser.add_argument("--mirostat-eta", type=float, default=kwargs["mirostat-eta"])
@@ -249,10 +280,16 @@ def do_inference_tts( progress=gr.Progress(track_tqdm=True), *args, **kwargs ):
 	parser.add_argument("--layer-skip-exit-layer", type=int, default=kwargs["layer-skip-exit-layer"])
 	parser.add_argument("--layer-skip-entropy-threshold", type=int, default=kwargs["layer-skip-entropy-threshold"])
 	parser.add_argument("--layer-skip-varentropy-threshold", type=int, default=kwargs["layer-skip-varentropy-threshold"])
+	"""
 	parser.add_argument("--refine-on-stop", action="store_true")
 	parser.add_argument("--denoise-start", type=float, default=0.0)
 	parser.add_argument("--cfg-strength", type=float, default=kwargs['cfg-strength'])
 	parser.add_argument("--cfg-rescale", type=float, default=kwargs['cfg-rescale'])
+	
+	parser.add_argument("--sampling-scores-masked-only", action="store_true")
+	parser.add_argument("--sampling-scores-flatten", action="store_true")
+	parser.add_argument("--sampling-scores-remask", action="store_true")
+
 	args, unknown = parser.parse_known_args()
 
 	if is_windows:
@@ -274,6 +311,21 @@ def do_inference_tts( progress=gr.Progress(track_tqdm=True), *args, **kwargs ):
 	if kwargs.pop("refine-on-stop", False):
 		args.refine_on_stop = True
 
+	if kwargs.pop("no-phonemize", False):
+		args.no_phonemize = True
+	
+	if kwargs.pop("play", False):
+		args.play = True
+	
+	if kwargs.pop("sampling-scores-masked-only", False):
+		args.sampling_scores_masked_only = True
+	
+	if kwargs.pop("sampling-scores-flatten", False):
+		args.sampling_scores_flatten = True
+	
+	if kwargs.pop("sampling-scores-remask", False):
+		args.sampling_scores_remask = True
+
 	if args.split_text_by == "lines":
 		args.split_text_by = "\n"
 	elif args.split_text_by == "none":
@@ -289,30 +341,35 @@ def do_inference_tts( progress=gr.Progress(track_tqdm=True), *args, **kwargs ):
 	sampling_kwargs = dict(
 		split_text_by=args.split_text_by,
 		context_history=args.context_history,
+		phonemize=not args.no_phonemize,
 		voice_convert=args.voice_convert,
 		max_steps=args.max_steps,
-		max_levels=args.max_levels,
+		#max_levels=args.max_levels,
 		max_duration=args.max_duration,
 		ar_temperature=args.ar_temperature, nar_temperature=args.nar_temperature,
 		min_ar_temperature=args.min_ar_temperature, min_nar_temperature=args.min_nar_temperature,
 		top_p=args.top_p, top_k=args.top_k, min_p=args.min_p, top_no=args.top_no,
 		repetition_penalty=args.repetition_penalty, repetition_penalty_decay=args.repetition_penalty_decay,
 		length_penalty=args.length_penalty,
-		beam_width=args.beam_width,
-		mirostat_tau=args.mirostat_tau, mirostat_eta=args.mirostat_eta,
-		dry_multiplier=args.dry_multiplier, dry_base=args.dry_base, dry_allowed_length=args.dry_allowed_length,
-		entropix_sampling=args.entropix_sampling,
-		layer_skip=args.layer_skip,
-		layer_skip_exit_layer=args.layer_skip_exit_layer,
-		layer_skip_entropy_threshold=args.layer_skip_entropy_threshold,
-		layer_skip_varentropy_threshold=args.layer_skip_varentropy_threshold,
-		refine_on_stop=args.refine_on_stop,
+		#beam_width=args.beam_width,
+		#mirostat_tau=args.mirostat_tau, mirostat_eta=args.mirostat_eta,
+		#dry_multiplier=args.dry_multiplier, dry_base=args.dry_base, dry_allowed_length=args.dry_allowed_length,
+		#entropix_sampling=args.entropix_sampling,
+		#layer_skip=args.layer_skip,
+		#layer_skip_exit_layer=args.layer_skip_exit_layer,
+		#layer_skip_entropy_threshold=args.layer_skip_entropy_threshold,
+		#layer_skip_varentropy_threshold=args.layer_skip_varentropy_threshold,
+		#refine_on_stop=args.refine_on_stop,
 		denoise_start=args.denoise_start,
-		prefix_silence=args.prefix_silence,
-		input_prompt_prefix=args.input_prompt_prefix,
+		#prefix_silence=args.prefix_silence,
+		#input_prompt_prefix=args.input_prompt_prefix,
 		input_prompt_length=args.input_prompt_length,
 		cfg_strength=args.cfg_strength,
 		cfg_rescale=args.cfg_rescale,
+
+		sampling_scores_masked_only=args.sampling_scores_masked_only,
+		sampling_scores_flatten=args.sampling_scores_flatten,
+		sampling_scores_remask=args.sampling_scores_remask,
 	)
 
 	with timer("Inferenced in", callback=lambda msg: gr.Info( msg )) as t:
@@ -321,6 +378,7 @@ def do_inference_tts( progress=gr.Progress(track_tqdm=True), *args, **kwargs ):
 			language=args.language,
 			text_language=args.text_language,
 			task=args.task,
+			play=args.play,
 			modality=args.modality.lower(),
 			references=args.references.split(";") if args.references is not None else [],
 			**sampling_kwargs,
@@ -416,6 +474,7 @@ def do_training( progress=gr.Progress(track_tqdm=True), *args, **kwargs ):
 parser = argparse.ArgumentParser(allow_abbrev=False)
 parser.add_argument("--yaml", type=Path, default=os.environ.get('VALLE_YAML', None)) # os environ so it can be specified in a HuggingFace Space too
 parser.add_argument("--model", type=Path, default=os.environ.get('VALLE_MODEL', None)) # os environ so it can be specified in a HuggingFace Space too
+parser.add_argument("--lora", type=Path, default=os.environ.get('VALLE_LORA', None)) # os environ so it can be specified in a HuggingFace Space too
 parser.add_argument("--listen", default=None, help="Path for Gradio to listen on")
 parser.add_argument("--share", action="store_true")
 parser.add_argument("--render_markdown", action="store_true", default="VALLE_YAML" in os.environ)
@@ -469,6 +528,9 @@ with ui:
 						with gr.Row():
 							layout["inference_tts"]["inputs"]["split-text-by"] = gr.Dropdown(choices=["sentences", "lines"], label="Text Delimiter", info="How to split the text into utterances.", value="sentences")
 							layout["inference_tts"]["inputs"]["context-history"] = gr.Slider(value=0, minimum=0, maximum=4, step=1, label="(Rolling) Context History", info="How many prior lines to serve as the context/prefix (0 to disable).")
+						with gr.Row():
+							layout["inference_tts"]["inputs"]["no-phonemize"] = gr.Checkbox(label="No Phonemize", info="Use raw text rather than phonemize the text as the input prompt.")
+							layout["inference_tts"]["inputs"]["play"] = gr.Checkbox(label="Auto Play", info="Auto play on generation (using sounddevice).")
 					with gr.Tab("Sampler Settings"):
 						with gr.Row():
 							layout["inference_tts"]["inputs"]["ar-temperature"] = gr.Slider(value=1.0, minimum=0.0, maximum=1.5, step=0.05, label="Temperature (AR/NAR-len)", info="Adjusts the probabilities in the AR/NAR-len. (0 to greedy* sample)")
@@ -486,7 +548,12 @@ with ui:
 							layout["inference_tts"]["inputs"]["repetition-penalty"] = gr.Slider(value=1.0, minimum=0.0, maximum=5.0, step=0.05, label="Repetition Penalty", info="Incurs a penalty to tokens based on how often they appear in a sequence.")
 							layout["inference_tts"]["inputs"]["repetition-penalty-decay"] = gr.Slider(value=0.0, minimum=-2.0, maximum=2.0, step=0.05, label="Repetition Penalty Length Decay", info="Modifies the reptition penalty based on how far back in time the token appeared in the sequence.")
 							layout["inference_tts"]["inputs"]["length-penalty"] = gr.Slider(value=0.0, minimum=-2.0, maximum=2.0, step=0.05, label="Length Penalty", info="(AR only) Modifies the probability of a stop token based on the current length of the sequence.")
+						with gr.Row():
+							layout["inference_tts"]["inputs"]["sampling-scores-masked-only"] = gr.Checkbox(label="Sampled Scores: Masked Only", info="(NAR-len only) Update scores for newly generated tokens only")
+							layout["inference_tts"]["inputs"]["sampling-scores-flattened"] = gr.Checkbox(label="Sampled Scores: Flattened", info="(NAR-len only) Flattens the scores for all codebook levels")
+							layout["inference_tts"]["inputs"]["sampling-scores-remask"] = gr.Checkbox(label="Sampled Scores: Remask", info="(NAR-len only) Remasks P%% of existing tokens randomly after each step.")
 					# These settings are pretty much not supported anyways
+					"""
 					with gr.Tab("Experimental Settings", visible=cfg.experimental):
 						with gr.Row():
 							layout["inference_tts"]["inputs"]["max-levels"] = gr.Slider(value=7, minimum=0, maximum=7, step=1, label="Max NAR Levels", info="Limits how many steps to perform in the NAR pass.")
@@ -509,6 +576,7 @@ with ui:
 							layout["inference_tts"]["inputs"]["layer-skip-exit-layer"] = gr.Slider(value=11, minimum=0, maximum=11, step=1, label="Layer Skip Exit Layer", info="Maximum model layer to exit early from.")
 							layout["inference_tts"]["inputs"]["layer-skip-entropy-threshold"] = gr.Slider(value=0.1, minimum=0, maximum=1.0, step=0.01, label="Layer Skip Entropy Threshold", info="Entropy threshold for early-exit")
 							layout["inference_tts"]["inputs"]["layer-skip-varentropy-threshold"] = gr.Slider(value=0.1, minimum=0, maximum=1.0, step=0.01, label="Layer Skip Varentropy Threshold", info="Varentropy threshold for early-exit")
+					"""
 
 		layout["inference_tts"]["buttons"]["inference"].click(
 			fn=do_inference_tts,
@@ -554,6 +622,7 @@ with ui:
 							layout["inference_stt"]["inputs"]["repetition-penalty"] = gr.Slider(value=1.0, minimum=-2.0, maximum=2.0, step=0.05, label="Repetition Penalty", info="Incurs a penalty to tokens based on how often they appear in a sequence.")
 							layout["inference_stt"]["inputs"]["repetition-penalty-decay"] = gr.Slider(value=0.0, minimum=-2.0, maximum=2.0, step=0.05, label="Repetition Penalty Length Decay", info="Modifies the reptition penalty based on how far back in time the token appeared in the sequence.")
 							layout["inference_stt"]["inputs"]["length-penalty"] = gr.Slider(value=0.0, minimum=-2.0, maximum=2.0, step=0.05, label="Length Penalty", info="(AR only) Modifies the probability of a stop token based on the current length of the sequence.")
+						"""
 						with gr.Row():
 							layout["inference_stt"]["inputs"]["dynamic-sampling"] = gr.Checkbox(label="Dynamic Temperature", info="Dynamically adjusts the temperature based on the highest confident predicted token per sampling step.")
 							layout["inference_stt"]["inputs"]["mirostat-tau"] = gr.Slider(value=0.0, minimum=0.0, maximum=8.0, step=0.05, label="Mirostat τ (Tau)", info="The \"surprise\" value when performing mirostat sampling. 0 to disable.")
@@ -562,6 +631,7 @@ with ui:
 							layout["inference_stt"]["inputs"]["dry-multiplier"] = gr.Slider(value=0.0, minimum=0.0, maximum=8.0, step=0.05, label="DRY Multiplier", info="The multiplying factor for the DRY score penalty (0 to disable DRY sampling).")
 							layout["inference_stt"]["inputs"]["dry-base"] = gr.Slider(value=1.75, minimum=0.0, maximum=8.0, step=0.05, label="DRY Base", info="The base of the exponent in the DRY score penalty")
 							layout["inference_stt"]["inputs"]["dry-allowed-length"] = gr.Slider(value=2, minimum=0, maximum=75, step=1, label="Allowed Length", info="The maximimum length a token can be to perform DRY penalty with.")
+						"""
 
 		layout["inference_stt"]["buttons"]["inference"].click(
 			fn=do_inference_stt,
@@ -609,8 +679,9 @@ with ui:
 				with gr.Column(scale=7):
 					with gr.Row():
 						layout["settings"]["inputs"]["models"] = gr.Dropdown(choices=get_model_paths(), value=args.yaml or args.model, label="Model", info="Model to load. Can load from a config YAML or the weights itself.")
-						layout["settings"]["inputs"]["device"] = gr.Dropdown(choices=get_devices(), value="cuda:0", label="Device", info="Device to load the weights onto.")
+						layout["settings"]["inputs"]["loras"] = gr.Dropdown(choices=get_lora_paths(), value=args.yaml or args.lora, label="LoRA", info="LoRA to load. Can load from a config YAML or the weights itself.")
 					with gr.Row():
+						layout["settings"]["inputs"]["device"] = gr.Dropdown(choices=get_devices(), value="cuda:0", label="Device", info="Device to load the weights onto.")
 						layout["settings"]["inputs"]["dtype"] = gr.Dropdown(choices=get_dtypes(), value="auto", label="Precision", info="Tensor type to load the model under.")
 						layout["settings"]["inputs"]["attentions"] = gr.Dropdown(choices=get_attentions(), value="auto", label="Attentions", info="Attention mechanism to utilize.")