Update main.py

main.py

@@ -1,12 +1,15 @@
-from fastapi import FastAPI, WebSocket
+from fastapi import FastAPI, Request, WebSocket, WebSocketDisconnect
 from fastapi.staticfiles import StaticFiles
 from fastapi.responses import FileResponse
+from pydantic import BaseModel
 import torch
 import torch.nn as nn
 import pickle
 import numpy as np
+from music21 import note
 import random
 import asyncio
+import math
 import os
 
 app = FastAPI(title="Bio-Vibe AI Inference")
@@ -17,132 +20,76 @@ app.mount("/music", StaticFiles(directory="music"), name="music")
 async def serve_frontend():
     return FileResponse("index.html")
 
+class HeartRateData(BaseModel):
+    bpm: int
+
 # ==========================================
-# MODEL
+# 1. KHÔI PHỤC "BỘ NÃO" (PYTORCH MODEL)
 # ==========================================
+# Cấu trúc class phải y hệt lúc train
 class BioVibeAI(nn.Module):
     def __init__(self, vocab_size):
         super().__init__()
+
         self.embedding = nn.Embedding(vocab_size, 256)
-        self.lstm = nn.LSTM(256, 512, num_layers=3, batch_first=True, dropout=0.3)
+
+        self.lstm = nn.LSTM(
+            256,
+            512,
+            num_layers=3,
+            batch_first=True,
+            dropout=0.3
+        )
+
         self.norm = nn.LayerNorm(512)
         self.fc = nn.Linear(512, vocab_size)
 
-    def forward(self, x, hidden=None):
+    def forward(self, x):
         x = self.embedding(x)
-        out, hidden = self.lstm(x, hidden)
+        out, _ = self.lstm(x)
         out = out[:, -1, :]
         out = self.norm(out)
-        return self.fc(out), hidden
+        return self.fc(out)
 
+# Nạp từ điển
 with open('mapping_dict.pkl', 'rb') as f:
     int_to_note = pickle.load(f)
-note_to_int = {n: i for i, n in int_to_note.items()}
+note_to_int = {note: number for number, note in int_to_note.items()}
 vocab_size = len(int_to_note)
 
-device = torch.device("cpu")
+# Load trọng số mô hình
+device = torch.device("cpu") # Server chạy CPU cho nhẹ
 model = BioVibeAI(vocab_size=vocab_size)
+# Load trọng số mô hình từ file
 state_dict = torch.load("model.pth", map_location=device)
-clean_state_dict = {
-    k[7:] if k.startswith('module.') else k: v
-    for k, v in state_dict.items()
-}
+
+# --- THÊM ĐOẠN NÀY: Lột bỏ vỏ bọc 'module.' của DataParallel ---
+clean_state_dict = {}
+for key, value in state_dict.items():
+    if key.startswith('module.'):
+        # Cắt bỏ 7 ký tự đầu tiên ('m', 'o', 'd', 'u', 'l', 'e', '.')
+        clean_key = key[7:] 
+        clean_state_dict[clean_key] = value
+    else:
+        clean_state_dict[key] = value
+
+# Nạp bộ não đã được gọt dũa sạch sẽ vào model
 model.load_state_dict(clean_state_dict)
-model.eval()
+model.eval() # Chuyển sang chế độ suy luận
 
+# Nạp kho dữ liệu cũ để lấy Seed (đoạn nhạc mồi)
 with open('notes_corpus.pkl', 'rb') as f:
     corpus = pickle.load(f)
