Update beat_generator/riffusion_cpu.py

beat_generator/riffusion_cpu.py

@@ -3,35 +3,51 @@ import soundfile as sf
 import os
 import uuid
 
+
 def generate_audio_loop(style, bpm):
+    # Ensure directory exists
+    outdir = "assets/generated_beats"
+    os.makedirs(outdir, exist_ok=True)
+
+    # Audio settings
     sr = 44100
-    sec = 8  # loop length
+    sec = 5   # shorter = safer for HF
+    samples = int(sr * sec)
 
-    t = np.linspace(0, sec, int(sr * sec))
+    t = np.linspace(0, sec, samples)
 
-    # style-based fake synth freq
+    # Pick a frequency based on style
     if style == "Country Rap":
-        freq = 90 + bpm
+        freq = 80 + bpm * 0.5
     elif style == "Trap":
-        freq = 140 + bpm / 2
+        freq = 120 + bpm * 0.3
     elif style == "Dark Outlaw":
-        freq = 50 + bpm / 3
+        freq = 60 + bpm * 0.4
     elif style == "Drill":
-        freq = 160 + bpm / 1.5
+        freq = 150 + bpm * 0.2
     else:
-        freq = 110 + bpm
+        freq = 100 + bpm * 0.2
 
-    wave = 0.2 * np.sin(2 * np.pi * freq * t)
+    # Main tone
+    waveform = 0.3 * np.sin(2 * np.pi * freq * t)
 
-    # add percussive clicks
-    clicks = np.random.choice([0, 1], size=wave.shape, p=[0.995, 0.005])
-    wave = wave + (clicks * 0.5)
+    # Add simple hi-hat clicks
+    for i in range(0, samples, int(sr * 0.2)):
+        if i < samples:
+            waveform[i:i+200] += 0.5 * np.random.uniform(-1, 1, 200)
 
-    wave = wave.astype(np.float32)
+    # Normalize to float32
+    waveform = waveform.astype(np.float32)
 
-    os.makedirs("assets/generated_beats", exist_ok=True)
-    outpath = f"assets/generated_beats/{uuid.uuid4().hex}.wav"
+    # Output file path
+    filename = f"{uuid.uuid4().hex}.wav"
+    output_path = os.path.join(outdir, filename)
 
-    sf.write(outpath, wave, sr)
+    # Write audio
+    try:
+        sf.write(output_path, waveform, sr)
+    except Exception as e:
+        print("AUDIO WRITE ERROR:", e)
+        return None
 
-    return outpath
+    return output_path