Update main.py

main.py

@@ -2,18 +2,15 @@ from fastapi import FastAPI, Request
 from fastapi.staticfiles import StaticFiles
 from fastapi.responses import FileResponse
 from pydantic import BaseModel
-from pydub import AudioSegment
 import torch
-import torch.nn as nn
-import pickle
-import numpy as np
-from music21 import note, chord, stream, instrument
+from transformers import AutoProcessor, MusicgenForConditionalGeneration
+import scipy.io.wavfile
 import os
-os.environ['KMP_DUPLICATE_LIB_OK'] = 'True'
-import random
-import subprocess
+import time
 
-app = FastAPI(title="Bio-Vibe AI Inference")
+app = FastAPI(title="Bio-Vibe MusicGen Inference")
+
+# Tạo thư mục chứa nhạc
 os.makedirs("music", exist_ok=True)
 app.mount("/music", StaticFiles(directory="music"), name="music")
 
@@ -25,232 +22,77 @@ class HeartRateData(BaseModel):
     bpm: int
 
 # ==========================================
-# 1. KHÔI PHỤC "BỘ NÃO" (PYTORCH MODEL)
+# 1. TẢI BỘ NÃO MUSICGEN (META)
 # ==========================================
-# Cấu trúc class phải y hệt lúc train
-class BioVibeAI(nn.Module):
-    def __init__(self, vocab_size, embed_size=256, hidden_size=512, num_layers=2):
-        super(BioVibeAI, self).__init__()
-        self.embedding = nn.Embedding(vocab_size, embed_size)
-        self.lstm = nn.LSTM(embed_size, hidden_size, num_layers, batch_first=True)
-        self.fc = nn.Linear(hidden_size, vocab_size)
-        
-    def forward(self, x):
-        embedded = self.embedding(x)
-        out, _ = self.lstm(embedded)
-        out = self.fc(out[:, -1, :])
-        return out
-
-# Nạp từ điển
-with open('mapping_dict.pkl', 'rb') as f:
-    int_to_note = pickle.load(f)
-note_to_int = {note: number for number, note in int_to_note.items()}
-vocab_size = len(int_to_note)
-
-# 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("biovibe_brain.pth", map_location=device)
+print("Đang tải bộ não MusicGen-Small từ Meta (Khoảng 1.5GB)...")
+# Server Hugging Face Free chạy CPU nên ta dùng CPU
+device = "cpu"
 
-# --- 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() # Chuyển sang chế độ suy luận
+processor = AutoProcessor.from_pretrained("facebook/musicgen-small")
+model = MusicgenForConditionalGeneration.from_pretrained("facebook/musicgen-small")
+model.to(device)
+print("✅ Tải xong! Hệ thống đã sẵn sàng.")
 
-# 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)
+# ==========================================
+# 2. HÀM DỌN RÁC (Tối ưu ổ cứng)
+# ==========================================
+def cleanup_old_audio(folder_path="music", max_age_seconds=600):
+    now = time.time()
+    for filename in os.listdir(folder_path):
+        file_path = os.path.join(folder_path, filename)
+        if filename.endswith('.wav'):
+            if os.path.getmtime(file_path) < now - max_age_seconds:
+                try:
+                    os.remove(file_path)
+                    print(f"🔥 Đã tự động dọn dẹp: {filename}")
+                except:
+                    pass
 
 # ==========================================
-# 2. HÀM SÁNG TÁC BẰNG NHIỆT ĐỘ (TEMPERATURE)
+# 3. KỊCH BẢN CHỮA LÀNH BẰNG NHỊP TIM
 # ==========================================
-def generate_ai_midi(bpm: int, output_file: str, num_notes: int = 60):
-    # --- PHẦN THÊM MỚI: ĐỊNH NGHĨA THANG NGŨ CUNG (C Pentatonic) ---
-    # Các nốt: C, D, E, G, A ở nhiều quãng tám (từ trầm đến trung)
-    # Chỉ dùng các nốt từ C2 (36) đến C4 (60) - Tuyệt đối không để vượt quá 60
-    # PENTATONIC_MS = [36, 38, 40, 43, 45, 48, 50, 52, 55, 57, 60]
-    
-    # def get_closest_pentatonic(pitch_input):
-    #     """Hàm ép một nốt bất kỳ về nốt Ngũ cung gần nhất"""
-    #     return min(PENTATONIC_MS, key=lambda x: abs(x - pitch_input))
-    # -----------------------------------------------------------
-    # Logic phản hồi sinh học
-    if bpm > 80:
-        temperature = 0.7  # An toàn, ổn định
-        base_offset = 1.0  # Nốt đánh chậm
+def get_prompt_from_bpm(bpm: int) -> str:
+    if bpm < 70:
+        return "Calm relaxing lo-fi piano music, slow tempo, ambient, soothing, healing therapy"
+    elif bpm <= 85:
+        return "Chillout electronic ambient music, smooth synth, moderate tempo, focus"
+    elif bpm <= 100:
+        return "Upbeat synthwave, energetic bassline, moderate fast tempo, driving rhythm"
     else:
