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 from langsmith import traceable 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) stem_list = {key: data["stems"] for key, data in variants.items()} 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, stem_list ) except Exception as e: return f"Error generating variants: {str(e)}", None, None, None, None, None, None @traceable(run_type="chain", name="groq_call") def select_variant(section_type, variant_num, append): """Select a variant for a specific section, with an option to append.""" global ALL_VARIANTS, SELECTED_VARIANTS try: if section_type not in ALL_VARIANTS: return f"No variants generated for {section_type} yet", "None" variant_key = f"variant{variant_num}" if variant_key not in ALL_VARIANTS[section_type]: return f"Variant {variant_num} not found for {section_type}", "None" # Store the selection variant_info = ALL_VARIANTS[section_type][variant_key] SELECTED_VARIANTS[section_type] = variant_info["config"] # Create a display string for the selected variant display_text = f"Selected: Variant {variant_num}" if "description" in variant_info: display_text += f" - {variant_info['description']}" # Log selection details for tracing logger.info(f"Selected variant for {section_type}: {variant_key}") logger.info(f"Variant description: {variant_info.get('description', 'No description')}") logger.info(f"Stems used: {variant_info.get('config', {}).get('stems', [])}") return f"Selected variant {variant_num} for {section_type}", display_text except Exception as e: logger.error(f"Error selecting variant: {str(e)}") return f"Error selecting variant: {str(e)}", "None" def generate_full_track( crossfade_ms, output_track_name, *section_flags, progress=gr.Progress(), ): """Generate the full track from selected variants""" global TEMP_DIR, SELECTED_VARIANTS, UPLOADED_STEMS, sections if not TEMP_DIR or not os.path.exists(TEMP_DIR): return "Error: No stems loaded", None, None, None try: progress(0.1, desc="Preparing to generate full track...") # Get section names from the configuration section_names = list(sections.keys()) # Map section flags to section names sections_to_include = {} for section_name, include_flag in zip(section_names, section_flags): if include_flag and section_name in SELECTED_VARIANTS: sections_to_include[section_name] = SELECTED_VARIANTS[section_name] if not sections_to_include: return "Error: No sections selected or available", None, None, None progress(0.3, desc="Creating track structure...") # Create the final track final_track = None # Process each section in order from the configuration for section_name in section_names: if section_name not in sections_to_include: continue progress( 0.4 + 0.1 * section_names.index(section_name) / len(section_names), 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", "Section Details": { section: { "BPM": sections[section]["bpm"], "Bars": sections[section]["bars"], "Volume Automation": sections[section]["volume_automation"], "Curve": sections[section]["curve"] } for section in sections_list } } 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, params): """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") gr.Markdown(f"**Volume Automation**: {params['volume_automation']}") gr.Markdown(f"**Curve Type**: {params['curve']}") bpm_slider = gr.Slider( label="BPM (Beats Per Minute)", minimum=60, maximum=180, value=params['bpm'], step=1, ) bars_slider = gr.Slider( label="Number of Bars", minimum=4, maximum=64, value=params['bars'], step=4 ) p_slider = gr.Slider( label="Variation Parameter (p)", minimum=0, maximum=1, value=params['p'], 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") stem_list = gr.JSON(label="Stem List") selected_variant = gr.Text(label="Selected Variant", value="None", interactive=False) with gr.Row(): variant_audio_list = [] select_btn_list = [] for i in range(1, 5): with gr.Column(): gr.Markdown(f"### Variant {i}") variant_audio = gr.Audio(label=f"Variant {i}", interactive=True) variant_audio_list.append(variant_audio) select_btn = gr.Radio( choices=[f"Select Variant {i}"], label="", value=None, interactive=True ) 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, "select_btn": select_btn_list, "selected_variant": selected_variant, "stem_list": stem_list } def setup_section_event_handlers(section_name, section_ui, selected_variants_display): """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"], section_ui["stem_list"] ], ) # 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"], section_ui["selected_variant"]], ) # Update selected variants display whenever a variant is selected def update_selected_variants_display(): selected = {} for s in SELECTED_VARIANTS: if s in ALL_VARIANTS: for v in range(1, 5): variant_key = f"variant{v}" if variant_key in ALL_VARIANTS[s]: if s not in selected: selected[s] = {} selected[s]["Selected Variant"] = ALL_VARIANTS[s][variant_key].get("description", f"Variant {v}") return json.dumps(selected, indent=2) select_btn.change( fn=update_selected_variants_display, inputs=[], outputs=[selected_variants_display], ) # 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"): gr.Markdown("### Generate and Select Variants for Each Section") gr.Markdown("Generate variants for each section and select which one to use in the final track") # Create UI for sections dynamically based on final_arrangement.json section_uis = {} # Create selected variants display first selected_variants_display = gr.JSON( label="Selected Variants", value={"No variants selected yet": "Generate and select variants in sections below"} ) for section_name, params in sections.items(): section_uis[section_name] = create_section_ui(section_name, params) setup_section_event_handlers(section_name, section_uis[section_name], selected_variants_display) 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") section_checkboxes = {} for section_name in sections.keys(): section_checkboxes[section_name] = gr.Checkbox( label=f"Include {section_name.capitalize()}", value=True ) gr.Markdown("### Selected Variants Summary") selected_variants_display = gr.JSON( label="Selected Variants", value={"No variants selected yet": "Generate and select variants in Section 2"} ) 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, *[section_checkboxes[section] for section in sections.keys()] ], outputs=[track_status, full_track_audio, track_summary], ) if __name__ == "__main__": demo.launch()