updated

app.py

@@ -12,7 +12,6 @@ os.environ['ELASTIC_LOG_LEVEL'] = 'DEBUG'
 from transformers import AutoProcessor, pipeline
 from elastic_models.transformers import MusicgenForConditionalGeneration
 
-
 MODEL_CONFIG = {
     'cost_per_hour': 1.8,  # $1.8 per hour
 }
@@ -207,9 +206,9 @@ def generate_music(text_prompt, duration=10, guidance_scale=3.0):
         max_val = np.max(np.abs(audio_data))
         if max_val > 0:
             audio_data = audio_data / max_val * 0.95
-        
+
         audio_data = (audio_data * 32767).astype(np.int16)
-        
+
         print(f"[GENERATION] Final audio shape: {audio_data.shape}")
         print(f"[GENERATION] Audio range: [{np.min(audio_data)}, {np.max(audio_data)}]")
         print(f"[GENERATION] Audio dtype: {audio_data.dtype}")
@@ -225,7 +224,7 @@ def generate_music(text_prompt, duration=10, guidance_scale=3.0):
             file_size = os.path.getsize(temp_path)
             print(f"[GENERATION] Audio saved to: {temp_path}")
             print(f"[GENERATION] File size: {file_size} bytes")
-            
+
             # Try returning numpy format instead
             print(f"[GENERATION] Returning numpy tuple: ({sample_rate}, audio_array)")
             return (sample_rate, audio_data)
@@ -265,7 +264,7 @@ def get_cache_key(prompt, duration, guidance_scale):
 def generate_music_batch(text_prompt, duration=10, guidance_scale=3.0, model_mode="compressed"):
     try:
         cache_key = get_cache_key(text_prompt, duration, guidance_scale)
-        
+
         generator, processor = load_model()
         model_name = "Compressed (S)"
 