-        temperature = 1.0  # Bay bổng, ngẫu hứng
-        base_offset = 0.5  # Nốt đánh nhanh hơn một chút
-        
-    # Chọn bừa 100 nốt từ kho dữ liệu để làm "mồi" cho AI
-    start_idx = random.randint(0, len(corpus) - 100)
-    seed_sequence = corpus[start_idx:start_idx + 100]
-    pattern = [note_to_int[char] for char in seed_sequence]
-    
-    generated_notes = []
-    
-    print(f"Đang sáng tác với Temp={temperature}...")
-    
-    # AI bắt đầu dự đoán
-    for i in range(num_notes):
-        sequence_tensor = torch.tensor([pattern], dtype=torch.long).to(device)
-        
-        with torch.no_grad():
-            prediction = model(sequence_tensor)
-            
-        # ÁP DỤNG TEMPERATURE
-        prediction = prediction / temperature
-        probabilities = torch.softmax(prediction, dim=1).numpy()[0]
-        
-        # Chọn nốt dựa trên phân phối xác suất
-        index = np.random.choice(len(probabilities), p=probabilities)
-        
-        result_note = int_to_note[index]
-        # --- PHẦN SỬA ĐỔI: ÉP NỐT VỀ PENTATONIC ---
-        # 1. Nếu là nốt đơn (ví dụ: 'C4', 'E5')
-        # if not ('.' in result_note or result_note.isdigit()):
-        #     # Trong hàm generate_ai_midi, phần xử lý nốt:
-        #     temp_note = note.Note(result_note)
-        #     pitch_val = temp_note.pitch.ps
-            
-        #     # Nếu nốt cao hơn nốt 60 (C4), hãy hạ nó xuống 1 hoặc 2 quãng tám
-        #     while pitch_val > 60:
-        #         pitch_val -= 12
-            
-        #     # Sau đó mới tìm nốt gần nhất trong danh sách Pentatonic trầm
-        #     new_pitch = get_closest_pentatonic(pitch_val)
-        #     temp_note.pitch.ps = new_pitch
-            
-        # # 2. Nếu là hợp âm hoặc nốt dạng số (ví dụ: '60.64.67')
-        # else:
-        #     parts = result_note.split('.')
-        #     new_parts = []
-        #     for p in parts:
-        #         p_int = int(p)
-        #         new_p = get_closest_pentatonic(p_int)
-        #         new_parts.append(str(int(new_p)))
-        #     result_note = '.'.join(new_parts)
-        # -----------------------------------------
-        generated_notes.append(result_note)
-        
-        # Cập nhật cửa sổ trượt (thêm nốt mới, xóa nốt cũ nhất)
-        pattern.append(index)
-        pattern = pattern[1:]
-        
-    # ==========================================
-    # 3. CHUYỂN NGỮ AI THÀNH FILE MIDI
-    # ==========================================
-    offset = 0
-    output_stream = stream.Stream()
-    output_stream.append(instrument.Piano()) # Chọn nhạc cụ
-    
-    for pattern_note in generated_notes:
-        # Nếu là hợp âm (có dấu chấm, ví dụ '4.7.11')
-        try:
-            # Xử lý Hợp âm
-            if ('.' in pattern_note) or pattern_note.isdigit():
-                notes_in_chord = pattern_note.split('.')
-                chord_notes = []
-                for current_note in notes_in_chord:
-                    n = note.Note(int(current_note))
-                    # --- CAN THIỆP TẠI ĐÂY ---
-                    # 1. Ép vùng âm trầm (Lower Pitch)
-                    if n.pitch.ps > 65: 
-                        n.pitch.ps -= 12 
-                    
-                    # 2. KỸ THUẬT 2: Humanize Velocity (Lực nhấn cực khẽ)
-                    n.volume.velocity = random.randint(30, 45)
-                    # -------------------------
-                    chord_notes.append(n)
-                new_obj = chord.Chord(chord_notes)
-
-            # Xử lý Nốt đơn
-            else:
-                new_obj = note.Note(pattern_note)
-                # --- CAN THIỆP TẠI ĐÂY ---
-                if new_obj.pitch.ps > 65:
-                    new_obj.pitch.ps -= 12
-                
-                # 2. KỸ THUẬT 2: Humanize Velocity
-                new_obj.volume.velocity = random.randint(30, 45)
-
-            # 3. KỸ THUẬT 2: Micro-timing (Lệch nhịp nhẹ để tạo sự uyển chuyển)
-            # Thay vì offset cứng nhắc, ta cộng thêm một lượng siêu nhỏ ngẫu nhiên
-            humanized_offset = offset + random.uniform(-0.03, 0.03)
-            new_obj.offset = humanized_offset
-            
-            # Kéo dài nốt một chút (Duration) để chúng gối đầu lên nhau (Legato)
-            new_obj.quarterLength = base_offset * 1.1
-            output_stream.append(new_obj)
-            offset += base_offset
-        except:
-            continue
-        
-    output_stream.write('midi', fp=output_file)
+        return "Fast tempo, tense cinematic action music, heavy drums, energetic orchestra"
 
 # ==========================================
 # 4. API ENDPOINT
 # ==========================================
