updated

app.py

@@ -6,25 +6,11 @@ import random
 import os
 import tempfile
 import soundfile as sf
-import time
 
 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 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}
-
 
 def set_seed(seed: int = 42):
     random.seed(seed)
@@ -37,6 +23,7 @@ def set_seed(seed: int = 42):
 
 
 def cleanup_gpu():
+    """Clean up GPU memory to avoid TensorRT conflicts."""
     if torch.cuda.is_available():
         torch.cuda.empty_cache()
         torch.cuda.synchronize()
@@ -44,6 +31,7 @@ def cleanup_gpu():
 
 
 def cleanup_temp_files():
+    """Clean up old temporary audio files."""
     import glob
     import time
     temp_dir = tempfile.gettempdir()
@@ -59,8 +47,6 @@ def cleanup_temp_files():
 
 _generator = None
 _processor = None
-_original_generator = None
-_original_processor = None
 
 
 def load_model():
@@ -102,43 +88,6 @@ def load_model():
     return _generator, _processor
 
 
-def load_original_model():
-    global _original_generator, _original_processor
-
-    if _original_generator is None:
-        print("[ORIGINAL MODEL] Starting original model initialization...")
-        cleanup_gpu()
-
-        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-        print(f"[ORIGINAL MODEL] Using device: {device}")
-
-        print("[ORIGINAL MODEL] Loading processor...")
-        _original_processor = AutoProcessor.from_pretrained(
-            "facebook/musicgen-large"
-        )
-        from transformers import MusicgenForConditionalGeneration as HFMusicgenForConditionalGeneration
-
-        print("[ORIGINAL MODEL] Loading original model...")
-        model = HFMusicgenForConditionalGeneration.from_pretrained(
-            "facebook/musicgen-large",
-            torch_dtype=torch.float16,
-        ).to(device)
-
-        model.eval()
-
-        print("[ORIGINAL MODEL] Creating pipeline...")
-        _original_generator = pipeline(
-            task="text-to-audio",
-            model=model,
-            tokenizer=_original_processor.tokenizer,
-            device=device,
-        )
-
-        print("[ORIGINAL MODEL] Original model initialization completed successfully")
-
-    return _original_generator, _original_processor
-
-
 def calculate_max_tokens(duration_seconds):
     token_rate = 50
     max_new_tokens = int(duration_seconds * token_rate)
@@ -211,9 +160,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 / max_val * 0.95  # Scale to 95% to avoid clipping
 
-        audio_data = (audio_data * 32767).astype(np.int16)
+        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)}]")
@@ -231,7 +180,6 @@ def generate_music(text_prompt, duration=10, guidance_scale=3.0):
             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)
         else:
@@ -244,194 +192,56 @@ def generate_music(text_prompt, duration=10, guidance_scale=3.0):
         return None
 
 
-def calculate_generation_cost(generation_time_seconds, mode='S'):
-    hours = generation_time_seconds / 3600
-    cost_per_hour = MODEL_CONFIG['cost_per_hour']
-    return hours * cost_per_hour
-
-
-def calculate_cost_savings(compressed_time, original_time):
-    compressed_cost = calculate_generation_cost(compressed_time, 'S')
-    original_cost = calculate_generation_cost(original_time, 'original')
-    savings = original_cost - compressed_cost
-    savings_percent = (savings / original_cost * 100) if original_cost > 0 else 0
-    return {
-        'compressed_cost': compressed_cost,
-        'original_cost': original_cost,
-        'savings': savings,
-        'savings_percent': savings_percent
-    }
-
-
-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}"
-
-
-def generate_music_batch(text_prompt, duration=10, guidance_scale=3.0, model_mode="compressed"):
-    try:
-
-        generator, processor = load_model()
-        model_name = "Compressed (S)"
-
-        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)
-        print(f"[GENERATION] Using seed: 42")
-
-        max_new_tokens = calculate_max_tokens(duration)
-
-        generation_params = {
-            'do_sample': True,
-            'guidance_scale': guidance_scale,
-            'max_new_tokens': max_new_tokens,
-            'min_new_tokens': max_new_tokens,
-            'cache_implementation': 'paged',
-        }
-
-        batch_size = 4 if MODEL_CONFIG['batch_mode'] else 1
-        prompts = [text_prompt] * batch_size
-        
-        start_time = time.time()
-        outputs = generator(
-            prompts,
-            batch_size=batch_size,
-            generate_kwargs=generation_params
-        )
-        generation_time = time.time() - start_time
-
-        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
-                if audio_data.shape[1] > 1:
-                    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}")
-
-        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:
-        print(f"[ERROR] Batch generation failed: {str(e)}")
-        cleanup_gpu()
-        error_msg = f"❌ Generation failed: {str(e)}"
-        return None, None, None, None, error_msg
-
-
-with gr.Blocks(title="MusicGen Large - Music Generation", theme=gr.themes.Soft()) as demo:
+with gr.Blocks(title="MusicGen Large - Music Generation") as demo:
     gr.Markdown("# 🎵 MusicGen Large Music Generator")
-    
-    gr.Markdown(
-        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)"
+    gr.Markdown("Generate music from text descriptions using Facebook's MusicGen Large model with elastic compression.")
+
+    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"
             )
-            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"
+
+            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():
+            audio_output = gr.Audio(
+                label="Generated Music",
+                type="numpy"
             )
-        
-        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")
+
+            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
+                """)
 
     generate_btn.click(
-        fn=generate_simple,
+        fn=generate_music,
         inputs=[text_input, duration, guidance_scale],
-        outputs=[audio_output1, audio_output2, audio_output3, audio_output4, generation_info],
+        outputs=audio_output,
         show_progress=True
     )
 
@@ -453,12 +263,7 @@ with gr.Blocks(title="MusicGen Large - Music Generation", theme=gr.themes.Soft()
     gr.Markdown("---")
     gr.Markdown("""
     <div style="text-align: center; color: #666; font-size: 12px; margin-top: 2rem;">
-        <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>
+        <strong>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>