import os import gradio as gr import json import logging import shutil import tempfile from pathlib import Path import numpy as np from utils import ( rename_files_remove_spaces, load_audio_files, get_stems, generate_section_variants, export_section_variants, edm_arrangement_tab, ) from dotenv import load_dotenv load_dotenv() logger = logging.getLogger(__name__) logger.setLevel(logging.INFO) # Global variable to store the temporary directory TEMP_DIR = "test" # Global variable to store the selected variants SELECTED_VARIANTS = {} # Global variable to store all variants for each section ALL_VARIANTS = {} # Global variable to store the uploaded stems UPLOADED_STEMS = {} def process_uploaded_files(files, progress=gr.Progress()): """Process uploaded files and return basic info""" global TEMP_DIR, UPLOADED_STEMS try: if not files: return "Error: No files uploaded", [] progress(0, desc="Starting process...") # Create a temporary directory for processing TEMP_DIR = tempfile.mkdtemp() try: # Copy uploaded files to temp directory progress(0.2, desc="Copying uploaded files...") for file in files: if file.name.lower().endswith(".wav"): shutil.copy2(file.name, TEMP_DIR) # First rename all files to remove spaces progress(0.5, desc="Renaming files...") rename_files_remove_spaces(TEMP_DIR) # Load audio files progress(0.8, desc="Loading audio files...") UPLOADED_STEMS = load_audio_files(TEMP_DIR) if not UPLOADED_STEMS: return "Error: No stems loaded", [] # Get stem names stem_names = get_stems(TEMP_DIR) progress(1.0, desc="Complete!") return f"Successfully loaded {len(stem_names)} stems", stem_names except Exception as e: if os.path.exists(TEMP_DIR): shutil.rmtree(TEMP_DIR) TEMP_DIR = None raise e except Exception as e: return f"Error occurred: {str(e)}", [] def generate_section_variants_handler( section_type, bpm_value, bars_value, p_value, progress=gr.Progress() ): """Handler function for generating section variants""" global TEMP_DIR, ALL_VARIANTS, UPLOADED_STEMS if not TEMP_DIR or not os.path.exists(TEMP_DIR): return ( "Error: No stems loaded. Please upload stems first.", None, None, None, None, ) try: progress(0.1, desc=f"Getting arrangements for {section_type}...") # Generate variants variants = generate_section_variants( TEMP_DIR, UPLOADED_STEMS, section_type, bpm=int(bpm_value), bars=int(bars_value), p=float(p_value), progress=progress ) logger.info(f"S1: VARIANTS of section {section_type} : {variants}") progress(0.4, desc="Variants generated") # Store variants for later use ALL_VARIANTS[section_type] = variants # Export variants to audio files variant_output_dir = os.path.join(TEMP_DIR, section_type + "_variants") audio_paths = export_section_variants( variants, variant_output_dir, section_type ) progress(0.6, desc="Exporting audio files...") logger.info(f"S2: AUDIO_PATHS of section {section_type} : {audio_paths}") # Create audio elements for each variant variant1_audio = audio_paths.get("variant1") variant2_audio = audio_paths.get("variant2") variant3_audio = audio_paths.get("variant3") variant4_audio = audio_paths.get("variant4") # Descriptions descriptions = {key: data["description"] for key, data in variants.items()} descriptions_json = json.dumps(descriptions, indent=2) logger.info(f"S3: DESCRIPTIONS of section {section_type} : {descriptions_json}") progress(1.0, desc="Complete!") return ( f"Generated {len(variants)} variants for {section_type}", variant1_audio, variant2_audio, variant3_audio, variant4_audio, descriptions_json, ) except Exception as e: return f"Error generating variants: {str(e)}", None, None, None, None, None def select_variant(section_type, variant_num, append): """Select a variant for a specific section, with an option to append. Retry up to 5 times on error.""" global ALL_VARIANTS, SELECTED_VARIANTS max_retries = 5 for attempt in range(1, max_retries + 1): try: if section_type not in ALL_VARIANTS: msg = f"No variants generated for {section_type} yet" if attempt == max_retries: return msg continue variant_key = f"variant{variant_num}" if variant_key not in ALL_VARIANTS[section_type]: msg = f"Variant {variant_num} not found for {section_type}" if attempt == max_retries: return msg continue # If appending, add to the list of selected variants if append: # For now, just overwrite with the latest selection (no append logic) SELECTED_VARIANTS[section_type] = ALL_VARIANTS[section_type][variant_key]["config"] else: # Otherwise, replace the selection SELECTED_VARIANTS[section_type] = ALL_VARIANTS[section_type][variant_key]["config"] return f"Selected variant {variant_num} for {section_type} (Append: {append})" except Exception as e: msg = f"Error selecting variant (attempt {attempt}): {str(e)}" if attempt == max_retries: return msg continue def generate_full_track( crossfade_ms, output_track_name, include_intro, include_buildup, include_breakdown, include_drop, include_outro, progress=gr.Progress(), ): """Generate the full track from selected variants""" global TEMP_DIR, SELECTED_VARIANTS, UPLOADED_STEMS if not TEMP_DIR or not os.path.exists(TEMP_DIR): return "Error: No stems loaded", None, None try: progress(0.1, desc="Preparing to generate full track...") # Check which sections to include based on user selections and available variants sections_to_include = {} logger.info(f"SELECTED_VARIANTS: {SELECTED_VARIANTS}") if include_intro and "intro" in SELECTED_VARIANTS: sections_to_include["intro"] = SELECTED_VARIANTS["intro"] if include_buildup and "buildup" in SELECTED_VARIANTS: sections_to_include["buildup"] = SELECTED_VARIANTS["buildup"] # We always include the full loop if it exists if "full_loop" in SELECTED_VARIANTS: sections_to_include["full_loop"] = SELECTED_VARIANTS["full_loop"] if include_breakdown and "breakdown" in SELECTED_VARIANTS: sections_to_include["breakdown"] = SELECTED_VARIANTS["breakdown"] if include_drop and "drop" in SELECTED_VARIANTS: sections_to_include["drop"] = SELECTED_VARIANTS["drop"] if include_outro and "outro" in SELECTED_VARIANTS: sections_to_include["outro"] = SELECTED_VARIANTS["outro"] if not sections_to_include: return "Error: No sections selected or available", None, None progress(0.3, desc="Creating track structure...") # Create the final track final_track = None # Define the order of sections section_order = ["intro", "buildup", "full_loop", "breakdown", "drop", "outro"] # Process each section in order for section_name in section_order: if section_name not in sections_to_include: continue progress( 0.4 + 0.1 * section_order.index(section_name) / len(section_order), desc=f"Processing {section_name}...", ) # Get the selected variant config variant_config = sections_to_include[section_name] # Create temporary copy of stems to avoid modifying the originals stems_copy = {k: v for k, v in UPLOADED_STEMS.items()} # Create audio for this section from utils import create_section_from_json section_audio = create_section_from_json(variant_config, stems_copy) # Add to final track if final_track is None: final_track = section_audio else: final_track = final_track.append(section_audio, crossfade=crossfade_ms) progress(0.9, desc="Exporting final track...") # Export the final track full_track_path = os.path.join(TEMP_DIR, output_track_name) final_track.export(full_track_path, format="wav") # Create track summary sections_list = list(sections_to_include.keys()) track_duration = len(final_track) / 1000 # in seconds track_summary = { "Sections included": sections_list, "Total sections": len(sections_list), "Duration": f"{int(track_duration // 60)}:{int(track_duration % 60):02d}", "Crossfade": f"{crossfade_ms} ms", } progress(1.0, desc="Complete!") return ( "Track generated successfully!", full_track_path, json.dumps(track_summary, indent=2), ) except Exception as e: return f"Error generating track: {str(e)}", None, None def generate_full_loop_variants(bpm_value, bars_value, p_value, progress=gr.Progress()): """Generate variants for the full loop section""" return generate_section_variants_handler( "full_loop", bpm_value, bars_value, p_value, progress ) def create_section_ui(section_name, bpm_default, bars_default, p_default): """Helper function to create UI elements for a section.""" with gr.Accordion(f"Generate {section_name.capitalize()} Variants", open=False): with gr.Row(): with gr.Column(scale=1): gr.Markdown(f"### {section_name.capitalize()} Parameters") bpm_slider = gr.Slider( label="BPM (Beats Per Minute)", minimum=60, maximum=180, value=bpm_default, step=1, ) bars_slider = gr.Slider( label="Number of Bars", minimum=4, maximum=64, value=bars_default, step=4 ) p_slider = gr.Slider( label="Variation Parameter (p)", minimum=0, maximum=1, value=p_default, step=0.1, ) generate_btn = gr.Button( f"Generate {section_name.capitalize()} Variants", variant="primary" ) with gr.Column(scale=2): status = gr.Textbox(label="Status", interactive=False) descriptions = gr.JSON(label="Variant Descriptions") # Create empty lists to store the audio and checkbox components variant_audio_list = [] select_btn_list = [] with gr.Row(): for i in range(1, 5): with gr.Column(): gr.Markdown(f"### Variant {i}") # Create and immediately append each component to its respective list variant_audio = gr.Audio(label=f"Variant {i}", interactive=True) variant_audio_list.append(variant_audio) select_btn = gr.Checkbox(label=f"Select Variant {i}") select_btn_list.append(select_btn) return { "bpm_slider": bpm_slider, "bars_slider": bars_slider, "p_slider": p_slider, "generate_btn": generate_btn, "status": status, "descriptions": descriptions, "variant_audio": variant_audio_list, # Return the complete list of audio components "select_btn": select_btn_list, # Return the complete list of checkbox components } def setup_section_event_handlers(section_name, section_ui): """Setup event handlers for a given section.""" section_type = gr.State(section_name) # Generate button click event section_ui["generate_btn"].click( fn=generate_section_variants_handler, inputs=[ section_type, section_ui["bpm_slider"], section_ui["bars_slider"], section_ui["p_slider"], ], outputs=[ section_ui["status"], *section_ui["variant_audio"], section_ui["descriptions"], ], ) # Selection buttons change events for i, select_btn in enumerate(section_ui["select_btn"], start=1): variant_num = gr.State(i) select_btn.change( fn=select_variant, inputs=[section_type, variant_num, gr.State(False)], outputs=[section_ui["status"]], ) # Load section configuration from a JSON file or string with open("final_arrangement.json", "r") as f: section_config_json = f.read() sections = json.loads(section_config_json) # Create Gradio interface with gr.Blocks(title="Interactive Music Track Generator") as demo: gr.Markdown("# Interactive Music Track Generator") gr.Markdown( "Upload your WAV stems, generate variants for each section, and create a full track" ) # Global variables for UI state stem_list = gr.State([]) with gr.Tab("1. Upload Stems"): with gr.Row(): with gr.Column(): # File upload section gr.Markdown("### Upload Files") gr.Markdown("Drag and drop your WAV stem files here") file_input = gr.File( label="WAV Stems", file_count="multiple", file_types=[".wav"] ) upload_btn = gr.Button("Upload and Process Files", variant="primary") with gr.Column(): upload_status = gr.Textbox(label="Upload Status", interactive=False) stem_display = gr.JSON(label="Available Stems") with gr.Tab("1.1. ⏳ Finalise Sections Arrangement (In Progress)"): gr.Markdown("### 🎶 Finalise Sections Arrangement") gr.Markdown( "🎛️ Use the diagram below to adjust the arrangement of your sections. Click on a section to edit its properties." ) # Editable diagram with progress indicator with gr.Row(): gr.Markdown("⏳ In Progress: Adjusting sections...") edm_arrangement_tab() with gr.Tab("2. Generate Section Variants"): # Example of creating UI for sections dynamically section_uis = {} for section_name, params in sections.items(): section_uis[section_name] = create_section_ui( section_name, params["bpm"], params["bars"], params["p"] ) setup_section_event_handlers(section_name, section_uis[section_name]) with gr.Tab("3. Create Full Track"): with gr.Row(): with gr.Column(): gr.Markdown("### Track Settings") crossfade_ms = gr.Slider( label="Crossfade Duration (ms)", minimum=0, maximum=2000, value=500, step=100, ) output_track_name = gr.Textbox( label="Output Filename", value="full_track_output.wav", placeholder="e.g., full_track_output.wav", ) gr.Markdown("### Sections to Include") include_intro = gr.Checkbox(label="Include Intro", value=True) include_buildup = gr.Checkbox(label="Include buildup", value=True) include_breakdown = gr.Checkbox(label="Include breakdown", value=True) include_drop = gr.Checkbox(label="Include drop", value=True) include_outro = gr.Checkbox(label="Include Outro", value=True) generate_track_btn = gr.Button( "Generate Full Track", variant="primary", scale=2 ) with gr.Column(): track_status = gr.Textbox(label="Status", interactive=False) track_summary = gr.JSON(label="Track Summary") full_track_audio = gr.Audio(label="Generated Full Track") # Event handlers upload_btn.click( fn=process_uploaded_files, inputs=[file_input], outputs=[upload_status, stem_display], ) # Generate full track generate_track_btn.click( fn=generate_full_track, inputs=[ crossfade_ms, output_track_name, include_intro, include_breakdown, include_buildup, include_drop, include_outro, ], outputs=[track_status, full_track_audio, track_summary], ) if __name__ == "__main__": demo.launch()