-
-# Low-pass filter
-def apply_audio_filters(wav_path):
-    sound = AudioSegment.from_wav(wav_path)
-    
-    # Cắt toàn bộ tần số trên 600Hz (Đây là mức "Dark/Warm" - cực kỳ ấm áp)
-    # 600Hz sẽ loại bỏ hoàn toàn tiếng chói, chỉ để lại tiếng piano mờ ảo.
-    sound = sound.low_pass_filter(600) 
+@app.post("/generate-vibe")
+async def generate_vibe(data: HeartRateData, request: Request):
+    cleanup_old_audio()
     
-    # Giảm thêm âm lượng nốt để tránh nén tiếng
-    sound = sound - 7 
+    unique_id = int(time.time())
+    wav_file = f"music/vibe_{data.bpm}_{unique_id}.wav"
     
-    # Làm mượt đầu cuối
-    sound = sound.fade_in(2000).fade_out(3000)
+    # Dịch nhịp tim thành câu lệnh văn bản
+    music_prompt = get_prompt_from_bpm(data.bpm)
+    print(f"Nhịp tim {data.bpm} BPM -> Sáng tác theo chủ đề: '{music_prompt}'")
     
-    sound.export(wav_path, format="wav")
-@app.post("/generate-vibe")
-async def generate_vibe(data: HeartRateData, request: Request):
-    # Đường dẫn file
-    midi_file = f"music/ai_vibe_{data.bpm}.mid"
-    wav_file = f"music/ai_vibe_{data.bpm}.wav"
-    soundfont_path = "soundfonts/Piano.sf2" # Đường dẫn tới file sf2 bạn vừa tải
+    # Tiền xử lý văn bản
+    inputs = processor(
+        text=[music_prompt],
+        padding=True,
+        return_tensors="pt",
+    ).to(device)
     
-    # Bước 1: PyTorch sinh ra cấu trúc nốt nhạc (MIDI)
-    # Bước 1: PyTorch sinh ra cấu trúc nốt nhạc (MIDI)
-    generate_ai_midi(data.bpm, midi_file)
+    # Cho AI sinh âm thanh
+    # max_new_tokens=256 tương đương khoảng 5 giây âm thanh (phù hợp để demo trên CPU)
+    # Tăng số này lên nếu muốn nhạc dài hơn (nhưng sẽ gen lâu hơn)
+    print("Đang ép xung CPU để render âm thanh... Vui lòng đợi!")
+    audio_values = model.generate(**inputs, do_sample=True, guidance_scale=3, max_new_tokens=256)
     
-    # Bước 2: Tự viết lệnh Render bằng Subprocess chuẩn xác cho FluidSynth 2.5+
-    print("Đang render MIDI thành âm thanh Piano...")
-    try:
-        subprocess.run([
-            "fluidsynth", 
-            "-ni",               # Chạy ẩn không cần giao diện
-            "-g", "0.6",         # [NÂNG CẤP] Giảm âm lượng tổng để tiếng bớt "đanh"
-            "-r", "44100",       # Chất lượng Sample rate
-            "--reverb", "yes",   # [MỚI] Bật vang số để các nốt nối đuôi nhau (giảm rời rạc)
-            "--chorus", "yes",   # [MỚI] Làm dày âm thanh, tạo cảm giác êm ái hơn
-            "-F", wav_file,      # Xuất ra file thay vì phát ra loa laptop!
-            soundfont_path,      # Dàn nhạc (Soundfont)
-            midi_file            # Kịch bản (MIDI)
-        ], check=True)
-
-        # --- BƯỚC MỚI: LỌC ÂM THANH NHỨC ÓC ---
-        apply_audio_filters(wav_file)
-        # -------------------------------------
-    except Exception as e:
-        print(f"Lỗi khi render âm thanh: {e}")
-        return {"status": "error", "message": "Render failed"}
+    # Xuất ra file .wav chuẩn
+    sampling_rate = model.config.audio_encoder.sampling_rate
+    audio_data = audio_values[0, 0].cpu().numpy()
+    scipy.io.wavfile.write(wav_file, rate=sampling_rate, data=audio_data)
+    print("✅ Render hoàn tất!")
     
     server_url = str(request.base_url).rstrip("/")
-    
-    # Bước 3: Trả về link file âm thanh xịn xò cho App Flutter
     return {"status": "success", "audio_url": f"{server_url}/{wav_file}"}