Update app.py

app.py

@@ -3,7 +3,7 @@ import json
 import asyncio
 import base64
 import numpy as np
-from flask import Flask, request, send_file, jsonify
+from flask import Flask, request, Response, jsonify, stream_with_context
 from io import BytesIO
 from utils import decode, pcm_to_wav
 from datetime import datetime
@@ -30,6 +30,7 @@ class MockLiveMusicSession:
         self.model = model
         self.callbacks = None
         self.is_playing = False
+        self.setup_complete = False
 
     async def setWeightedPrompts(self, params):
         print(f"Setting prompts: {params['weightedPrompts']}")
@@ -40,6 +41,8 @@ class MockLiveMusicSession:
     def play(self):
         self.is_playing = True
         print("Starting music generation")
+        if self.callbacks and self.callbacks.get('onmessage'):
+            self.callbacks['onmessage']({'setupComplete': True})
 
     def close(self):
         self.is_playing = False
@@ -53,6 +56,8 @@ ai = GoogleGenAI({
 })
 model = 'lyria-realtime-exp'
 sample_rate = 48000
+channels = 2
+bits_per_sample = 16
 
 # Genre-specific parameters
 GENRE_PARAMS = {
@@ -63,8 +68,30 @@ GENRE_PARAMS = {
     "slowed": {"base_freq": 55, "mod_freq": 0.2, "amplitude": 0.3}
 }
 
+def generate_audio_chunk(prompts, config, total_duration):
+    slowed_factor = config.get('slowed_factor', 1.0)
+    chunk_duration = 5 * slowed_factor  # 5 seconds per chunk
+    samples_per_chunk = int(sample_rate * chunk_duration * channels)
+    t = np.linspace(0, chunk_duration, samples_per_chunk // channels, False)
+    
+    # Weighted average of genre parameters
+    total_weight = sum(p['weight'] for p in prompts)
+    base_freq = sum(p['weight'] * GENRE_PARAMS.get(p['text'], GENRE_PARAMS["Synthwave"])['base_freq'] for p in prompts) / total_weight
+    mod_freq = sum(p['weight'] * GENRE_PARAMS.get(p['text'], GENRE_PARAMS["Synthwave"])['mod_freq'] for p in prompts) / total_weight
+    amplitude = sum(p['weight'] * GENRE_PARAMS.get(p['text'], GENRE_PARAMS["Synthwave"])['amplitude'] for p in prompts) / total_weight
+    amplitude *= 0.5 if slowed_factor < 1 else 1.0  # Reduce for slowed effect
+    
+    # Generate layered audio with 3 frequencies
+    chunk = np.zeros(samples_per_chunk, dtype=np.float32)
+    for _ in range(3):
+        freq_offset = np.random.uniform(-10, 10)
+        chunk[:samples_per_chunk//channels] += amplitude * np.sin(2 * np.pi * (base_freq + freq_offset + mod_freq * np.sin(2 * np.pi * 0.1 * t)) * t / sample_rate)
+    chunk = np.tile(chunk, channels)  # Duplicate for stereo
+    chunk = np.clip(chunk * 32768, -32768, 32767).astype(np.int16)  # Convert to 16-bit
+    return chunk.tobytes()
+
 @app.route('/generate', methods=['POST'])
-async def generate_music():
+def generate_music():
     try:
         data = request.get_json()
         if not data:
@@ -90,75 +117,36 @@ async def generate_music():
             } for i, prompt in enumerate(prompts)
         ]
 
-        collected_chunks = []
-        total_duration = 0
-        target_duration = 60  # 1 minute
-
-        async def mock_generate_chunks():
-            nonlocal collected_chunks, total_duration
-            slowed_factor = config.get('slowed_factor', 1.0)
-            chunk_duration = 5 * slowed_factor  # Increased to 5 seconds per chunk
-            samples_per_chunk = int(sample_rate * chunk_duration * 2)  # Stereo
-            t = np.linspace(0, chunk_duration, samples_per_chunk, False)
-            np.random.seed(None)  # Ensure randomness per run
-            while total_duration < target_duration:
-                total_weight = sum(p['weight'] for p in weighted_prompts)
-                base_freq = sum(p['weight'] * GENRE_PARAMS.get(p['text'], GENRE_PARAMS["Synthwave"])['base_freq'] for p in weighted_prompts) / total_weight
-                mod_freq = sum(p['weight'] * GENRE_PARAMS.get(p['text'], GENRE_PARAMS["Synthwave"])['mod_freq'] for p in weighted_prompts) / total_weight
-                amplitude = sum(p['weight'] * GENRE_PARAMS.get(p['text'], GENRE_PARAMS["Synthwave"])['amplitude'] for p in weighted_prompts) / total_weight
-                amplitude *= 0.5 if slowed_factor < 1 else 1.0  # Reduce for slowed effect
-                # Layer multiple frequencies
-                chunk = np.zeros(samples_per_chunk, dtype=np.int16)
-                for _ in range(3):  # Layer 3 sine waves
-                    freq_offset = np.random.uniform(-10, 10)  # Slight random variation
-                    chunk += (amplitude * 32767 * np.sin(2 * np.pi * (base_freq + freq_offset + mod_freq * np.sin(2 * np.pi * 0.1 * t)) * t / sample_rate)).astype(np.int16)
-                chunk = np.clip(chunk, -32768, 32767)  # Prevent overflow
-                encoded_chunk = base64.b64encode(chunk.tobytes()).decode('utf-8')
-                decoded_chunk = decode(encoded_chunk)
-                collected_chunks.append(decoded_chunk)
-                total_duration += chunk_duration
-                yield {'serverContent': {'audioChunks': [{'data': encoded_chunk}]}}
-                await asyncio.sleep(0.1 * slowed_factor)
-            yield {'close': True}
-
-        session = await ai.live.music.connect({
-            'model': model,
-            'callbacks': {
-                'onmessage': lambda e: None,
-                'onerror': lambda e: print(f"Error: {e}"),
-                'onclose': lambda: print("Session closed")
-            }
-        })
+        session = MockLiveMusicSession(model)
+        session.callbacks = {
+            'onmessage': lambda msg: None,
+            'onerror': lambda e: print(f"Error: {e}"),
+            'onclose': lambda: print("Session closed")
+        }
+
+        def generate_stream():
+            total_duration = 0
+            target_duration = 60  # 1 minute
+            session.setup_complete = True
+            yield json.dumps({'setupComplete': True}) + '\n'
+
+            while total_duration < target_duration and session.is_playing:
+                chunk_data = generate_audio_chunk(weighted_prompts, config, total_duration)
+                encoded_chunk = base64.b64encode(chunk_data).decode('utf-8')
+                message = {
+                    'serverContent': {
+                        'audioChunks': [{'data': encoded_chunk}]
+                    }
+                }
+                yield json.dumps(message) + '\n'
+                total_duration += 5 * config.get('slowed_factor', 1.0)
+                asyncio.run(asyncio.sleep(0.1))  # Simulate real-time generation
+
+            if session.callbacks and session.callbacks.get('onclose'):
+                session.callbacks['onclose']()
 
-        await session.setWeightedPrompts({'weightedPrompts': weighted_prompts})
-        await session.setMusicGenerationConfig({'musicGenerationConfig': config})
         session.play()
-
-        async for message in mock_generate_chunks():
-            if message.get('close'):
-                break
-            if message.get('serverContent', {}).get('audioChunks'):
-                chunk_data = message['serverContent']['audioChunks'][0]['data']
-                decoded_chunk = decode(chunk_data)
-                collected_chunks.append(decoded_chunk)
-                chunk_duration = len(decoded_chunk) / (sample_rate * 2 * 2)
-                total_duration += chunk_duration
-
-        total_length = sum(len(chunk) for chunk in collected_chunks)
-        combined_pcm = np.concatenate([np.frombuffer(chunk, dtype=np.int16) for chunk in collected_chunks])
-        combined_pcm_bytes = combined_pcm.tobytes()
-        wav_blob = pcm_to_wav(combined_pcm_bytes, 2, sample_rate, 16)
-
-        output = BytesIO(wav_blob)
-        timestamp = datetime.now().isoformat().replace(':', '-')
-        filename = f"prompt-dj-music-1min-{timestamp}.wav"
-
-        return send_file(
-            output,
-            mimetype='audio/wav',
-            as_attachment=True,
-            download_name=filename
-        )
+        return Response(stream_with_context(generate_stream()), mimetype='text/event-stream')
 
     except Exception as e:
         return jsonify({'error': str(e)}), 500