""" Gradio UI Event Handlers Module Main entry point for setting up all event handlers """ import gradio as gr from typing import Optional # Import handler modules from . import generation_handlers as gen_h from . import results_handlers as res_h from . import training_handlers as train_h from acestep.gradio_ui.i18n import t def setup_event_handlers(demo, dit_handler, llm_handler, dataset_handler, dataset_section, generation_section, results_section): """Setup event handlers connecting UI components and business logic""" # ========== Dataset Handlers ========== dataset_section["import_dataset_btn"].click( fn=dataset_handler.import_dataset, inputs=[dataset_section["dataset_type"]], outputs=[dataset_section["data_status"]] ) # ========== Service Initialization ========== generation_section["refresh_btn"].click( fn=lambda: gen_h.refresh_checkpoints(dit_handler), outputs=[generation_section["checkpoint_dropdown"]] ) generation_section["config_path"].change( fn=gen_h.update_model_type_settings, inputs=[generation_section["config_path"]], outputs=[ generation_section["inference_steps"], generation_section["guidance_scale"], generation_section["use_adg"], generation_section["shift"], generation_section["cfg_interval_start"], generation_section["cfg_interval_end"], generation_section["task_type"], ] ) generation_section["init_btn"].click( fn=lambda *args: gen_h.init_service_wrapper(dit_handler, llm_handler, *args), inputs=[ generation_section["checkpoint_dropdown"], generation_section["config_path"], generation_section["device"], generation_section["init_llm_checkbox"], generation_section["lm_model_path"], generation_section["backend_dropdown"], generation_section["use_flash_attention_checkbox"], generation_section["offload_to_cpu_checkbox"], generation_section["offload_dit_to_cpu_checkbox"], ], outputs=[ generation_section["init_status"], generation_section["generate_btn"], generation_section["service_config_accordion"], # Model type settings (updated based on actual loaded model) generation_section["inference_steps"], generation_section["guidance_scale"], generation_section["use_adg"], generation_section["shift"], generation_section["cfg_interval_start"], generation_section["cfg_interval_end"], generation_section["task_type"], ] ) # ========== LoRA Handlers ========== generation_section["load_lora_btn"].click( fn=dit_handler.load_lora, inputs=[generation_section["lora_path"]], outputs=[generation_section["lora_status"]] ).then( # Update checkbox to enabled state after loading fn=lambda: gr.update(value=True), outputs=[generation_section["use_lora_checkbox"]] ) generation_section["unload_lora_btn"].click( fn=dit_handler.unload_lora, outputs=[generation_section["lora_status"]] ).then( # Update checkbox to disabled state after unloading fn=lambda: gr.update(value=False), outputs=[generation_section["use_lora_checkbox"]] ) generation_section["use_lora_checkbox"].change( fn=dit_handler.set_use_lora, inputs=[generation_section["use_lora_checkbox"]], outputs=[generation_section["lora_status"]] ) # ========== UI Visibility Updates ========== generation_section["init_llm_checkbox"].change( fn=gen_h.update_negative_prompt_visibility, inputs=[generation_section["init_llm_checkbox"]], outputs=[generation_section["lm_negative_prompt"]] ) generation_section["init_llm_checkbox"].change( fn=gen_h.update_audio_cover_strength_visibility, inputs=[generation_section["task_type"], generation_section["init_llm_checkbox"]], outputs=[generation_section["audio_cover_strength"]] ) generation_section["task_type"].change( fn=gen_h.update_audio_cover_strength_visibility, inputs=[generation_section["task_type"], generation_section["init_llm_checkbox"]], outputs=[generation_section["audio_cover_strength"]] ) generation_section["batch_size_input"].change( fn=gen_h.update_audio_components_visibility, inputs=[generation_section["batch_size_input"]], outputs=[ results_section["audio_col_1"], results_section["audio_col_2"], results_section["audio_col_3"], results_section["audio_col_4"], results_section["audio_row_5_8"], results_section["audio_col_5"], results_section["audio_col_6"], results_section["audio_col_7"], results_section["audio_col_8"], ] ) # ========== Audio Conversion ========== generation_section["convert_src_to_codes_btn"].click( fn=lambda src: gen_h.convert_src_audio_to_codes_wrapper(dit_handler, src), inputs=[generation_section["src_audio"]], outputs=[generation_section["text2music_audio_code_string"]] ) # ========== Instruction UI Updates ========== for trigger in [generation_section["task_type"], generation_section["track_name"], generation_section["complete_track_classes"]]: trigger.change( fn=lambda *args: gen_h.update_instruction_ui(dit_handler, *args), inputs=[ generation_section["task_type"], generation_section["track_name"], generation_section["complete_track_classes"], generation_section["text2music_audio_code_string"], generation_section["init_llm_checkbox"] ], outputs=[ generation_section["instruction_display_gen"], generation_section["track_name"], generation_section["complete_track_classes"], generation_section["audio_cover_strength"], generation_section["repainting_group"], generation_section["text2music_audio_codes_group"], ] ) # ========== Sample/Transcribe Handlers ========== # Load random example from ./examples/text2music directory generation_section["sample_btn"].click( fn=lambda task: gen_h.load_random_example(task) + (True,), inputs=[ generation_section["task_type"], ], outputs=[ generation_section["captions"], generation_section["lyrics"], generation_section["think_checkbox"], generation_section["bpm"], generation_section["audio_duration"], generation_section["key_scale"], generation_section["vocal_language"], generation_section["time_signature"], results_section["is_format_caption_state"] ] ) generation_section["text2music_audio_code_string"].change( fn=gen_h.update_transcribe_button_text, inputs=[generation_section["text2music_audio_code_string"]], outputs=[generation_section["transcribe_btn"]] ) generation_section["transcribe_btn"].click( fn=lambda codes, debug: gen_h.transcribe_audio_codes(llm_handler, codes, debug), inputs=[ generation_section["text2music_audio_code_string"], generation_section["constrained_decoding_debug"] ], outputs=[ results_section["status_output"], generation_section["captions"], generation_section["lyrics"], generation_section["bpm"], generation_section["audio_duration"], generation_section["key_scale"], generation_section["vocal_language"], generation_section["time_signature"], results_section["is_format_caption_state"] ] ) # ========== Reset Format Caption Flag ========== for trigger in [generation_section["captions"], generation_section["lyrics"], generation_section["bpm"], generation_section["key_scale"], generation_section["time_signature"], generation_section["vocal_language"], generation_section["audio_duration"]]: trigger.change( fn=gen_h.reset_format_caption_flag, inputs=[], outputs=[results_section["is_format_caption_state"]] ) # ========== Audio Uploads Accordion ========== for trigger in [generation_section["reference_audio"], generation_section["src_audio"]]: trigger.change( fn=gen_h.update_audio_uploads_accordion, inputs=[generation_section["reference_audio"], generation_section["src_audio"]], outputs=[generation_section["audio_uploads_accordion"]] ) # ========== Instrumental Checkbox ========== generation_section["instrumental_checkbox"].change( fn=gen_h.handle_instrumental_checkbox, inputs=[generation_section["instrumental_checkbox"], generation_section["lyrics"]], outputs=[generation_section["lyrics"]] ) # ========== Format Button ========== # Note: cfg_scale and negative_prompt are not supported in format mode generation_section["format_btn"].click( fn=lambda caption, lyrics, bpm, duration, key_scale, time_sig, temp, top_k, top_p, debug: gen_h.handle_format_sample( llm_handler, caption, lyrics, bpm, duration, key_scale, time_sig, temp, top_k, top_p, debug ), inputs=[ generation_section["captions"], generation_section["lyrics"], generation_section["bpm"], generation_section["audio_duration"], generation_section["key_scale"], generation_section["time_signature"], generation_section["lm_temperature"], generation_section["lm_top_k"], generation_section["lm_top_p"], generation_section["constrained_decoding_debug"], ], outputs=[ generation_section["captions"], generation_section["lyrics"], generation_section["bpm"], generation_section["audio_duration"], generation_section["key_scale"], generation_section["vocal_language"], generation_section["time_signature"], results_section["is_format_caption_state"], results_section["status_output"], ] ) # ========== Simple/Custom Mode Toggle ========== generation_section["generation_mode"].change( fn=gen_h.handle_generation_mode_change, inputs=[generation_section["generation_mode"]], outputs=[ generation_section["simple_mode_group"], generation_section["caption_accordion"], generation_section["lyrics_accordion"], generation_section["generate_btn"], generation_section["simple_sample_created"], generation_section["optional_params_accordion"], ] ) # ========== Simple Mode Instrumental Checkbox ========== # When instrumental is checked, disable vocal language and set to ["unknown"] generation_section["simple_instrumental_checkbox"].change( fn=gen_h.handle_simple_instrumental_change, inputs=[generation_section["simple_instrumental_checkbox"]], outputs=[generation_section["simple_vocal_language"]] ) # ========== Random Description Button ========== generation_section["random_desc_btn"].click( fn=gen_h.load_random_simple_description, inputs=[], outputs=[ generation_section["simple_query_input"], generation_section["simple_instrumental_checkbox"], generation_section["simple_vocal_language"], ] ) # ========== Create Sample Button (Simple Mode) ========== # Note: cfg_scale and negative_prompt are not supported in create_sample mode generation_section["create_sample_btn"].click( fn=lambda query, instrumental, vocal_lang, temp, top_k, top_p, debug: gen_h.handle_create_sample( llm_handler, query, instrumental, vocal_lang, temp, top_k, top_p, debug ), inputs=[ generation_section["simple_query_input"], generation_section["simple_instrumental_checkbox"], generation_section["simple_vocal_language"], generation_section["lm_temperature"], generation_section["lm_top_k"], generation_section["lm_top_p"], generation_section["constrained_decoding_debug"], ], outputs=[ generation_section["captions"], generation_section["lyrics"], generation_section["bpm"], generation_section["audio_duration"], generation_section["key_scale"], generation_section["vocal_language"], generation_section["simple_vocal_language"], generation_section["time_signature"], generation_section["instrumental_checkbox"], generation_section["caption_accordion"], generation_section["lyrics_accordion"], generation_section["generate_btn"], generation_section["simple_sample_created"], generation_section["think_checkbox"], results_section["is_format_caption_state"], results_section["status_output"], ] ) # ========== Load/Save Metadata ========== generation_section["load_file"].upload( fn=gen_h.load_metadata, inputs=[generation_section["load_file"]], outputs=[ generation_section["task_type"], generation_section["captions"], generation_section["lyrics"], generation_section["vocal_language"], generation_section["bpm"], generation_section["key_scale"], generation_section["time_signature"], generation_section["audio_duration"], generation_section["batch_size_input"], generation_section["inference_steps"], generation_section["guidance_scale"], generation_section["seed"], generation_section["random_seed_checkbox"], generation_section["use_adg"], generation_section["cfg_interval_start"], generation_section["cfg_interval_end"], generation_section["shift"], generation_section["infer_method"], generation_section["custom_timesteps"], generation_section["audio_format"], generation_section["lm_temperature"], generation_section["lm_cfg_scale"], generation_section["lm_top_k"], generation_section["lm_top_p"], generation_section["lm_negative_prompt"], generation_section["use_cot_metas"], # Added: use_cot_metas generation_section["use_cot_caption"], generation_section["use_cot_language"], generation_section["audio_cover_strength"], generation_section["think_checkbox"], generation_section["text2music_audio_code_string"], generation_section["repainting_start"], generation_section["repainting_end"], generation_section["track_name"], generation_section["complete_track_classes"], generation_section["instrumental_checkbox"], # Added: instrumental_checkbox results_section["is_format_caption_state"] ] ) # Save buttons for all 8 audio outputs download_existing_js = """(current_audio, batch_files) => { // Debug: print what the input actually is console.log("👉 [Debug] Current Audio Input:", current_audio); // 1. Safety check if (!current_audio) { console.warn("⚠️ No audio selected or audio is empty."); return; } if (!batch_files || !Array.isArray(batch_files)) { console.warn("⚠️ Batch file list is empty/not ready."); return; } // 2. Smartly extract path string let pathString = ""; if (typeof current_audio === "string") { // Case A: direct path string received pathString = current_audio; } else if (typeof current_audio === "object") { // Case B: an object is received, try common properties // Gradio file objects usually have path, url, or name pathString = current_audio.path || current_audio.name || current_audio.url || ""; } if (!pathString) { console.error("❌ Error: Could not extract a valid path string from input.", current_audio); return; } // 3. Extract Key (UUID) // Path could be /tmp/.../uuid.mp3 or url like /file=.../uuid.mp3 let filename = pathString.split(/[\\\\/]/).pop(); // get the filename let key = filename.split('.')[0]; // get UUID without extension console.log(`🔑 Key extracted: ${key}`); // 4. Find matching file(s) in the list let targets = batch_files.filter(f => { // Also extract names from batch_files objects // f usually contains name (backend path) and orig_name (download name) const fPath = f.name || f.path || ""; return fPath.includes(key); }); if (targets.length === 0) { console.warn("❌ No matching files found in batch list for key:", key); alert("Batch list does not contain this file yet. Please wait for generation to finish."); return; } // 5. Trigger download(s) console.log(`🎯 Found ${targets.length} files to download.`); targets.forEach((f, index) => { setTimeout(() => { const a = document.createElement('a'); // Prefer url (frontend-accessible link), otherwise try data a.href = f.url || f.data; a.download = f.orig_name || "download"; a.style.display = 'none'; document.body.appendChild(a); a.click(); document.body.removeChild(a); }, index * 1000); // 300ms interval to avoid browser blocking }); } """ for btn_idx in range(1, 9): results_section[f"save_btn_{btn_idx}"].click( fn=None, inputs=[ results_section[f"generated_audio_{btn_idx}"], results_section["generated_audio_batch"], ], js=download_existing_js # Run the above JS ) # ========== Send to SRC Handlers ========== for btn_idx in range(1, 9): results_section[f"send_to_src_btn_{btn_idx}"].click( fn=res_h.send_audio_to_src_with_metadata, inputs=[ results_section[f"generated_audio_{btn_idx}"], results_section["lm_metadata_state"] ], outputs=[ generation_section["src_audio"], generation_section["bpm"], generation_section["captions"], generation_section["lyrics"], generation_section["audio_duration"], generation_section["key_scale"], generation_section["vocal_language"], generation_section["time_signature"], results_section["is_format_caption_state"] ] ) # ========== Score Calculation Handlers ========== # Use default argument to capture btn_idx value at definition time (Python closure fix) def make_score_handler(idx): return lambda scale, batch_idx, queue: res_h.calculate_score_handler_with_selection( dit_handler, llm_handler, idx, scale, batch_idx, queue ) for btn_idx in range(1, 9): results_section[f"score_btn_{btn_idx}"].click( fn=make_score_handler(btn_idx), inputs=[ generation_section["score_scale"], results_section["current_batch_index"], results_section["batch_queue"], ], outputs=[ results_section[f"score_display_{btn_idx}"], results_section[f"details_accordion_{btn_idx}"], results_section["batch_queue"] ] ) # ========== LRC Timestamp Handlers ========== # Use default argument to capture btn_idx value at definition time (Python closure fix) def make_lrc_handler(idx): return lambda batch_idx, queue, vocal_lang, infer_steps: res_h.generate_lrc_handler( dit_handler, idx, batch_idx, queue, vocal_lang, infer_steps ) for btn_idx in range(1, 9): results_section[f"lrc_btn_{btn_idx}"].click( fn=make_lrc_handler(btn_idx), inputs=[ results_section["current_batch_index"], results_section["batch_queue"], generation_section["vocal_language"], generation_section["inference_steps"], ], outputs=[ results_section[f"lrc_display_{btn_idx}"], results_section[f"details_accordion_{btn_idx}"], # NOTE: Removed generated_audio output! # Audio subtitles are now updated via lrc_display.change() event. results_section["batch_queue"] ] ) def generation_wrapper(*args): yield from res_h.generate_with_batch_management(dit_handler, llm_handler, *args) # ========== Generation Handler ========== generation_section["generate_btn"].click( fn=generation_wrapper, inputs=[ generation_section["captions"], generation_section["lyrics"], generation_section["bpm"], generation_section["key_scale"], generation_section["time_signature"], generation_section["vocal_language"], generation_section["inference_steps"], generation_section["guidance_scale"], generation_section["random_seed_checkbox"], generation_section["seed"], generation_section["reference_audio"], generation_section["audio_duration"], generation_section["batch_size_input"], generation_section["src_audio"], generation_section["text2music_audio_code_string"], generation_section["repainting_start"], generation_section["repainting_end"], generation_section["instruction_display_gen"], generation_section["audio_cover_strength"], generation_section["task_type"], generation_section["use_adg"], generation_section["cfg_interval_start"], generation_section["cfg_interval_end"], generation_section["shift"], generation_section["infer_method"], generation_section["custom_timesteps"], generation_section["audio_format"], generation_section["lm_temperature"], generation_section["think_checkbox"], generation_section["lm_cfg_scale"], generation_section["lm_top_k"], generation_section["lm_top_p"], generation_section["lm_negative_prompt"], generation_section["use_cot_metas"], generation_section["use_cot_caption"], generation_section["use_cot_language"], results_section["is_format_caption_state"], generation_section["constrained_decoding_debug"], generation_section["allow_lm_batch"], generation_section["auto_score"], generation_section["auto_lrc"], generation_section["score_scale"], generation_section["lm_batch_chunk_size"], generation_section["track_name"], generation_section["complete_track_classes"], generation_section["autogen_checkbox"], results_section["current_batch_index"], results_section["total_batches"], results_section["batch_queue"], results_section["generation_params_state"], ], outputs=[ results_section["generated_audio_1"], results_section["generated_audio_2"], results_section["generated_audio_3"], results_section["generated_audio_4"], results_section["generated_audio_5"], results_section["generated_audio_6"], results_section["generated_audio_7"], results_section["generated_audio_8"], results_section["generated_audio_batch"], results_section["generation_info"], results_section["status_output"], generation_section["seed"], results_section["score_display_1"], results_section["score_display_2"], results_section["score_display_3"], results_section["score_display_4"], results_section["score_display_5"], results_section["score_display_6"], results_section["score_display_7"], results_section["score_display_8"], results_section["codes_display_1"], results_section["codes_display_2"], results_section["codes_display_3"], results_section["codes_display_4"], results_section["codes_display_5"], results_section["codes_display_6"], results_section["codes_display_7"], results_section["codes_display_8"], results_section["details_accordion_1"], results_section["details_accordion_2"], results_section["details_accordion_3"], results_section["details_accordion_4"], results_section["details_accordion_5"], results_section["details_accordion_6"], results_section["details_accordion_7"], results_section["details_accordion_8"], results_section["lrc_display_1"], results_section["lrc_display_2"], results_section["lrc_display_3"], results_section["lrc_display_4"], results_section["lrc_display_5"], results_section["lrc_display_6"], results_section["lrc_display_7"], results_section["lrc_display_8"], results_section["lm_metadata_state"], results_section["is_format_caption_state"], results_section["current_batch_index"], results_section["total_batches"], results_section["batch_queue"], results_section["generation_params_state"], results_section["batch_indicator"], results_section["prev_batch_btn"], results_section["next_batch_btn"], results_section["next_batch_status"], results_section["restore_params_btn"], ] ).then( fn=lambda *args: res_h.generate_next_batch_background(dit_handler, llm_handler, *args), inputs=[ generation_section["autogen_checkbox"], results_section["generation_params_state"], results_section["current_batch_index"], results_section["total_batches"], results_section["batch_queue"], results_section["is_format_caption_state"], ], outputs=[ results_section["batch_queue"], results_section["total_batches"], results_section["next_batch_status"], results_section["next_batch_btn"], ] ) # ========== Batch Navigation Handlers ========== results_section["prev_batch_btn"].click( fn=res_h.navigate_to_previous_batch, inputs=[ results_section["current_batch_index"], results_section["batch_queue"], ], outputs=[ results_section["generated_audio_1"], results_section["generated_audio_2"], results_section["generated_audio_3"], results_section["generated_audio_4"], results_section["generated_audio_5"], results_section["generated_audio_6"], results_section["generated_audio_7"], results_section["generated_audio_8"], results_section["generated_audio_batch"], results_section["generation_info"], results_section["current_batch_index"], results_section["batch_indicator"], results_section["prev_batch_btn"], results_section["next_batch_btn"], results_section["status_output"], results_section["score_display_1"], results_section["score_display_2"], results_section["score_display_3"], results_section["score_display_4"], results_section["score_display_5"], results_section["score_display_6"], results_section["score_display_7"], results_section["score_display_8"], results_section["codes_display_1"], results_section["codes_display_2"], results_section["codes_display_3"], results_section["codes_display_4"], results_section["codes_display_5"], results_section["codes_display_6"], results_section["codes_display_7"], results_section["codes_display_8"], results_section["lrc_display_1"], results_section["lrc_display_2"], results_section["lrc_display_3"], results_section["lrc_display_4"], results_section["lrc_display_5"], results_section["lrc_display_6"], results_section["lrc_display_7"], results_section["lrc_display_8"], results_section["details_accordion_1"], results_section["details_accordion_2"], results_section["details_accordion_3"], results_section["details_accordion_4"], results_section["details_accordion_5"], results_section["details_accordion_6"], results_section["details_accordion_7"], results_section["details_accordion_8"], results_section["restore_params_btn"], ] ) results_section["next_batch_btn"].click( fn=res_h.capture_current_params, inputs=[ generation_section["captions"], generation_section["lyrics"], generation_section["bpm"], generation_section["key_scale"], generation_section["time_signature"], generation_section["vocal_language"], generation_section["inference_steps"], generation_section["guidance_scale"], generation_section["random_seed_checkbox"], generation_section["seed"], generation_section["reference_audio"], generation_section["audio_duration"], generation_section["batch_size_input"], generation_section["src_audio"], generation_section["text2music_audio_code_string"], generation_section["repainting_start"], generation_section["repainting_end"], generation_section["instruction_display_gen"], generation_section["audio_cover_strength"], generation_section["task_type"], generation_section["use_adg"], generation_section["cfg_interval_start"], generation_section["cfg_interval_end"], generation_section["shift"], generation_section["infer_method"], generation_section["custom_timesteps"], generation_section["audio_format"], generation_section["lm_temperature"], generation_section["think_checkbox"], generation_section["lm_cfg_scale"], generation_section["lm_top_k"], generation_section["lm_top_p"], generation_section["lm_negative_prompt"], generation_section["use_cot_metas"], generation_section["use_cot_caption"], generation_section["use_cot_language"], generation_section["constrained_decoding_debug"], generation_section["allow_lm_batch"], generation_section["auto_score"], generation_section["auto_lrc"], generation_section["score_scale"], generation_section["lm_batch_chunk_size"], generation_section["track_name"], generation_section["complete_track_classes"], ], outputs=[results_section["generation_params_state"]] ).then( fn=res_h.navigate_to_next_batch, inputs=[ generation_section["autogen_checkbox"], results_section["current_batch_index"], results_section["total_batches"], results_section["batch_queue"], ], outputs=[ results_section["generated_audio_1"], results_section["generated_audio_2"], results_section["generated_audio_3"], results_section["generated_audio_4"], results_section["generated_audio_5"], results_section["generated_audio_6"], results_section["generated_audio_7"], results_section["generated_audio_8"], results_section["generated_audio_batch"], results_section["generation_info"], results_section["current_batch_index"], results_section["batch_indicator"], results_section["prev_batch_btn"], results_section["next_batch_btn"], results_section["status_output"], results_section["next_batch_status"], results_section["score_display_1"], results_section["score_display_2"], results_section["score_display_3"], results_section["score_display_4"], results_section["score_display_5"], results_section["score_display_6"], results_section["score_display_7"], results_section["score_display_8"], results_section["codes_display_1"], results_section["codes_display_2"], results_section["codes_display_3"], results_section["codes_display_4"], results_section["codes_display_5"], results_section["codes_display_6"], results_section["codes_display_7"], results_section["codes_display_8"], results_section["lrc_display_1"], results_section["lrc_display_2"], results_section["lrc_display_3"], results_section["lrc_display_4"], results_section["lrc_display_5"], results_section["lrc_display_6"], results_section["lrc_display_7"], results_section["lrc_display_8"], results_section["details_accordion_1"], results_section["details_accordion_2"], results_section["details_accordion_3"], results_section["details_accordion_4"], results_section["details_accordion_5"], results_section["details_accordion_6"], results_section["details_accordion_7"], results_section["details_accordion_8"], results_section["restore_params_btn"], ] ).then( fn=lambda *args: res_h.generate_next_batch_background(dit_handler, llm_handler, *args), inputs=[ generation_section["autogen_checkbox"], results_section["generation_params_state"], results_section["current_batch_index"], results_section["total_batches"], results_section["batch_queue"], results_section["is_format_caption_state"], ], outputs=[ results_section["batch_queue"], results_section["total_batches"], results_section["next_batch_status"], results_section["next_batch_btn"], ] ) # ========== Restore Parameters Handler ========== results_section["restore_params_btn"].click( fn=res_h.restore_batch_parameters, inputs=[ results_section["current_batch_index"], results_section["batch_queue"] ], outputs=[ generation_section["text2music_audio_code_string"], generation_section["captions"], generation_section["lyrics"], generation_section["bpm"], generation_section["key_scale"], generation_section["time_signature"], generation_section["vocal_language"], generation_section["audio_duration"], generation_section["batch_size_input"], generation_section["inference_steps"], generation_section["lm_temperature"], generation_section["lm_cfg_scale"], generation_section["lm_top_k"], generation_section["lm_top_p"], generation_section["think_checkbox"], generation_section["use_cot_caption"], generation_section["use_cot_language"], generation_section["allow_lm_batch"], generation_section["track_name"], generation_section["complete_track_classes"], ] ) # ========== LRC Display Change Handlers ========== # NEW APPROACH: Use lrc_display.change() to update audio subtitles # This decouples audio value updates from subtitle updates, avoiding flickering. # # When lrc_display text changes (from generate, LRC button, or manual edit): # 1. lrc_display.change() is triggered # 2. update_audio_subtitles_from_lrc() parses LRC and updates audio subtitles # 3. Audio value is NEVER updated here - only subtitles for lrc_idx in range(1, 9): results_section[f"lrc_display_{lrc_idx}"].change( fn=res_h.update_audio_subtitles_from_lrc, inputs=[ results_section[f"lrc_display_{lrc_idx}"], # audio_duration not needed - parse_lrc_to_subtitles calculates end time from timestamps ], outputs=[ results_section[f"generated_audio_{lrc_idx}"], # Only updates subtitles, not value ] ) def setup_training_event_handlers(demo, dit_handler, llm_handler, training_section): """Setup event handlers for the training tab (dataset builder and LoRA training)""" # ========== Load Existing Dataset (Top Section) ========== # Load existing dataset JSON at the top of Dataset Builder training_section["load_json_btn"].click( fn=train_h.load_existing_dataset_for_preprocess, inputs=[ training_section["load_json_path"], training_section["dataset_builder_state"], ], outputs=[ training_section["load_json_status"], training_section["audio_files_table"], training_section["sample_selector"], training_section["dataset_builder_state"], # Also update preview fields with first sample training_section["preview_audio"], training_section["preview_filename"], training_section["edit_caption"], training_section["edit_lyrics"], training_section["edit_bpm"], training_section["edit_keyscale"], training_section["edit_timesig"], training_section["edit_duration"], training_section["edit_language"], training_section["edit_instrumental"], ] ) # ========== Dataset Builder Handlers ========== # Scan directory for audio files training_section["scan_btn"].click( fn=lambda dir, name, tag, pos, instr, state: train_h.scan_directory( dir, name, tag, pos, instr, state ), inputs=[ training_section["audio_directory"], training_section["dataset_name"], training_section["custom_tag"], training_section["tag_position"], training_section["all_instrumental"], training_section["dataset_builder_state"], ], outputs=[ training_section["audio_files_table"], training_section["scan_status"], training_section["sample_selector"], training_section["dataset_builder_state"], ] ) # Auto-label all samples training_section["auto_label_btn"].click( fn=lambda state, skip: train_h.auto_label_all(dit_handler, llm_handler, state, skip), inputs=[ training_section["dataset_builder_state"], training_section["skip_metas"], ], outputs=[ training_section["audio_files_table"], training_section["label_progress"], training_section["dataset_builder_state"], ] ) # Sample selector change - update preview training_section["sample_selector"].change( fn=train_h.get_sample_preview, inputs=[ training_section["sample_selector"], training_section["dataset_builder_state"], ], outputs=[ training_section["preview_audio"], training_section["preview_filename"], training_section["edit_caption"], training_section["edit_lyrics"], training_section["edit_bpm"], training_section["edit_keyscale"], training_section["edit_timesig"], training_section["edit_duration"], training_section["edit_language"], training_section["edit_instrumental"], ] ) # Save sample edit training_section["save_edit_btn"].click( fn=train_h.save_sample_edit, inputs=[ training_section["sample_selector"], training_section["edit_caption"], training_section["edit_lyrics"], training_section["edit_bpm"], training_section["edit_keyscale"], training_section["edit_timesig"], training_section["edit_language"], training_section["edit_instrumental"], training_section["dataset_builder_state"], ], outputs=[ training_section["audio_files_table"], training_section["edit_status"], training_section["dataset_builder_state"], ] ) # Update settings when changed for trigger in [training_section["custom_tag"], training_section["tag_position"], training_section["all_instrumental"]]: trigger.change( fn=train_h.update_settings, inputs=[ training_section["custom_tag"], training_section["tag_position"], training_section["all_instrumental"], training_section["dataset_builder_state"], ], outputs=[training_section["dataset_builder_state"]] ) # Save dataset training_section["save_dataset_btn"].click( fn=train_h.save_dataset, inputs=[ training_section["save_path"], training_section["dataset_name"], training_section["dataset_builder_state"], ], outputs=[training_section["save_status"]] ) # ========== Preprocess Handlers ========== # Load existing dataset JSON for preprocessing # This also updates the preview section so users can view/edit samples training_section["load_existing_dataset_btn"].click( fn=train_h.load_existing_dataset_for_preprocess, inputs=[ training_section["load_existing_dataset_path"], training_section["dataset_builder_state"], ], outputs=[ training_section["load_existing_status"], training_section["audio_files_table"], training_section["sample_selector"], training_section["dataset_builder_state"], # Also update preview fields with first sample training_section["preview_audio"], training_section["preview_filename"], training_section["edit_caption"], training_section["edit_lyrics"], training_section["edit_bpm"], training_section["edit_keyscale"], training_section["edit_timesig"], training_section["edit_duration"], training_section["edit_language"], training_section["edit_instrumental"], ] ) # Preprocess dataset to tensor files training_section["preprocess_btn"].click( fn=lambda output_dir, state: train_h.preprocess_dataset( output_dir, dit_handler, state ), inputs=[ training_section["preprocess_output_dir"], training_section["dataset_builder_state"], ], outputs=[training_section["preprocess_progress"]] ) # ========== Training Tab Handlers ========== # Load preprocessed tensor dataset training_section["load_dataset_btn"].click( fn=train_h.load_training_dataset, inputs=[training_section["training_tensor_dir"]], outputs=[training_section["training_dataset_info"]] ) # Start training from preprocessed tensors def training_wrapper(tensor_dir, r, a, d, lr, ep, bs, ga, se, sh, sd, od, ts): try: for progress, log, plot, state in train_h.start_training( tensor_dir, dit_handler, r, a, d, lr, ep, bs, ga, se, sh, sd, od, ts ): yield progress, log, plot, state except Exception as e: logger.exception("Training wrapper error") yield f"❌ Error: {str(e)}", str(e), None, ts training_section["start_training_btn"].click( fn=training_wrapper, inputs=[ training_section["training_tensor_dir"], training_section["lora_rank"], training_section["lora_alpha"], training_section["lora_dropout"], training_section["learning_rate"], training_section["train_epochs"], training_section["train_batch_size"], training_section["gradient_accumulation"], training_section["save_every_n_epochs"], training_section["training_shift"], training_section["training_seed"], training_section["lora_output_dir"], training_section["training_state"], ], outputs=[ training_section["training_progress"], training_section["training_log"], training_section["training_loss_plot"], training_section["training_state"], ] ) # Stop training training_section["stop_training_btn"].click( fn=train_h.stop_training, inputs=[training_section["training_state"]], outputs=[ training_section["training_progress"], training_section["training_state"], ] ) # Export LoRA training_section["export_lora_btn"].click( fn=train_h.export_lora, inputs=[ training_section["export_path"], training_section["lora_output_dir"], ], outputs=[training_section["export_status"]] )