Update app.py

app.py

@@ -55,15 +55,23 @@ def load_musicgen_model():
 
 processor, music_model, device = load_musicgen_model()
 
-# Function to chunk audio into 10-second segments
+# Function to chunk audio into segments
 def chunk_audio(audio_path, chunk_duration=10):
-    """Split audio into 10-second chunks and return list of chunk file paths"""
+    """Split audio into chunks and return list of chunk file paths"""
     try:
         # Load audio file
         audio = AudioSegment.from_file(audio_path)
         duration_ms = len(audio)
         chunk_ms = chunk_duration * 1000
         
+        # Validate chunk duration
+        if chunk_duration <= 0:
+            raise ValueError("Chunk duration must be positive")
+        
+        if chunk_duration > duration_ms / 1000:
+            # If chunk duration is longer than audio, return the whole audio
+            return [audio_path], 1
+            
         chunks = []
         chunk_files = []
         
@@ -299,15 +307,15 @@ def process_chunk(chunk_path, chunk_idx, total_chunks, generate_audio=True):
         }
 
 # Function to get predictions for all chunks
-def get_predictions(audio_input, generate_audio=True):
-    # Chunk the audio into 10-second segments
-    chunk_files, total_chunks = chunk_audio(audio_input, chunk_duration=10)
+def get_predictions(audio_input, generate_audio=True, chunk_duration=10):
+    # Chunk the audio into segments
+    chunk_files, total_chunks = chunk_audio(audio_input, chunk_duration)
     
     results = []
     
     # Process each chunk
     for i, chunk_path in enumerate(chunk_files):
-        print(f"Processing chunk {i+1}/{total_chunks}")
+        print(f"Processing chunk {i+1}/{total_chunks} ({chunk_duration}s each)")
         result = process_chunk(chunk_path, i, total_chunks, generate_audio)
         results.append(result)
     
@@ -328,11 +336,20 @@ def clear_all():
 # Create the Gradio interface with proper output handling
 with gr.Blocks(title="Affective Virtual Environments - Chunked Processing") as interface:
     gr.Markdown("# Affective Virtual Environments")
-    gr.Markdown("Create an AVE using your voice. Audio is split into 5-second chunks, with separate predictions and generations for each segment.")
+    gr.Markdown("Create an AVE using your voice. Audio is split into chunks, with separate predictions and generations for each segment.")
     
     with gr.Row():
         audio_input = gr.Audio(label="Input Audio", type="filepath", sources=["microphone", "upload"])
         with gr.Column():
+            # Add chunk duration input
+            chunk_duration_input = gr.Number(
+                label="Chunk Duration (seconds)",
+                value=10,
+                minimum=1,
+                maximum=60,
+                step=1,
+                info="Duration of each audio segment to process (1-60 seconds)"
+            )
             # Add checkbox for audio generation
             generate_audio_checkbox = gr.Checkbox(
                 label="Generate Audio (may take longer)", 
@@ -356,8 +373,8 @@ with gr.Blocks(title="Affective Virtual Environments - Chunked Processing") as i
     output_containers = []
     group_components = []  # Store group components separately
     
-    # We'll create up to 10 chunk slots (adjust as needed)
-    for i in range(10):
+    # We'll create up to 20 chunk slots to accommodate different chunk durations
+    for i in range(20):
         with gr.Group(visible=False) as chunk_group:
             gr.Markdown(f"### Chunk {i+1} Results")
             with gr.Row():
@@ -378,17 +395,21 @@ with gr.Blocks(title="Affective Virtual Environments - Chunked Processing") as i
             'music': audio_output
         })
     
-    def process_and_display(audio_input, generate_audio):
+    def process_and_display(audio_input, generate_audio, chunk_duration):
+        # Validate chunk duration
+        if chunk_duration is None or chunk_duration <= 0:
+            chunk_duration = 10
+        
         # Show loading indicator
-        yield [gr.HTML("""
+        yield [gr.HTML(f"""
             <div style="text-align: center; margin: 20px;">
-                <p style="font-size: 18px; color: #4a4a4a;">Processing audio chunks...</p>
+                <p style="font-size: 18px; color: #4a4a4a;">Processing audio in {chunk_duration}-second chunks...</p>
                 <div style="border: 4px solid #f3f3f3; border-top: 4px solid #3498db; border-radius: 50%; width: 30px; height: 30px; animation: spin 2s linear infinite; margin: 0 auto;"></div>
                 <style>@keyframes spin {0% { transform: rotate(0deg); } 100% { transform: rotate(360deg); }}</style>
             </div>
         """)] + [gr.Group(visible=False)] * len(group_components) + [None] * (len(output_containers) * 5)
         
-        results = get_predictions(audio_input, generate_audio)
+        results = get_predictions(audio_input, generate_audio, chunk_duration)
         
         # Initialize outputs list
         outputs = []
@@ -421,7 +442,7 @@ with gr.Blocks(title="Affective Virtual Environments - Chunked Processing") as i
     # Set up the button click
     process_btn.click(
         fn=process_and_display,
-        inputs=[audio_input, generate_audio_checkbox],
+        inputs=[audio_input, generate_audio_checkbox, chunk_duration_input],
         outputs=[loading_indicator] + group_components + [comp for container in output_containers for comp in [
             container['emotion'],
             container['transcription'],
@@ -441,7 +462,7 @@ with gr.Blocks(title="Affective Virtual Environments - Chunked Processing") as i
             container['sentiment'],
             container['image'],
             container['music']
-        ]] + [loading_indicator]
+        ]] + [loading_indicator] + [gr.Number(value=10)]  # Reset chunk duration to default
     )
 
 interface.launch()