Update app.py

app.py

@@ -13,12 +13,31 @@ import spaces
 os.environ["HF_HUB_ENABLE_HF_TRANSFER"] = "1"
 
 checkpoints_path = snapshot_download("mrfakename/openf5-v2", allow_patterns=["model_*.pt", "vocab.txt"], token=os.getenv("HF_TOKEN"))
-
 models = {}
-for checkpoint_path in tqdm(os.listdir(checkpoints_path), desc="Loading models"):
-    if checkpoint_path.endswith(".pt"):
-        model_name = checkpoint_path.split("/")[-1].replace(".pt", "")
-        models[model_name] = F5TTS(ckpt_file=os.path.join(checkpoints_path, checkpoint_path), vocab_file=os.path.join(checkpoints_path, "vocab.txt"))
+checkpoint_files = [f for f in os.listdir(checkpoints_path) if f.endswith(".pt")]
+
+# Sort checkpoint files by step number
+def get_step_number(filename):
+    name = filename.replace(".pt", "")
+    if name == "model_last":
+        return float('-inf')  # Ensure model_last comes first
+    try:
+        return int(name.split("_")[1])
+    except (IndexError, ValueError):
+        return float('inf')  # Put non-standard names at the end
+
+sorted_checkpoints = sorted(checkpoint_files, key=get_step_number)
+
+# Load models in the sorted order
+for checkpoint_path in tqdm(sorted_checkpoints, desc="Loading models"):
+    model_name = checkpoint_path.replace(".pt", "")
+    # Load one model at a time to be memory efficient
+    models[model_name] = F5TTS(ckpt_file=os.path.join(checkpoints_path, checkpoint_path), 
+                              vocab_file=os.path.join(checkpoints_path, "vocab.txt"))
+
+
+
+
 
 @spaces.GPU
 def generate_audio(model_name, ref_file, ref_text, gen_text, progress=gr.Progress()):
@@ -36,7 +55,7 @@ def generate_audio(model_name, ref_file, ref_text, gen_text, progress=gr.Progres
 
 with gr.Blocks() as demo:
     gr.Markdown(ABOUT)
-    model_name = gr.Dropdown(label="Model", choices=list(models.keys()))
+    model_name = gr.Radio(label="Model", choices=list(models.keys()))
     ref_file = gr.Audio(label="Reference Audio", type="filepath")
     gen_text = gr.Textbox(label="Text")
     btn_generate = gr.Button("Generate Audio", variant="primary")