-
-# ==========================================
-# CONSTANTS
-# ==========================================
-SEQ_LEN = 128
-BARS_PER_BATCH = 2
-
-# Chord progressions — 2 màu sắc luân phiên để tránh lặp
-CHORD_SETS = {
-    "calm": [            # BPM cao — tông mềm, dreamy
-        [48, 52, 55, 59],    # Cmaj7
-        [45, 48, 52, 55],    # Am7
-        [41, 45, 48, 52],    # Fmaj7
-        [43, 47, 50, 55],    # G9
-    ],
-    "groove": [          # BPM thấp — tông rõ hơn, có chiều sâu
-        [48, 52, 55, 59],    # Cmaj7
-        [43, 47, 50, 53],    # G7
-        [45, 48, 52, 57],    # Am9
-        [41, 44, 48, 52],    # Fm7
-    ],
-}
-
-# ==========================================
-# ENTRAINMENT HELPERS
-# ==========================================
-
-def get_music_beat(bpm: int) -> float:
-    """
-    Tim đập nhanh -> nhac cham lai de keo nhip tim xuong (entrainment).
-    BPM [40->180] => music_tempo [85->55] BPM (nghich chieu).
-    """
-    music_tempo = 85.0 - (bpm - 40) * (30.0 / 140.0)
-    music_tempo = max(52.0, min(88.0, music_tempo))
-    return 60.0 / music_tempo
-
-
-def get_temperature(bpm: int) -> float:
-    """
-    BPM cao -> temp THAP -> melody on dinh, du doan duoc, calming.
-    BPM thap -> temp CAO -> melody da dang, groovy hon.
-    Range: [0.60, 0.92]
-    """
-    t = (bpm - 40) / (180 - 40)
-    return round(0.92 - t * 0.32, 3)
-
-
-def get_melody_params(bpm: int, music_beat: float) -> dict:
-    """
-    BPM cao -> not nhe hon, dai hon (legato), range cao (airy).
-    BPM thap -> not vua, ngan hon (rhythmic), range mid.
-    """
-    t = (bpm - 40) / (180 - 40)
-
-    vel_low  = int(75 - t * 28)    # 75 (bpm=40) -> 47 (bpm=180)
-    vel_high = int(90 - t * 22)    # 90 -> 68
-
-    dur_factor = 0.72 + t * 0.48   # 0.72 -> 1.20 (legato khi BPM cao)
-    duration   = round(music_beat * dur_factor, 4)
-
-    midi_low  = int(59 + t * 5)    # 59 -> 64
-    midi_high = int(71 + t * 5)    # 71 -> 76
-
-    return {
-        "vel_low": max(40, vel_low),
-        "vel_high": max(60, vel_high),
-        "duration": duration,
-        "midi_low": midi_low,
-        "midi_high": midi_high,
-    }
-
-
-def get_chord_set(bpm: int, bar_count: int) -> list:
-    """
-    BPM cao -> chord calm (Cmaj7/Am7/Fmaj7/G9).
-    BPM thap -> chord groove (co Fm7 them mau).
-    Doi chord set moi 16 bar.
-    """
-    use_groove = (bpm < 85) or (bar_count % 16 >= 8 and bpm < 110)
-    return CHORD_SETS["groove"] if use_groove else CHORD_SETS["calm"]
-
-
-# ==========================================
-# WEBSOCKET ENDPOINT
+    
 # ==========================================
+# 4. API ENDPOINT
+# ==========================================    
 @app.websocket("/ws/vibe")
 async def websocket_vibe(websocket: WebSocket):
     await websocket.accept()
-
+    
     current_bpm = 75
-    is_playing  = True
+    is_playing = True
 
     async def listen():
         nonlocal current_bpm, is_playing
@@ -152,129 +99,107 @@ async def websocket_vibe(websocket: WebSocket):
                 bpm = int(data.get("bpm", current_bpm))
                 if 40 <= bpm <= 180:
                     current_bpm = bpm
-        except Exception:
+        except:
             is_playing = False
 
     asyncio.create_task(listen())
 
-    # Chon seed uu tien LOFI
-    lofi_songs = [s for s in corpus if s and len(s) >= SEQ_LEN and s[0] == "STYLE_LOFI"]
-    pool = lofi_songs if lofi_songs else [s for s in corpus if len(s) >= SEQ_LEN]
-    song = random.choice(pool)
+    SEQ_LEN = 50
+    # 👉 chọn 1 bài nhạc random
+    song = random.choice(corpus)
+
+    # 👉 đảm bảo đủ độ dài
+    if len(song) < SEQ_LEN:
+        song = random.choice(corpus)
+
+    # 👉 lấy đoạn đầu
+    seed_tokens = song[:SEQ_LEN]
+
+    # 👉 ép style (QUAN TRỌNG)
+    seed_tokens[0] = "STYLE_LOFI"
 
-    seed = list(song[:SEQ_LEN])
-    seed[0] = "STYLE_LOFI"
-    pattern = [note_to_int.get(tok, 0) for tok in seed]
+    # 👉 convert sang int
+    pattern = [note_to_int[n] for n in seed_tokens]
 
-    # Warm-up LSTM mot lan duy nhat voi toan bo seed
-    with torch.no_grad():
-        seed_tensor = torch.tensor([pattern], dtype=torch.long)
-        _, hidden = model(seed_tensor)
+    inp = torch.zeros(1, SEQ_LEN, dtype=torch.long)
 
-    bar_count = 0
+    note_count = 0
+
+    CHORDS = [
+        [48,52,55,59],
+        [45,48,52,55],
+        [41,45,48,52],
+        [43,47,50,53]
+    ]
+
+    def snap(midi, chord):
+        return min(chord + [c+12 for c in chord], key=lambda x: abs(x - midi))
 
     try:
+        STEPS_PER_BATCH = 16
+
         while is_playing:
