Update demo.py

demo.py

@@ -16,12 +16,12 @@ with gr.Blocks(title="🔊", theme=gr.themes.Base(primary_hue="blue", neutral_hu
     with gr.Tabs():
         with gr.Tab("Inference"):
             with gr.Row():
-                # Initialize voice_model with proper choices
-                initial_names = sorted(names) if names else []
+                # Get initial model choices from original.py
+                initial_model_choices = sorted(names) if names else []
                 voice_model = gr.Dropdown(
                     label="Model Voice", 
-                    choices=initial_names, 
-                    value=initial_names[0] if initial_names else None, 
+                    choices=initial_model_choices, 
+                    value=initial_model_choices[0] if initial_model_choices else None, 
                     interactive=True
                 )
                 refresh_button = gr.Button("Refresh", variant="primary")
@@ -48,7 +48,7 @@ with gr.Blocks(title="🔊", theme=gr.themes.Base(primary_hue="blue", neutral_hu
                     with gr.Row():
                         input_audio0 = gr.Dropdown(
                             label="Input Path",
-                            value="",
+                            value=None,
                             choices=[],
                             allow_custom_value=True
                         )
@@ -69,7 +69,7 @@ with gr.Blocks(title="🔊", theme=gr.themes.Base(primary_hue="blue", neutral_hu
                         def handle_record(audio):
                             if audio:
                                 return audio
-                            return ""
+                            return None
                         
                         record_button.change(
                             fn=handle_record,
@@ -80,7 +80,7 @@ with gr.Blocks(title="🔊", theme=gr.themes.Base(primary_hue="blue", neutral_hu
                         def handle_upload(audio):
                             if audio:
                                 return audio
-                            return ""
+                            return None
                         
                         dropbox.change(
                             fn=handle_upload,
@@ -94,7 +94,7 @@ with gr.Blocks(title="🔊", theme=gr.themes.Base(primary_hue="blue", neutral_hu
                             label="Change Index",
                             choices=[],
                             interactive=True,
-                            value=""
+                            value=None
                         )
                         index_rate1 = gr.Slider(
                             minimum=0,
@@ -154,31 +154,30 @@ with gr.Blocks(title="🔊", theme=gr.themes.Base(primary_hue="blue", neutral_hu
 
                     # Consolidated refresh logic
                     def refresh_ui():
-                        # Get updated lists
+                        # Get updated lists from change_choices which returns dictionaries
                         try:
-                            model_choices, index_choices = change_choices()
+                            model_result, index_result = change_choices()
+                            model_choices = model_result["choices"]
+                            index_choices = index_result["choices"]
                         except Exception as e:
                             print(f"Error in change_choices: {e}")
                             model_choices = []
                             index_choices = []
                         
-                        # Ensure model_choices and index_choices are lists
-                        if not isinstance(model_choices, list):
-                            model_choices = []
-                        if not isinstance(index_choices, list):
-                            index_choices = []
-                        
                         audio_paths = get_audio_paths('audios')
                         
-                        # Get current values
-                        current_model = voice_model.value if hasattr(voice_model, 'value') else None
-                        current_index = file_index2.value if hasattr(file_index2, 'value') else None
-                        current_audio = input_audio0.value if hasattr(input_audio0, 'value') else None
+                        # Get current values to preserve selection when possible
+                        current_model = voice_model.value
+                        current_index = file_index2.value
+                        current_audio = input_audio0.value
                         
-                        # Set defaults
-                        default_model = current_model if current_model in model_choices else (model_choices[0] if model_choices else None)
-                        default_index = current_index if current_index in index_choices else (index_choices[0] if index_choices else None)
-                        default_audio = current_audio if current_audio in audio_paths else (audio_paths[0] if audio_paths else None)
+                        # Set defaults with fallback logic
+                        default_model = (current_model if current_model in model_choices 
+                                        else (model_choices[0] if model_choices else None))
+                        default_index = (current_index if current_index in index_choices 
+                                        else (index_choices[0] if index_choices else None))
+                        default_audio = (current_audio if current_audio in audio_paths 
+                                        else (audio_paths[0] if audio_paths else None))
 
                         return (
                             gr.update(choices=model_choices, value=default_model),  # voice_model
@@ -424,13 +423,20 @@ with gr.Blocks(title="🔊", theme=gr.themes.Base(primary_hue="blue", neutral_hu
                             interactive=True,
                         )
                         with gr.Accordion(label="Change pretrains", open=False):
-                            def get_pretrained_choices(sr, letter):
-                                pretrained_dir = 'assets/pretrained_v2'
-                                if not os.path.exists(pretrained_dir):
-                                    return []
-                                return [os.path.abspath(os.path.join(pretrained_dir, file)) 
-                                       for file in os.listdir(pretrained_dir) 
-                                       if file.endswith('.pth') and sr in file and letter in file]
+                            def get_pretrained_choices(sr, if_f0, version):
+                                # Use the original functions from original.py
+                                if version == "v1":
+                                    path_str = ""
+                                else:
+                                    path_str = "_v2"
+                                
+                                if if_f0:
+                                    f0_str = "f0"
+                                else:
+                                    f0_str = ""
+                                
+                                pretrained_G, pretrained_D = get_pretrained_models(path_str, f0_str, sr)
+                                return [pretrained_G] if pretrained_G else [], [pretrained_D] if pretrained_D else []
 
                             pretrained_G14 = gr.Dropdown(
                                 label="pretrained G",
@@ -447,17 +453,18 @@ with gr.Blocks(title="🔊", theme=gr.themes.Base(primary_hue="blue", neutral_hu
                                 interactive=True
                             )
 
-                            def update_pretrained_dropdowns(sr, f0, ver):
+                            def update_pretrained_dropdowns(sr, if_f0, ver):
                                 sr_str = sr if isinstance(sr, str) else str(sr)
-                                g_choices = get_pretrained_choices(sr_str, 'G')
-                                d_choices = get_pretrained_choices(sr_str, 'D')
+                                g_choices, d_choices = get_pretrained_choices(sr_str, if_f0, ver)
                                 return (
                                     gr.update(choices=g_choices, value=g_choices[0] if g_choices else ""),
                                     gr.update(choices=d_choices, value=d_choices[0] if d_choices else "")
                                 )
                             
+                            # Bind update function to changes
                             sr2.change(fn=update_pretrained_dropdowns, inputs=[sr2, if_f0_3, version19], outputs=[pretrained_G14, pretrained_D15])
                             version19.change(fn=update_pretrained_dropdowns, inputs=[sr2, if_f0_3, version19], outputs=[pretrained_G14, pretrained_D15])
+                            if_f0_3.change(fn=update_pretrained_dropdowns, inputs=[sr2, if_f0_3, version19], outputs=[pretrained_G14, pretrained_D15])
 
                     with gr.Row():
                         download_model = gr.Button('5.Download Model')
@@ -468,12 +475,12 @@ with gr.Blocks(title="🔊", theme=gr.themes.Base(primary_hue="blue", neutral_hu
                             if not name or name.strip() == "":
                                 return [], "Please enter a model name"
                             
-                            model_path = f'assets/weights/{name}'
-                            index_pattern = f'logs/{name.split(".")[0]}/added_*.index'
+                            model_path = f'logs/{name}'
+                            index_pattern = f'logs/{name}/added_*.index'
                             
                             files = []
                             if os.path.exists(model_path):
-                                files.extend([os.path.join(model_path, f) for f in os.listdir(model_path)])
+                                files.extend([os.path.join(model_path, f) for f in os.listdir(model_path) if f.endswith('.pth')])
                             files.extend(glob.glob(index_pattern))
                             
                             return files, f"Found {len(files)} files"
@@ -485,9 +492,9 @@ with gr.Blocks(title="🔊", theme=gr.themes.Base(primary_hue="blue", neutral_hu
                         )
                     
                     if_f0_3.change(
-                        change_f0,
-                        [if_f0_3, sr2, version19],
-                        [f0method8, pretrained_G14, pretrained_D15],
+                        fn=change_f0,
+                        inputs=[if_f0_3, sr2, version19],
+                        outputs=[f0method8, pretrained_G14, pretrained_D15],
                     )
                     
                     but5 = gr.Button("1 Click Training", variant="primary", visible=False)
@@ -540,8 +547,23 @@ with gr.Blocks(title="🔊", theme=gr.themes.Base(primary_hue="blue", neutral_hu
                     )
 
     # Populate UI on load
+    def on_load():
+        # Initial refresh
+        model_result, index_result = change_choices()
+        audio_paths = get_audio_paths('audios')
+        
+        default_model = model_result["choices"][0] if model_result["choices"] else None
+        default_index = index_result["choices"][0] if index_result["choices"] else None
+        default_audio = audio_paths[0] if audio_paths else None
+        
+        return (
+            gr.update(choices=model_result["choices"], value=default_model),  # voice_model
+            gr.update(choices=index_result["choices"], value=default_index),  # file_index2
+            gr.update(choices=audio_paths, value=default_audio)              # input_audio0
+        )
+    
     app.load(
-        fn=refresh_ui,
+        fn=on_load,
         inputs=[],
         outputs=[voice_model, file_index2, input_audio0]
     )