@@ -301,18 +300,18 @@ def generate_music_batch(text_prompt, duration=10, guidance_scale=3.0, model_mod
 
         audio_variants = []
         sample_rate = outputs[0]['sampling_rate']
-        
+
         for i, output in enumerate(outputs):
             audio_data = output['audio']
-            
-            print(f"[GENERATION] Processing variant {i+1} audio shape: {audio_data.shape}")
-            
+
+            print(f"[GENERATION] Processing variant {i + 1} audio shape: {audio_data.shape}")
+
             if hasattr(audio_data, 'cpu'):
                 audio_data = audio_data.cpu().numpy()
 
             if len(audio_data.shape) == 3:
                 audio_data = audio_data[0]
-                
+
             if len(audio_data.shape) == 2:
                 if audio_data.shape[0] < audio_data.shape[1]:
                     audio_data = audio_data.T
@@ -320,31 +319,31 @@ def generate_music_batch(text_prompt, duration=10, guidance_scale=3.0, model_mod
                     audio_data = audio_data[:, 0]
                 else:
                     audio_data = audio_data.flatten()
-            
+
             audio_data = audio_data.flatten()
-            
+
             max_val = np.max(np.abs(audio_data))
             if max_val > 0:
                 audio_data = audio_data / max_val * 0.95
-            
+
             audio_data = (audio_data * 32767).astype(np.int16)
             audio_variants.append((sample_rate, audio_data))
-            
-            print(f"[GENERATION] Variant {i+1} final shape: {audio_data.shape}")
+
+            print(f"[GENERATION] Variant {i + 1} final shape: {audio_data.shape}")
 
         comparison_message = ""
-        
+
         if cache_key in original_time_cache:
             original_time = original_time_cache[cache_key]
             cost_info = calculate_cost_savings(generation_time, original_time)
-            
+
             comparison_message = f"💰 Cost Savings: ${cost_info['savings']:.4f} ({cost_info['savings_percent']:.1f}%) - Compressed: ${cost_info['compressed_cost']:.4f} vs Original: ${cost_info['original_cost']:.4f}"
             print(f"[COST] Savings: ${cost_info['savings']:.4f} ({cost_info['savings_percent']:.1f}%)")
         else:
             try:
                 print(f"[TIMING] Measuring original model speed for comparison...")
                 original_generator, original_processor = load_original_model()
-                
+
                 original_start = time.time()
                 original_outputs = original_generator(
                     prompts,
@@ -352,25 +351,26 @@ def generate_music_batch(text_prompt, duration=10, guidance_scale=3.0, model_mod
                     generate_kwargs=generation_params
                 )
                 original_time = time.time() - original_start
-                
+
                 original_time_cache[cache_key] = original_time
-                
+
                 cost_info = calculate_cost_savings(generation_time, original_time)
                 comparison_message = f"💰 Cost Savings: ${cost_info['savings']:.4f} ({cost_info['savings_percent']:.1f}%) - Compressed: ${cost_info['compressed_cost']:.4f} vs Original: ${cost_info['original_cost']:.4f}"
-                print(f"[COST] First comparison - Savings: ${cost_info['savings']:.4f} ({cost_info['savings_percent']:.1f}%)")
+                print(
+                    f"[COST] First comparison - Savings: ${cost_info['savings']:.4f} ({cost_info['savings_percent']:.1f}%)")
                 print(f"[TIMING] Original: {original_time:.2f}s, Compressed: {generation_time:.2f}s")
-                
+
                 del original_generator, original_processor
                 cleanup_gpu()
                 print(f"[CLEANUP] Original model cleaned up after timing measurement")
-                
+
             except Exception as e:
                 print(f"[WARNING] Could not measure original timing: {e}")
                 compressed_cost = calculate_generation_cost(generation_time, 'S')
                 comparison_message = f"💸 Compressed Cost: ${compressed_cost:.4f} (could not compare with original)"
 
         generation_info = f"✅ Generated 4 variants in {generation_time:.2f}s\n{comparison_message}"
-        
+
         return audio_variants[0], audio_variants[1], audio_variants[2], audio_variants[3], generation_info
 
     except Exception as e:
@@ -382,7 +382,8 @@ def generate_music_batch(text_prompt, duration=10, guidance_scale=3.0, model_mod
 
 with gr.Blocks(title="MusicGen Large - Music Generation") as demo:
     gr.Markdown("# 🎵 MusicGen Large Music Generator")
-    gr.Markdown("Generate music from text descriptions using Facebook's MusicGen Large model accelerated by TheStage for 2.3x faster performance")
+    gr.Markdown(
+        "Generate music from text descriptions using Facebook's MusicGen Large model accelerated by TheStage for 2.3x faster performance")
 
     with gr.Row():
         with gr.Column():
@@ -392,7 +393,7 @@ with gr.Blocks(title="MusicGen Large - Music Generation") as demo:
                 lines=3,
                 value="A groovy funk bassline with a tight drum beat"
             )
-                
+
             with gr.Row():
                 duration = gr.Slider(
                     minimum=5,
@@ -410,15 +411,15 @@ with gr.Blocks(title="MusicGen Large - Music Generation") as demo:
                     info="Higher values follow prompt more closely"
                 )
 
-            generate_btn = gr.Button("🎵 Generate 4 Music Variants", variant="primary", size="lg")
+            generate_btn = gr.Button("🎵 Generate Music", variant="primary", size="lg")
 
         with gr.Column():
             generation_info = gr.Markdown("Ready to generate music variants with cost comparison vs original model")
-            
+
             with gr.Row():
                 audio_output1 = gr.Audio(label="Variant 1", type="numpy")
                 audio_output2 = gr.Audio(label="Variant 2", type="numpy")
-            
+
             with gr.Row():
                 audio_output3 = gr.Audio(label="Variant 3", type="numpy")
                 audio_output4 = gr.Audio(label="Variant 4", type="numpy")
@@ -431,9 +432,11 @@ with gr.Blocks(title="MusicGen Large - Music Generation") as demo:
                 - Duration is limited to 30 seconds for faster generation
                 """)
 
+
     def generate_simple(text_prompt, duration, guidance_scale):
         return generate_music_batch(text_prompt, duration, guidance_scale, "compressed")
 
+
     generate_btn.click(
         fn=generate_simple,
         inputs=[text_input, duration, guidance_scale],