-            bpm = current_bpm
-
-            music_beat = get_music_beat(bpm)
-            bar_dur    = music_beat * 4
-            batch_dur  = bar_dur * BARS_PER_BATCH
-            temp       = get_temperature(bpm)
-            mel_params = get_melody_params(bpm, music_beat)
-            chords     = get_chord_set(bpm, bar_count)
-
-            # ── SINH MELODY ──────────────────────────────────────────────
-            melody_notes = []
-            current_time = 0.0
-            last_midi    = None
-            MAX_TOKENS   = 180
-
-            for _ in range(MAX_TOKENS):
-                if current_time >= batch_dur:
-                    break
-
-                new_tok = torch.tensor([[pattern[-1]]], dtype=torch.long)
+            beat = 60.0 / current_bpm
+
+            step = beat / 2
+
+            batch = []
+
+            for i in range(STEPS_PER_BATCH):
+                inp[0] = torch.tensor(pattern)
+
                 with torch.no_grad():
-                    logits, hidden = model(new_tok, hidden)
-                    probs = torch.softmax(logits / temp, dim=1).numpy()[0]
-                    idx   = int(np.random.choice(len(probs), p=probs))
+                    pred = model(inp)
+                    probs = torch.softmax(pred / 0.8, dim=1).numpy()[0]
+                    idx = np.random.choice(len(probs), p=probs)
 
-                tok     = int_to_note[idx]
+                tok = int_to_note[idx]
                 pattern = pattern[1:] + [idx]
 
-                if tok.startswith("TIME_"):
-                    try:
-                        gap_beats    = float(tok.split("_")[1])
-                        current_time += gap_beats * music_beat
-                    except ValueError:
-                        pass
-
-                elif tok.startswith("NOTE_"):
-                    try:
-                        midi = int(tok.split("_")[1])
-                        lo, hi = mel_params["midi_low"], mel_params["midi_high"]
-
-                        while midi < lo: midi += 12
-                        while midi > hi: midi -= 12
-
-                        if midi == last_midi and random.random() < 0.40:
-                            step = random.choice([-2, 2, 3, -3, 5])
-                            midi = max(lo - 2, min(hi + 2, midi + step))
-
-                        last_midi = midi
-                        velocity  = random.randint(mel_params["vel_low"],
-                                                   mel_params["vel_high"])
-                        melody_notes.append({
-                            "note":     midi,
-                            "velocity": velocity,
-                            "duration": mel_params["duration"],
-                            "type":     "melody",
-                            "time":     round(current_time, 4),
+                midi = None
+                try:
+                    if tok.startswith("NOTE"):
+                        midi = note.Note(tok.split("_")[1]).pitch.midi
+                except:
+                    pass
+
+                chord = CHORDS[((note_count + i) // 8) % len(CHORDS)]
+
+                notes = []
+
+                # chord 
+                if (note_count + i) % 8 == 0:
+                    for c in chord:
+                        notes.append({
+                            "note": c,
+                            "velocity": 40,
+                            "duration": beat * 4,
+                            "type": "chord",
+                            "time": i * step
                         })
-                    except ValueError:
-                        pass
-
-            # Fallback pentatonic
-            if len(melody_notes) < 3:
-                scale = [60, 62, 64, 67, 69, 72]
-                t = 0.0
-                while t < batch_dur:
-                    melody_notes.append({
-                        "note":     random.choice(scale),
-                        "velocity": random.randint(mel_params["vel_low"],
-                                                   mel_params["vel_high"]),
-                        "duration": mel_params["duration"],
-                        "type":     "melody",
-                        "time":     round(t, 4),
-                    })
-                    t += music_beat * random.choice([0.5, 0.5, 1.0])
-
-            # ── CHORD PADS ───────────────────────────────────────────────
-            # BPM cao -> pad nhe, am thanh (ambient)
-            chord_vel = max(32, int(50 - (bpm - 40) * 0.06))
-
-            chord_notes = []
-            for bar_i in range(BARS_PER_BATCH):
-                chord = chords[(bar_count + bar_i) % len(chords)]
-                t = bar_i * bar_dur
-                for c in chord:
-                    chord_notes.append({
-                        "note":     c,
-                        "velocity": chord_vel,
-                        "duration": round(bar_dur * 0.95, 4),
-                        "type":     "chord",
-                        "time":     round(t, 4),
+
+                # melody
+                if midi:
+                    while midi < 55: midi += 12
+                    while midi > 75: midi -= 12
+
+                    notes.append({
+                        "note": midi,
+                        "velocity": random.randint(55, 80),
+                        "duration": step * 1.2,
+                        "type": "melody",
+                        "time": i * step
                     })
 
+                batch.extend(notes)
+
             await websocket.send_json({
-                "notes":          melody_notes + chord_notes,
-                "batch_duration": round(batch_dur, 4),
-                "music_bpm":      round(60.0 / music_beat, 1),
+                "notes": batch,
+                "step": step
             })
 
-            bar_count += BARS_PER_BATCH
-            await asyncio.sleep(batch_dur * 0.80)
+            note_count += STEPS_PER_BATCH
+            await asyncio.sleep(step * STEPS_PER_BATCH * 0.8)
 
-    except Exception:
-        pass
+    except:
+        pass