updated

app.py

@@ -6,12 +6,20 @@ 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
+}
+
+original_time_cache = {}
+
+
 def set_seed(seed: int = 42):
     random.seed(seed)
     np.random.seed(seed)
@@ -23,7 +31,6 @@ 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()
@@ -31,7 +38,6 @@ def cleanup_gpu():
 
 
 def cleanup_temp_files():
-    """Clean up old temporary audio files."""
     import glob
     import time
     temp_dir = tempfile.gettempdir()
@@ -47,6 +53,8 @@ def cleanup_temp_files():
 
 _generator = None
 _processor = None
+_original_generator = None
+_original_processor = None
 
 
 def load_model():
@@ -88,6 +96,44 @@ 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,
+            device=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)
@@ -160,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  # Scale to 95% to avoid clipping
+            audio_data = audio_data / max_val * 0.95
         
-        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)}]")
@@ -180,6 +226,7 @@ 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:
@@ -192,9 +239,150 @@ 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_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:
+        cache_key = get_cache_key(text_prompt, duration, guidance_scale)
+        
+        generator, processor = load_model()
+        model_name = "Compressed (S)"
+
+        print(f"[GENERATION] Starting batch generation using {model_name} model...")
+        print(f"[GENERATION] Prompt: '{text_prompt}'")
+        print(f"[GENERATION] Duration: {duration}s")
+        print(f"[GENERATION] Guidance scale: {guidance_scale}")
+
+        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',
+        }
+
+        prompts = [text_prompt] * 4
+        start_time = time.time()
+        outputs = generator(
+            prompts,
+            batch_size=4,
+            generate_kwargs=generation_params
+        )
+        generation_time = time.time() - start_time
+
+        print(f"[GENERATION] Batch 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}")
+
+        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,
+                    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}"
+        
+        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") as demo:
     gr.Markdown("# 🎵 MusicGen Large Music Generator")
-    gr.Markdown("Generate music from text descriptions using Facebook's MusicGen Large model with elastic compression.")
+    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():
@@ -204,7 +392,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,
@@ -222,13 +410,18 @@ with gr.Blocks(title="MusicGen Large - Music Generation") as demo:
                     info="Higher values follow prompt more closely"
                 )
 
-            generate_btn = gr.Button("🎵 Generate Music", variant="primary", size="lg")
+            generate_btn = gr.Button("🎵 Generate 4 Music Variants", variant="primary", size="lg")
 
         with gr.Column():
-            audio_output = gr.Audio(
-                label="Generated Music",
-                type="numpy"
-            )
+            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("""
@@ -238,10 +431,13 @@ 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_music,
+        fn=generate_simple,
         inputs=[text_input, duration, guidance_scale],
-        outputs=audio_output,
+        outputs=[audio_output1, audio_output2, audio_output3, audio_output4, generation_info],
         show_progress=True
     )
 

requirements.txt

@@ -4,6 +4,7 @@
 
 torch
 thestage
-elastic_models[nvidia]
+# elastic_models[nvidia]
 scipy
 transformers
+soundfile