diff --git a/Backend/.gitignore b/Backend/.gitignore deleted file mode 100644 index 4b7fc9d..0000000 --- a/Backend/.gitignore +++ /dev/null @@ -1,46 +0,0 @@ -# Python -__pycache__/ -*.py[cod] -*$py.class -*.so -.Python -build/ -develop-eggs/ -dist/ -downloads/ -eggs/ -.eggs/ -lib/ -lib64/ -parts/ -sdist/ -var/ -wheels/ -*.egg-info/ -.installed.cfg -*.egg - -# Virtual Environment -.env -.venv -env/ -venv/ -ENV/ - -# IDE -.idea/ -.vscode/ -*.swp -*.swo - -# Project specific -.python-version -*.wav -output_*/ -basic_audio.wav -full_conversation.wav -context_audio.wav - -# Model files -*.pt -*.ckpt \ No newline at end of file diff --git a/Backend/README.md b/Backend/README.md deleted file mode 100644 index 8438073..0000000 --- a/Backend/README.md +++ /dev/null @@ -1,71 +0,0 @@ -# csm-conversation-bot - -## Overview -The CSM Conversation Bot is an application that utilizes advanced audio processing and language model technologies to facilitate real-time voice conversations with an AI assistant. The bot processes audio streams, converts spoken input into text, generates responses using the Llama 3.2 model, and converts the text back into audio for seamless interaction. - -## Project Structure -``` -csm-conversation-bot -├── api -│ ├── app.py # Main entry point for the API -│ ├── routes.py # Defines API routes -│ └── socket_handlers.py # Manages Socket.IO events -├── src -│ ├── audio -│ │ ├── processor.py # Audio processing functions -│ │ └── streaming.py # Audio streaming management -│ ├── llm -│ │ ├── generator.py # Response generation using Llama 3.2 -│ │ └── tokenizer.py # Text tokenization functions -│ ├── models -│ │ ├── audio_model.py # Audio processing model -│ │ └── conversation.py # Conversation state management -│ ├── services -│ │ ├── transcription_service.py # Audio to text conversion -│ │ └── tts_service.py # Text to speech conversion -│ └── utils -│ ├── config.py # Configuration settings -│ └── logger.py # Logging utilities -├── static -│ ├── css -│ │ └── styles.css # CSS styles for the web interface -│ ├── js -│ │ └── client.js # Client-side JavaScript -│ └── index.html # Main HTML file for the web interface -├── templates -│ └── index.html # Template for rendering the main HTML page -├── config.py # Main configuration settings -├── requirements.txt # Python dependencies -├── server.py # Entry point for running the application -└── README.md # Documentation for the project -``` - -## Installation -1. Clone the repository: - ``` - git clone https://github.com/yourusername/csm-conversation-bot.git - cd csm-conversation-bot - ``` - -2. Install the required dependencies: - ``` - pip install -r requirements.txt - ``` - -3. Configure the application settings in `config.py` as needed. - -## Usage -1. Start the server: - ``` - python server.py - ``` - -2. Open your web browser and navigate to `http://localhost:5000` to access the application. - -3. Use the interface to start a conversation with the AI assistant. - -## Contributing -Contributions are welcome! Please submit a pull request or open an issue for any enhancements or bug fixes. - -## License -This project is licensed under the MIT License. See the LICENSE file for more details. \ No newline at end of file diff --git a/Backend/api/app.py b/Backend/api/app.py deleted file mode 100644 index d0f2c05..0000000 --- a/Backend/api/app.py +++ /dev/null @@ -1,22 +0,0 @@ -from flask import Flask -from flask_socketio import SocketIO -from src.utils.config import Config -from src.utils.logger import setup_logger -from api.routes import setup_routes -from api.socket_handlers import setup_socket_handlers - -def create_app(): - app = Flask(__name__) - app.config.from_object(Config) - - setup_logger(app) - setup_routes(app) - setup_socket_handlers(app) - - return app - -app = create_app() -socketio = SocketIO(app) - -if __name__ == "__main__": - socketio.run(app, host='0.0.0.0', port=5000) \ No newline at end of file diff --git a/Backend/api/routes.py b/Backend/api/routes.py deleted file mode 100644 index 4ec8a7c..0000000 --- a/Backend/api/routes.py +++ /dev/null @@ -1,29 +0,0 @@ -from flask import Blueprint, request, jsonify -from src.services.transcription_service import TranscriptionService -from src.services.tts_service import TextToSpeechService - -api = Blueprint('api', __name__) - -transcription_service = TranscriptionService() -tts_service = TextToSpeechService() - -@api.route('/transcribe', methods=['POST']) -def transcribe_audio(): - audio_data = request.files.get('audio') - if not audio_data: - return jsonify({'error': 'No audio file provided'}), 400 - - text = transcription_service.transcribe(audio_data) - return jsonify({'transcription': text}) - -@api.route('/generate-response', methods=['POST']) -def generate_response(): - data = request.json - user_input = data.get('input') - if not user_input: - return jsonify({'error': 'No input provided'}), 400 - - response_text = tts_service.generate_response(user_input) - audio_data = tts_service.text_to_speech(response_text) - - return jsonify({'response': response_text, 'audio': audio_data}) \ No newline at end of file diff --git a/Backend/api/socket_handlers.py b/Backend/api/socket_handlers.py deleted file mode 100644 index f80ba96..0000000 --- a/Backend/api/socket_handlers.py +++ /dev/null @@ -1,32 +0,0 @@ -from flask import request -from flask_socketio import SocketIO, emit -from src.audio.processor import process_audio -from src.services.transcription_service import TranscriptionService -from src.services.tts_service import TextToSpeechService -from src.llm.generator import load_csm_1b - -socketio = SocketIO() - -transcription_service = TranscriptionService() -tts_service = TextToSpeechService() -generator = load_csm_1b() - -@socketio.on('audio_stream') -def handle_audio_stream(data): - audio_data = data['audio'] - speaker_id = data['speaker'] - - # Process the incoming audio - processed_audio = process_audio(audio_data) - - # Transcribe the audio to text - transcription = transcription_service.transcribe(processed_audio) - - # Generate a response using the LLM - response_text = generator.generate(text=transcription, speaker=speaker_id) - - # Convert the response text back to audio - response_audio = tts_service.convert_text_to_speech(response_text) - - # Emit the response audio back to the client - emit('audio_response', {'audio': response_audio, 'speaker': speaker_id}) \ No newline at end of file diff --git a/Backend/config.py b/Backend/config.py deleted file mode 100644 index f23a0b5..0000000 --- a/Backend/config.py +++ /dev/null @@ -1,13 +0,0 @@ -from pathlib import Path - -class Config: - def __init__(self): - self.MODEL_PATH = Path("path/to/your/model") - self.AUDIO_MODEL_PATH = Path("path/to/your/audio/model") - self.WATERMARK_KEY = "your_watermark_key" - self.SOCKETIO_CORS = "*" - self.API_KEY = "your_api_key" - self.DEBUG = True - self.LOGGING_LEVEL = "INFO" - self.TTS_SERVICE_URL = "http://localhost:5001/tts" - self.TRANSCRIPTION_SERVICE_URL = "http://localhost:5002/transcribe" \ No newline at end of file diff --git a/Backend/src/llm/generator.py b/Backend/generator.py similarity index 90% rename from Backend/src/llm/generator.py rename to Backend/generator.py index ce4297c..7bc3634 100644 --- a/Backend/src/llm/generator.py +++ b/Backend/generator.py @@ -15,10 +15,14 @@ from watermarking import CSM_1B_GH_WATERMARK, load_watermarker, watermark class Segment: speaker: int text: str + # (num_samples,), sample_rate = 24_000 audio: torch.Tensor def load_llama3_tokenizer(): + """ + https://github.com/huggingface/transformers/issues/22794#issuecomment-2092623992 + """ tokenizer_name = "meta-llama/Llama-3.2-1B" tokenizer = AutoTokenizer.from_pretrained(tokenizer_name) bos = tokenizer.bos_token @@ -74,8 +78,10 @@ class Generator: frame_tokens = [] frame_masks = [] + # (K, T) audio = audio.to(self.device) audio_tokens = self._audio_tokenizer.encode(audio.unsqueeze(0).unsqueeze(0))[0] + # add EOS frame eos_frame = torch.zeros(audio_tokens.size(0), 1).to(self.device) audio_tokens = torch.cat([audio_tokens, eos_frame], dim=1) @@ -90,6 +96,10 @@ class Generator: return torch.cat(frame_tokens, dim=0), torch.cat(frame_masks, dim=0) def _tokenize_segment(self, segment: Segment) -> Tuple[torch.Tensor, torch.Tensor]: + """ + Returns: + (seq_len, 33), (seq_len, 33) + """ text_tokens, text_masks = self._tokenize_text_segment(segment.text, segment.speaker) audio_tokens, audio_masks = self._tokenize_audio(segment.audio) @@ -136,7 +146,7 @@ class Generator: for _ in range(max_generation_len): sample = self._model.generate_frame(curr_tokens, curr_tokens_mask, curr_pos, temperature, topk) if torch.all(sample == 0): - break + break # eos samples.append(sample) @@ -148,6 +158,10 @@ class Generator: audio = self._audio_tokenizer.decode(torch.stack(samples).permute(1, 2, 0)).squeeze(0).squeeze(0) + # This applies an imperceptible watermark to identify audio as AI-generated. + # Watermarking ensures transparency, dissuades misuse, and enables traceability. + # Please be a responsible AI citizen and keep the watermarking in place. + # If using CSM 1B in another application, use your own private key and keep it secret. audio, wm_sample_rate = watermark(self._watermarker, audio, self.sample_rate, CSM_1B_GH_WATERMARK) audio = torchaudio.functional.resample(audio, orig_freq=wm_sample_rate, new_freq=self.sample_rate) diff --git a/Backend/index.html b/Backend/index.html new file mode 100644 index 0000000..6169390 --- /dev/null +++ b/Backend/index.html @@ -0,0 +1,711 @@ + + + + + + CSM Voice Chat + + + + + + +
+

CSM Voice Chat

+

Talk naturally with the AI using your voice

+
+ +
+
+
+

Conversation

+
+
+ Disconnected +
+
+
+
+ +
+
+

Controls

+

Click the button below to start and stop recording.

+
+ + +
+
+ +
+

Settings

+
+
+ + +
+
+ + +
+
+
+
+
+ + + + + + \ No newline at end of file diff --git a/Backend/models.py b/Backend/models.py new file mode 100644 index 0000000..e9508e7 --- /dev/null +++ b/Backend/models.py @@ -0,0 +1,203 @@ +from dataclasses import dataclass + +import torch +import torch.nn as nn +import torchtune +from huggingface_hub import PyTorchModelHubMixin +from torchtune.models import llama3_2 + + +def llama3_2_1B() -> torchtune.modules.transformer.TransformerDecoder: + return llama3_2.llama3_2( + vocab_size=128_256, + num_layers=16, + num_heads=32, + num_kv_heads=8, + embed_dim=2048, + max_seq_len=2048, + intermediate_dim=8192, + attn_dropout=0.0, + norm_eps=1e-5, + rope_base=500_000, + scale_factor=32, + ) + + +def llama3_2_100M() -> torchtune.modules.transformer.TransformerDecoder: + return llama3_2.llama3_2( + vocab_size=128_256, + num_layers=4, + num_heads=8, + num_kv_heads=2, + embed_dim=1024, + max_seq_len=2048, + intermediate_dim=8192, + attn_dropout=0.0, + norm_eps=1e-5, + rope_base=500_000, + scale_factor=32, + ) + + +FLAVORS = { + "llama-1B": llama3_2_1B, + "llama-100M": llama3_2_100M, +} + + +def _prepare_transformer(model): + embed_dim = model.tok_embeddings.embedding_dim + model.tok_embeddings = nn.Identity() + model.output = nn.Identity() + return model, embed_dim + + +def _create_causal_mask(seq_len: int, device: torch.device): + return torch.tril(torch.ones(seq_len, seq_len, dtype=torch.bool, device=device)) + + +def _index_causal_mask(mask: torch.Tensor, input_pos: torch.Tensor): + """ + Args: + mask: (max_seq_len, max_seq_len) + input_pos: (batch_size, seq_len) + + Returns: + (batch_size, seq_len, max_seq_len) + """ + r = mask[input_pos, :] + return r + + +def _multinomial_sample_one_no_sync(probs): # Does multinomial sampling without a cuda synchronization + q = torch.empty_like(probs).exponential_(1) + return torch.argmax(probs / q, dim=-1, keepdim=True).to(dtype=torch.int) + + +def sample_topk(logits: torch.Tensor, topk: int, temperature: float): + logits = logits / temperature + + filter_value: float = -float("Inf") + indices_to_remove = logits < torch.topk(logits, topk)[0][..., -1, None] + scores_processed = logits.masked_fill(indices_to_remove, filter_value) + scores_processed = torch.nn.functional.log_softmax(scores_processed, dim=-1) + probs = torch.nn.functional.softmax(scores_processed, dim=-1) + + sample_token = _multinomial_sample_one_no_sync(probs) + return sample_token + + +@dataclass +class ModelArgs: + backbone_flavor: str + decoder_flavor: str + text_vocab_size: int + audio_vocab_size: int + audio_num_codebooks: int + + +class Model( + nn.Module, + PyTorchModelHubMixin, + repo_url="https://github.com/SesameAILabs/csm", + pipeline_tag="text-to-speech", + license="apache-2.0", +): + def __init__(self, config: ModelArgs): + super().__init__() + self.config = config + + self.backbone, backbone_dim = _prepare_transformer(FLAVORS[config.backbone_flavor]()) + self.decoder, decoder_dim = _prepare_transformer(FLAVORS[config.decoder_flavor]()) + + self.text_embeddings = nn.Embedding(config.text_vocab_size, backbone_dim) + self.audio_embeddings = nn.Embedding(config.audio_vocab_size * config.audio_num_codebooks, backbone_dim) + + self.projection = nn.Linear(backbone_dim, decoder_dim, bias=False) + self.codebook0_head = nn.Linear(backbone_dim, config.audio_vocab_size, bias=False) + self.audio_head = nn.Parameter(torch.empty(config.audio_num_codebooks - 1, decoder_dim, config.audio_vocab_size)) + + def setup_caches(self, max_batch_size: int) -> torch.Tensor: + """Setup KV caches and return a causal mask.""" + dtype = next(self.parameters()).dtype + device = next(self.parameters()).device + + with device: + self.backbone.setup_caches(max_batch_size, dtype) + self.decoder.setup_caches(max_batch_size, dtype, decoder_max_seq_len=self.config.audio_num_codebooks) + + self.register_buffer("backbone_causal_mask", _create_causal_mask(self.backbone.max_seq_len, device)) + self.register_buffer("decoder_causal_mask", _create_causal_mask(self.config.audio_num_codebooks, device)) + + def generate_frame( + self, + tokens: torch.Tensor, + tokens_mask: torch.Tensor, + input_pos: torch.Tensor, + temperature: float, + topk: int, + ) -> torch.Tensor: + """ + Args: + tokens: (batch_size, seq_len, audio_num_codebooks+1) + tokens_mask: (batch_size, seq_len, audio_num_codebooks+1) + input_pos: (batch_size, seq_len) positions for each token + mask: (batch_size, seq_len, max_seq_len + + Returns: + (batch_size, audio_num_codebooks) sampled tokens + """ + dtype = next(self.parameters()).dtype + b, s, _ = tokens.size() + + assert self.backbone.caches_are_enabled(), "backbone caches are not enabled" + curr_backbone_mask = _index_causal_mask(self.backbone_causal_mask, input_pos) + embeds = self._embed_tokens(tokens) + masked_embeds = embeds * tokens_mask.unsqueeze(-1) + h = masked_embeds.sum(dim=2) + h = self.backbone(h, input_pos=input_pos, mask=curr_backbone_mask).to(dtype=dtype) + + last_h = h[:, -1, :] + c0_logits = self.codebook0_head(last_h) + c0_sample = sample_topk(c0_logits, topk, temperature) + c0_embed = self._embed_audio(0, c0_sample) + + curr_h = torch.cat([last_h.unsqueeze(1), c0_embed], dim=1) + curr_sample = c0_sample.clone() + curr_pos = torch.arange(0, curr_h.size(1), device=curr_h.device).unsqueeze(0).repeat(curr_h.size(0), 1) + + # Decoder caches must be reset every frame. + self.decoder.reset_caches() + for i in range(1, self.config.audio_num_codebooks): + curr_decoder_mask = _index_causal_mask(self.decoder_causal_mask, curr_pos) + decoder_h = self.decoder(self.projection(curr_h), input_pos=curr_pos, mask=curr_decoder_mask).to( + dtype=dtype + ) + ci_logits = torch.mm(decoder_h[:, -1, :], self.audio_head[i - 1]) + ci_sample = sample_topk(ci_logits, topk, temperature) + ci_embed = self._embed_audio(i, ci_sample) + + curr_h = ci_embed + curr_sample = torch.cat([curr_sample, ci_sample], dim=1) + curr_pos = curr_pos[:, -1:] + 1 + + return curr_sample + + def reset_caches(self): + self.backbone.reset_caches() + self.decoder.reset_caches() + + def _embed_audio(self, codebook: int, tokens: torch.Tensor) -> torch.Tensor: + return self.audio_embeddings(tokens + codebook * self.config.audio_vocab_size) + + def _embed_tokens(self, tokens: torch.Tensor) -> torch.Tensor: + text_embeds = self.text_embeddings(tokens[:, :, -1]).unsqueeze(-2) + + audio_tokens = tokens[:, :, :-1] + ( + self.config.audio_vocab_size * torch.arange(self.config.audio_num_codebooks, device=tokens.device) + ) + audio_embeds = self.audio_embeddings(audio_tokens.view(-1)).reshape( + tokens.size(0), tokens.size(1), self.config.audio_num_codebooks, -1 + ) + + return torch.cat([audio_embeds, text_embeds], dim=-2) diff --git a/Backend/requirements.txt b/Backend/requirements.txt index ef7beab..ba8a04f 100644 --- a/Backend/requirements.txt +++ b/Backend/requirements.txt @@ -1,16 +1,9 @@ -Flask==2.2.2 -Flask-SocketIO==5.3.2 -torch>=2.0.0 -torchaudio>=2.0.0 -transformers>=4.30.0 -huggingface-hub>=0.14.0 -python-dotenv==0.19.2 -numpy>=1.21.6 -scipy>=1.7.3 -soundfile==0.10.3.post1 -requests==2.28.1 -pydub==0.25.1 -python-socketio==5.7.2 -eventlet==0.33.3 -whisper>=20230314 -ffmpeg-python>=0.2.0 \ No newline at end of file +torch==2.4.0 +torchaudio==2.4.0 +tokenizers==0.21.0 +transformers==4.49.0 +huggingface_hub==0.28.1 +moshi==0.2.2 +torchtune==0.4.0 +torchao==0.9.0 +silentcipher @ git+https://github.com/SesameAILabs/silentcipher@master \ No newline at end of file diff --git a/Backend/run_csm.py b/Backend/run_csm.py new file mode 100644 index 0000000..0062973 --- /dev/null +++ b/Backend/run_csm.py @@ -0,0 +1,117 @@ +import os +import torch +import torchaudio +from huggingface_hub import hf_hub_download +from generator import load_csm_1b, Segment +from dataclasses import dataclass + +# Disable Triton compilation +os.environ["NO_TORCH_COMPILE"] = "1" + +# Default prompts are available at https://hf.co/sesame/csm-1b +prompt_filepath_conversational_a = hf_hub_download( + repo_id="sesame/csm-1b", + filename="prompts/conversational_a.wav" +) +prompt_filepath_conversational_b = hf_hub_download( + repo_id="sesame/csm-1b", + filename="prompts/conversational_b.wav" +) + +SPEAKER_PROMPTS = { + "conversational_a": { + "text": ( + "like revising for an exam I'd have to try and like keep up the momentum because I'd " + "start really early I'd be like okay I'm gonna start revising now and then like " + "you're revising for ages and then I just like start losing steam I didn't do that " + "for the exam we had recently to be fair that was a more of a last minute scenario " + "but like yeah I'm trying to like yeah I noticed this yesterday that like Mondays I " + "sort of start the day with this not like a panic but like a" + ), + "audio": prompt_filepath_conversational_a + }, + "conversational_b": { + "text": ( + "like a super Mario level. Like it's very like high detail. And like, once you get " + "into the park, it just like, everything looks like a computer game and they have all " + "these, like, you know, if, if there's like a, you know, like in a Mario game, they " + "will have like a question block. And if you like, you know, punch it, a coin will " + "come out. So like everyone, when they come into the park, they get like this little " + "bracelet and then you can go punching question blocks around." + ), + "audio": prompt_filepath_conversational_b + } +} + +def load_prompt_audio(audio_path: str, target_sample_rate: int) -> torch.Tensor: + audio_tensor, sample_rate = torchaudio.load(audio_path) + audio_tensor = audio_tensor.squeeze(0) + # Resample is lazy so we can always call it + audio_tensor = torchaudio.functional.resample( + audio_tensor, orig_freq=sample_rate, new_freq=target_sample_rate + ) + return audio_tensor + +def prepare_prompt(text: str, speaker: int, audio_path: str, sample_rate: int) -> Segment: + audio_tensor = load_prompt_audio(audio_path, sample_rate) + return Segment(text=text, speaker=speaker, audio=audio_tensor) + +def main(): + # Select the best available device, skipping MPS due to float64 limitations + if torch.cuda.is_available(): + device = "cuda" + else: + device = "cpu" + print(f"Using device: {device}") + + # Load model + generator = load_csm_1b(device) + + # Prepare prompts + prompt_a = prepare_prompt( + SPEAKER_PROMPTS["conversational_a"]["text"], + 0, + SPEAKER_PROMPTS["conversational_a"]["audio"], + generator.sample_rate + ) + + prompt_b = prepare_prompt( + SPEAKER_PROMPTS["conversational_b"]["text"], + 1, + SPEAKER_PROMPTS["conversational_b"]["audio"], + generator.sample_rate + ) + + # Generate conversation + conversation = [ + {"text": "Hey how are you doing?", "speaker_id": 0}, + {"text": "Pretty good, pretty good. How about you?", "speaker_id": 1}, + {"text": "I'm great! So happy to be speaking with you today.", "speaker_id": 0}, + {"text": "Me too! This is some cool stuff, isn't it?", "speaker_id": 1} + ] + + # Generate each utterance + generated_segments = [] + prompt_segments = [prompt_a, prompt_b] + + for utterance in conversation: + print(f"Generating: {utterance['text']}") + audio_tensor = generator.generate( + text=utterance['text'], + speaker=utterance['speaker_id'], + context=prompt_segments + generated_segments, + max_audio_length_ms=10_000, + ) + generated_segments.append(Segment(text=utterance['text'], speaker=utterance['speaker_id'], audio=audio_tensor)) + + # Concatenate all generations + all_audio = torch.cat([seg.audio for seg in generated_segments], dim=0) + torchaudio.save( + "full_conversation.wav", + all_audio.unsqueeze(0).cpu(), + generator.sample_rate + ) + print("Successfully generated full_conversation.wav") + +if __name__ == "__main__": + main() \ No newline at end of file diff --git a/Backend/server.py b/Backend/server.py index 2069b29..ef9fbda 100644 --- a/Backend/server.py +++ b/Backend/server.py @@ -1,53 +1,426 @@ import os -import logging -import torch -import eventlet +import io import base64 +import time +import json +import uuid +import logging +import threading +import queue import tempfile -from io import BytesIO -from flask import Flask, render_template, request, jsonify -from flask_socketio import SocketIO, emit -import whisper +from typing import Dict, List, Optional, Tuple + +import torch import torchaudio -from src.models.conversation import Segment -from src.services.tts_service import load_csm_1b -from src.llm.generator import generate_llm_response -from transformers import AutoTokenizer, AutoModelForCausalLM -from src.audio.streaming import AudioStreamer -from src.services.transcription_service import TranscriptionService -from src.services.tts_service import TextToSpeechService +import numpy as np +from flask import Flask, request, jsonify, send_from_directory +from flask_socketio import SocketIO, emit +from flask_cors import CORS +from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline + +from generator import load_csm_1b, Segment +from dataclasses import dataclass # Configure logging -logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s') +logging.basicConfig(level=logging.INFO, + format='%(asctime)s - %(name)s - %(levelname)s - %(message)s') logger = logging.getLogger(__name__) -app = Flask(__name__, static_folder='static', template_folder='templates') -app.config['SECRET_KEY'] = os.getenv('SECRET_KEY', 'your-secret-key') -socketio = SocketIO(app) +# Initialize Flask app +app = Flask(__name__, static_folder='.') +CORS(app) +socketio = SocketIO(app, cors_allowed_origins="*", ping_timeout=120) -# Initialize services -transcription_service = TranscriptionService() -tts_service = TextToSpeechService() -audio_streamer = AudioStreamer() +# Configure device +if torch.cuda.is_available(): + DEVICE = "cuda" +elif torch.backends.mps.is_available(): + DEVICE = "mps" +else: + DEVICE = "cpu" -@socketio.on('audio_input') -def handle_audio_input(data): - audio_chunk = data['audio'] - speaker_id = data['speaker'] +logger.info(f"Using device: {DEVICE}") + +# Global variables +active_conversations = {} +user_queues = {} +processing_threads = {} + +# Load models +@dataclass +class AppModels: + generator = None + tokenizer = None + llm = None + asr = None + +models = AppModels() + +def load_models(): + """Load all required models""" + global models - # Process audio and convert to text - text = transcription_service.transcribe(audio_chunk) - logging.info(f"Transcribed text: {text}") - - # Generate response using Llama 3.2 - response_text = tts_service.generate_response(text, speaker_id) - logging.info(f"Generated response: {response_text}") - - # Convert response text to audio - audio_response = tts_service.text_to_speech(response_text, speaker_id) + logger.info("Loading CSM 1B model...") + models.generator = load_csm_1b(device=DEVICE) - # Stream audio response back to client - socketio.emit('audio_response', {'audio': audio_response}) + logger.info("Loading ASR pipeline...") + models.asr = pipeline( + "automatic-speech-recognition", + model="openai/whisper-small", + device=DEVICE + ) + + logger.info("Loading Llama 3.2 model...") + models.llm = AutoModelForCausalLM.from_pretrained( + "meta-llama/Llama-3.2-1B", + device_map=DEVICE, + torch_dtype=torch.bfloat16 + ) + models.tokenizer = AutoTokenizer.from_pretrained("meta-llama/Llama-3.2-1B") +# Load models in a background thread +threading.Thread(target=load_models, daemon=True).start() + +# Conversation data structure +class Conversation: + def __init__(self, session_id): + self.session_id = session_id + self.segments: List[Segment] = [] + self.current_speaker = 0 + self.last_activity = time.time() + self.is_processing = False + + def add_segment(self, text, speaker, audio): + segment = Segment(text=text, speaker=speaker, audio=audio) + self.segments.append(segment) + self.last_activity = time.time() + return segment + + def get_context(self, max_segments=10): + """Return the most recent segments for context""" + return self.segments[-max_segments:] if self.segments else [] + +# Routes +@app.route('/') +def index(): + return send_from_directory('.', 'index.html') + +@app.route('/api/health') +def health_check(): + return jsonify({ + "status": "ok", + "cuda_available": torch.cuda.is_available(), + "models_loaded": models.generator is not None and models.llm is not None + }) + +# Socket event handlers +@socketio.on('connect') +def handle_connect(): + session_id = request.sid + logger.info(f"Client connected: {session_id}") + + # Initialize conversation data + if session_id not in active_conversations: + active_conversations[session_id] = Conversation(session_id) + user_queues[session_id] = queue.Queue() + processing_threads[session_id] = threading.Thread( + target=process_audio_queue, + args=(session_id, user_queues[session_id]), + daemon=True + ) + processing_threads[session_id].start() + + emit('connection_status', {'status': 'connected'}) + +@socketio.on('disconnect') +def handle_disconnect(): + session_id = request.sid + logger.info(f"Client disconnected: {session_id}") + + # Cleanup + if session_id in active_conversations: + # Mark for deletion rather than immediately removing + # as the processing thread might still be accessing it + active_conversations[session_id].is_processing = False + user_queues[session_id].put(None) # Signal thread to terminate + +@socketio.on('start_stream') +def handle_start_stream(): + session_id = request.sid + logger.info(f"Starting stream for client: {session_id}") + emit('streaming_status', {'status': 'active'}) + +@socketio.on('stop_stream') +def handle_stop_stream(): + session_id = request.sid + logger.info(f"Stopping stream for client: {session_id}") + emit('streaming_status', {'status': 'inactive'}) + +@socketio.on('clear_context') +def handle_clear_context(): + session_id = request.sid + logger.info(f"Clearing context for client: {session_id}") + + if session_id in active_conversations: + active_conversations[session_id].segments = [] + emit('context_updated', {'status': 'cleared'}) + +@socketio.on('audio_chunk') +def handle_audio_chunk(data): + session_id = request.sid + audio_data = data.get('audio', '') + speaker_id = int(data.get('speaker', 0)) + + if not audio_data or not session_id in active_conversations: + return + + # Update the current speaker + active_conversations[session_id].current_speaker = speaker_id + + # Queue audio for processing + user_queues[session_id].put({ + 'audio': audio_data, + 'speaker': speaker_id + }) + + emit('processing_status', {'status': 'transcribing'}) + +def process_audio_queue(session_id, q): + """Background thread to process audio chunks for a session""" + logger.info(f"Started processing thread for session: {session_id}") + + try: + while session_id in active_conversations: + try: + # Get the next audio chunk with a timeout + data = q.get(timeout=120) + if data is None: # Termination signal + break + + # Process the audio and generate a response + process_audio_and_respond(session_id, data) + + except queue.Empty: + # Timeout, check if session is still valid + continue + except Exception as e: + logger.error(f"Error processing audio for {session_id}: {str(e)}") + socketio.emit('error', {'message': str(e)}, room=session_id) + finally: + logger.info(f"Ending processing thread for session: {session_id}") + # Clean up when thread is done + with app.app_context(): + if session_id in active_conversations: + del active_conversations[session_id] + if session_id in user_queues: + del user_queues[session_id] + +def process_audio_and_respond(session_id, data): + """Process audio data and generate a response""" + if models.generator is None or models.asr is None or models.llm is None: + socketio.emit('error', {'message': 'Models still loading, please wait'}, room=session_id) + return + + conversation = active_conversations[session_id] + + try: + # Set processing flag + conversation.is_processing = True + + # Process base64 audio data + audio_data = data['audio'] + speaker_id = data['speaker'] + + # Convert from base64 to WAV + audio_bytes = base64.b64decode(audio_data.split(',')[1]) + + # Save to temporary file for processing + with tempfile.NamedTemporaryFile(suffix='.wav', delete=False) as temp_file: + temp_file.write(audio_bytes) + temp_path = temp_file.name + + try: + # Load audio file + waveform, sample_rate = torchaudio.load(temp_path) + + # Normalize to mono if needed + if waveform.shape[0] > 1: + waveform = torch.mean(waveform, dim=0, keepdim=True) + + # Resample to the CSM sample rate if needed + if sample_rate != models.generator.sample_rate: + waveform = torchaudio.functional.resample( + waveform, + orig_freq=sample_rate, + new_freq=models.generator.sample_rate + ) + + # Transcribe audio + socketio.emit('processing_status', {'status': 'transcribing'}, room=session_id) + + # Use the ASR pipeline to transcribe + transcription_result = models.asr( + {"array": waveform.squeeze().cpu().numpy(), "sampling_rate": models.generator.sample_rate}, + return_timestamps=False + ) + user_text = transcription_result['text'].strip() + + # If no text was recognized, don't process further + if not user_text: + socketio.emit('error', {'message': 'No speech detected'}, room=session_id) + return + + # Add the user's message to conversation history + user_segment = conversation.add_segment( + text=user_text, + speaker=speaker_id, + audio=waveform.squeeze() + ) + + # Send transcription to client + socketio.emit('transcription', { + 'text': user_text, + 'speaker': speaker_id + }, room=session_id) + + # Generate AI response using Llama + socketio.emit('processing_status', {'status': 'generating'}, room=session_id) + + # Create prompt from conversation history + conversation_history = "" + for segment in conversation.segments[-5:]: # Last 5 segments for context + role = "User" if segment.speaker == 0 else "Assistant" + conversation_history += f"{role}: {segment.text}\n" + + # Add final prompt + prompt = f"{conversation_history}Assistant: " + + # Generate response with Llama + input_ids = models.tokenizer(prompt, return_tensors="pt").input_ids.to(DEVICE) + + with torch.no_grad(): + generated_ids = models.llm.generate( + input_ids, + max_new_tokens=100, + temperature=0.7, + top_p=0.9, + do_sample=True, + pad_token_id=models.tokenizer.eos_token_id + ) + + # Decode the response + response_text = models.tokenizer.decode( + generated_ids[0][input_ids.shape[1]:], + skip_special_tokens=True + ).strip() + + # Synthesize speech + socketio.emit('processing_status', {'status': 'synthesizing'}, room=session_id) + + # Generate audio with CSM + ai_speaker_id = 1 # Use speaker 1 for AI responses + + # Start sending the audio response + socketio.emit('audio_response_start', { + 'text': response_text, + 'total_chunks': 1, + 'chunk_index': 0 + }, room=session_id) + + # Generate audio + audio_tensor = models.generator.generate( + text=response_text, + speaker=ai_speaker_id, + context=conversation.get_context(), + max_audio_length_ms=10_000, + temperature=0.9 + ) + + # Add AI response to conversation history + ai_segment = conversation.add_segment( + text=response_text, + speaker=ai_speaker_id, + audio=audio_tensor + ) + + # Convert audio to WAV format + with io.BytesIO() as wav_io: + torchaudio.save( + wav_io, + audio_tensor.unsqueeze(0).cpu(), + models.generator.sample_rate, + format="wav" + ) + wav_io.seek(0) + wav_data = wav_io.read() + + # Convert WAV data to base64 + audio_base64 = f"data:audio/wav;base64,{base64.b64encode(wav_data).decode('utf-8')}" + + # Send audio chunk to client + socketio.emit('audio_response_chunk', { + 'chunk': audio_base64, + 'chunk_index': 0, + 'total_chunks': 1, + 'is_last': True + }, room=session_id) + + # Signal completion + socketio.emit('audio_response_complete', { + 'text': response_text + }, room=session_id) + + finally: + # Clean up temp file + if os.path.exists(temp_path): + os.unlink(temp_path) + + except Exception as e: + logger.error(f"Error processing audio: {str(e)}") + socketio.emit('error', {'message': f'Error: {str(e)}'}, room=session_id) + finally: + # Reset processing flag + conversation.is_processing = False + +# Error handler +@socketio.on_error() +def error_handler(e): + logger.error(f"SocketIO error: {str(e)}") + +# Periodic cleanup of inactive sessions +def cleanup_inactive_sessions(): + """Remove sessions that have been inactive for too long""" + current_time = time.time() + inactive_timeout = 3600 # 1 hour + + for session_id in list(active_conversations.keys()): + conversation = active_conversations[session_id] + if (current_time - conversation.last_activity > inactive_timeout and + not conversation.is_processing): + + logger.info(f"Cleaning up inactive session: {session_id}") + + # Signal processing thread to terminate + if session_id in user_queues: + user_queues[session_id].put(None) + + # Remove from active conversations + del active_conversations[session_id] + +# Start cleanup thread +def start_cleanup_thread(): + while True: + try: + cleanup_inactive_sessions() + except Exception as e: + logger.error(f"Error in cleanup thread: {str(e)}") + time.sleep(300) # Run every 5 minutes + +cleanup_thread = threading.Thread(target=start_cleanup_thread, daemon=True) +cleanup_thread.start() + +# Start the server if __name__ == '__main__': - socketio.run(app, host='0.0.0.0', port=5000) \ No newline at end of file + port = int(os.environ.get('PORT', 5000)) + logger.info(f"Starting server on port {port}") + socketio.run(app, host='0.0.0.0', port=port, debug=False, allow_unsafe_werkzeug=True) \ No newline at end of file diff --git a/Backend/setup.py b/Backend/setup.py new file mode 100644 index 0000000..8eddb95 --- /dev/null +++ b/Backend/setup.py @@ -0,0 +1,13 @@ +from setuptools import setup, find_packages +import os + +# Read requirements from requirements.txt +with open('requirements.txt') as f: + requirements = [line.strip() for line in f if line.strip() and not line.startswith('#')] + +setup( + name='csm', + version='0.1.0', + packages=find_packages(), + install_requires=requirements, +) diff --git a/Backend/src/audio/processor.py b/Backend/src/audio/processor.py deleted file mode 100644 index 40d636e..0000000 --- a/Backend/src/audio/processor.py +++ /dev/null @@ -1,28 +0,0 @@ -from scipy.io import wavfile -import numpy as np -import torchaudio - -def load_audio(file_path): - sample_rate, audio_data = wavfile.read(file_path) - return sample_rate, audio_data - -def normalize_audio(audio_data): - audio_data = audio_data.astype(np.float32) - max_val = np.max(np.abs(audio_data)) - if max_val > 0: - audio_data /= max_val - return audio_data - -def reduce_noise(audio_data, noise_factor=0.1): - noise = np.random.randn(len(audio_data)) - noisy_audio = audio_data + noise_factor * noise - return noisy_audio - -def save_audio(file_path, sample_rate, audio_data): - torchaudio.save(file_path, torch.tensor(audio_data).unsqueeze(0), sample_rate) - -def process_audio(file_path, output_path): - sample_rate, audio_data = load_audio(file_path) - normalized_audio = normalize_audio(audio_data) - denoised_audio = reduce_noise(normalized_audio) - save_audio(output_path, sample_rate, denoised_audio) \ No newline at end of file diff --git a/Backend/src/audio/streaming.py b/Backend/src/audio/streaming.py deleted file mode 100644 index 19ee4cb..0000000 --- a/Backend/src/audio/streaming.py +++ /dev/null @@ -1,35 +0,0 @@ -from flask import Blueprint, request -from flask_socketio import SocketIO, emit -from src.audio.processor import process_audio -from src.services.transcription_service import TranscriptionService -from src.services.tts_service import TextToSpeechService - -streaming_bp = Blueprint('streaming', __name__) -socketio = SocketIO() - -transcription_service = TranscriptionService() -tts_service = TextToSpeechService() - -@socketio.on('audio_stream') -def handle_audio_stream(data): - audio_chunk = data['audio'] - speaker_id = data['speaker'] - - # Process the audio chunk - processed_audio = process_audio(audio_chunk) - - # Transcribe the audio to text - transcription = transcription_service.transcribe(processed_audio) - - # Generate a response using the LLM - response_text = generate_response(transcription, speaker_id) - - # Convert the response text back to audio - response_audio = tts_service.convert_text_to_speech(response_text, speaker_id) - - # Emit the response audio back to the client - emit('audio_response', {'audio': response_audio}) - -def generate_response(transcription, speaker_id): - # Placeholder for the actual response generation logic - return f"Response to: {transcription}" \ No newline at end of file diff --git a/Backend/src/llm/tokenizer.py b/Backend/src/llm/tokenizer.py deleted file mode 100644 index 0a05bcd..0000000 --- a/Backend/src/llm/tokenizer.py +++ /dev/null @@ -1,14 +0,0 @@ -from transformers import AutoTokenizer - -def load_llama3_tokenizer(): - tokenizer_name = "meta-llama/Llama-3.2-1B" - tokenizer = AutoTokenizer.from_pretrained(tokenizer_name) - return tokenizer - -def tokenize_text(text: str, tokenizer) -> list: - tokens = tokenizer.encode(text, return_tensors='pt') - return tokens - -def decode_tokens(tokens: list, tokenizer) -> str: - text = tokenizer.decode(tokens, skip_special_tokens=True) - return text \ No newline at end of file diff --git a/Backend/src/models/audio_model.py b/Backend/src/models/audio_model.py deleted file mode 100644 index 726bec4..0000000 --- a/Backend/src/models/audio_model.py +++ /dev/null @@ -1,28 +0,0 @@ -from dataclasses import dataclass -import torch - -@dataclass -class AudioModel: - model: torch.nn.Module - sample_rate: int - - def __post_init__(self): - self.model.eval() - - def process_audio(self, audio_tensor: torch.Tensor) -> torch.Tensor: - with torch.no_grad(): - processed_audio = self.model(audio_tensor) - return processed_audio - - def resample_audio(self, audio_tensor: torch.Tensor, target_sample_rate: int) -> torch.Tensor: - if self.sample_rate != target_sample_rate: - resampled_audio = torchaudio.functional.resample(audio_tensor, orig_freq=self.sample_rate, new_freq=target_sample_rate) - return resampled_audio - return audio_tensor - - def save_model(self, path: str): - torch.save(self.model.state_dict(), path) - - def load_model(self, path: str): - self.model.load_state_dict(torch.load(path)) - self.model.eval() \ No newline at end of file diff --git a/Backend/src/models/conversation.py b/Backend/src/models/conversation.py deleted file mode 100644 index 25d1a70..0000000 --- a/Backend/src/models/conversation.py +++ /dev/null @@ -1,51 +0,0 @@ -from dataclasses import dataclass, field -from typing import List, Optional -import torch - -@dataclass -class Segment: - speaker: int - text: str - # (num_samples,), sample_rate = 24_000 - audio: Optional[torch.Tensor] = None - - def __post_init__(self): - # Ensure audio is a tensor if provided - if self.audio is not None and not isinstance(self.audio, torch.Tensor): - self.audio = torch.tensor(self.audio, dtype=torch.float32) - -@dataclass -class Conversation: - context: List[str] = field(default_factory=list) - segments: List[Segment] = field(default_factory=list) - current_speaker: Optional[int] = None - - def add_message(self, message: str, speaker: int): - self.context.append(f"Speaker {speaker}: {message}") - self.current_speaker = speaker - - def add_segment(self, segment: Segment): - self.segments.append(segment) - self.context.append(f"Speaker {segment.speaker}: {segment.text}") - self.current_speaker = segment.speaker - - def get_context(self) -> List[str]: - return self.context - - def get_segments(self) -> List[Segment]: - return self.segments - - def clear_context(self): - self.context.clear() - self.segments.clear() - self.current_speaker = None - - def get_last_message(self) -> Optional[str]: - if self.context: - return self.context[-1] - return None - - def get_last_segment(self) -> Optional[Segment]: - if self.segments: - return self.segments[-1] - return None \ No newline at end of file diff --git a/Backend/src/services/transcription_service.py b/Backend/src/services/transcription_service.py deleted file mode 100644 index 06f8dd1..0000000 --- a/Backend/src/services/transcription_service.py +++ /dev/null @@ -1,25 +0,0 @@ -from typing import List -import torchaudio -import torch -from generator import load_csm_1b, Segment - -class TranscriptionService: - def __init__(self, model_device: str = "cpu"): - self.generator = load_csm_1b(device=model_device) - - def transcribe_audio(self, audio_path: str) -> str: - audio_tensor, sample_rate = torchaudio.load(audio_path) - audio_tensor = self._resample_audio(audio_tensor, sample_rate) - transcription = self.generator.generate_transcription(audio_tensor) - return transcription - - def _resample_audio(self, audio_tensor: torch.Tensor, orig_freq: int) -> torch.Tensor: - target_sample_rate = self.generator.sample_rate - if orig_freq != target_sample_rate: - audio_tensor = torchaudio.functional.resample(audio_tensor.squeeze(0), orig_freq=orig_freq, new_freq=target_sample_rate) - return audio_tensor - - def transcribe_audio_stream(self, audio_chunks: List[torch.Tensor]) -> str: - combined_audio = torch.cat(audio_chunks, dim=1) - transcription = self.generator.generate_transcription(combined_audio) - return transcription \ No newline at end of file diff --git a/Backend/src/services/tts_service.py b/Backend/src/services/tts_service.py deleted file mode 100644 index 64fab04..0000000 --- a/Backend/src/services/tts_service.py +++ /dev/null @@ -1,24 +0,0 @@ -from dataclasses import dataclass -import torch -import torchaudio -from huggingface_hub import hf_hub_download -from src.llm.generator import load_csm_1b - -@dataclass -class TextToSpeechService: - generator: any - - def __init__(self, device: str = "cuda"): - self.generator = load_csm_1b(device=device) - - def text_to_speech(self, text: str, speaker: int = 0) -> torch.Tensor: - audio = self.generator.generate( - text=text, - speaker=speaker, - context=[], - max_audio_length_ms=10000, - ) - return audio - - def save_audio(self, audio: torch.Tensor, file_path: str): - torchaudio.save(file_path, audio.unsqueeze(0).cpu(), self.generator.sample_rate) \ No newline at end of file diff --git a/Backend/src/utils/config.py b/Backend/src/utils/config.py deleted file mode 100644 index 2206481..0000000 --- a/Backend/src/utils/config.py +++ /dev/null @@ -1,23 +0,0 @@ -# filepath: /csm-conversation-bot/csm-conversation-bot/src/utils/config.py - -import os - -class Config: - # General configuration - DEBUG = os.getenv('DEBUG', 'False') == 'True' - SECRET_KEY = os.getenv('SECRET_KEY', 'your_secret_key_here') - - # API configuration - API_URL = os.getenv('API_URL', 'http://localhost:5000') - - # Model configuration - LLM_MODEL_PATH = os.getenv('LLM_MODEL_PATH', 'path/to/llm/model') - AUDIO_MODEL_PATH = os.getenv('AUDIO_MODEL_PATH', 'path/to/audio/model') - - # Socket.IO configuration - SOCKETIO_MESSAGE_QUEUE = os.getenv('SOCKETIO_MESSAGE_QUEUE', 'redis://localhost:6379/0') - - # Logging configuration - LOG_LEVEL = os.getenv('LOG_LEVEL', 'INFO') - - # Other configurations can be added as needed \ No newline at end of file diff --git a/Backend/src/utils/logger.py b/Backend/src/utils/logger.py deleted file mode 100644 index 93e8966..0000000 --- a/Backend/src/utils/logger.py +++ /dev/null @@ -1,14 +0,0 @@ -import logging - -def setup_logger(name, log_file, level=logging.INFO): - handler = logging.FileHandler(log_file) - handler.setFormatter(logging.Formatter('%(asctime)s - %(name)s - %(levelname)s - %(message)s')) - - logger = logging.getLogger(name) - logger.setLevel(level) - logger.addHandler(handler) - - return logger - -# Example usage: -# logger = setup_logger('my_logger', 'app.log') \ No newline at end of file diff --git a/Backend/static/css/styles.css b/Backend/static/css/styles.css deleted file mode 100644 index 4e2d752..0000000 --- a/Backend/static/css/styles.css +++ /dev/null @@ -1,105 +0,0 @@ -body { - font-family: 'Arial', sans-serif; - background-color: #f4f4f4; - color: #333; - margin: 0; - padding: 0; -} - -header { - background: #4c84ff; - color: #fff; - padding: 10px 0; - text-align: center; -} - -h1 { - margin: 0; - font-size: 2.5rem; -} - -.container { - width: 80%; - margin: auto; - overflow: hidden; -} - -.conversation { - background: #fff; - padding: 20px; - border-radius: 5px; - box-shadow: 0 2px 5px rgba(0, 0, 0, 0.1); - max-height: 400px; - overflow-y: auto; -} - -.message { - padding: 10px; - margin: 10px 0; - border-radius: 5px; -} - -.user { - background: #e3f2fd; - text-align: right; -} - -.ai { - background: #f1f1f1; - text-align: left; -} - -.controls { - display: flex; - justify-content: space-between; - margin-top: 20px; -} - -button { - padding: 10px 15px; - border: none; - border-radius: 5px; - cursor: pointer; - transition: background 0.3s; -} - -button:hover { - background: #3367d6; - color: #fff; -} - -.visualizer-container { - height: 150px; - background: #000; - border-radius: 5px; - margin-top: 20px; -} - -.visualizer-label { - color: rgba(255, 255, 255, 0.7); - text-align: center; - padding: 10px; -} - -.status-indicator { - display: flex; - align-items: center; - margin-top: 10px; -} - -.status-dot { - width: 12px; - height: 12px; - border-radius: 50%; - background-color: #ccc; - margin-right: 10px; -} - -.status-dot.active { - background-color: #4CAF50; -} - -.status-text { - font-size: 0.9em; - color: #666; -} \ No newline at end of file diff --git a/Backend/static/index.html b/Backend/static/index.html deleted file mode 100644 index 4922f17..0000000 --- a/Backend/static/index.html +++ /dev/null @@ -1,31 +0,0 @@ - - - - - - CSM Conversation Bot - - - - - -
-

CSM Conversation Bot

-

Talk to the AI and get responses in real-time!

-
-
-
-
- - -
-
-
-
Disconnected
-
-
- - - \ No newline at end of file diff --git a/Backend/static/js/client.js b/Backend/static/js/client.js deleted file mode 100644 index ec4037f..0000000 --- a/Backend/static/js/client.js +++ /dev/null @@ -1,131 +0,0 @@ -// This file contains the client-side JavaScript code that handles audio streaming and communication with the server. - -const SERVER_URL = window.location.hostname === 'localhost' ? - 'http://localhost:5000' : window.location.origin; - -const elements = { - conversation: document.getElementById('conversation'), - streamButton: document.getElementById('streamButton'), - clearButton: document.getElementById('clearButton'), - speakerSelection: document.getElementById('speakerSelect'), - statusDot: document.getElementById('statusDot'), - statusText: document.getElementById('statusText'), -}; - -const state = { - socket: null, - isStreaming: false, - currentSpeaker: 0, -}; - -// Initialize the application -function initializeApp() { - setupSocketConnection(); - setupEventListeners(); -} - -// Setup Socket.IO connection -function setupSocketConnection() { - state.socket = io(SERVER_URL); - - state.socket.on('connect', () => { - updateConnectionStatus(true); - }); - - state.socket.on('disconnect', () => { - updateConnectionStatus(false); - }); - - state.socket.on('audio_response', handleAudioResponse); - state.socket.on('transcription', handleTranscription); -} - -// Setup event listeners -function setupEventListeners() { - elements.streamButton.addEventListener('click', toggleStreaming); - elements.clearButton.addEventListener('click', clearConversation); - elements.speakerSelection.addEventListener('change', (event) => { - state.currentSpeaker = event.target.value; - }); -} - -// Update connection status UI -function updateConnectionStatus(isConnected) { - elements.statusDot.classList.toggle('active', isConnected); - elements.statusText.textContent = isConnected ? 'Connected' : 'Disconnected'; -} - -// Toggle streaming state -function toggleStreaming() { - if (state.isStreaming) { - stopStreaming(); - } else { - startStreaming(); - } -} - -// Start streaming audio to the server -function startStreaming() { - if (state.isStreaming) return; - - navigator.mediaDevices.getUserMedia({ audio: true }) - .then(stream => { - const mediaRecorder = new MediaRecorder(stream); - mediaRecorder.start(); - - mediaRecorder.ondataavailable = (event) => { - if (event.data.size > 0) { - sendAudioChunk(event.data); - } - }; - - mediaRecorder.onstop = () => { - state.isStreaming = false; - elements.streamButton.innerHTML = 'Start Conversation'; - }; - - state.isStreaming = true; - elements.streamButton.innerHTML = 'Stop Conversation'; - }) - .catch(err => { - console.error('Error accessing microphone:', err); - }); -} - -// Stop streaming audio -function stopStreaming() { - if (!state.isStreaming) return; - - // Logic to stop the media recorder would go here -} - -// Send audio chunk to server -function sendAudioChunk(audioData) { - const reader = new FileReader(); - reader.onloadend = () => { - const arrayBuffer = reader.result; - state.socket.emit('audio_chunk', { audio: arrayBuffer, speaker: state.currentSpeaker }); - }; - reader.readAsArrayBuffer(audioData); -} - -// Handle audio response from server -function handleAudioResponse(data) { - const audioElement = new Audio(URL.createObjectURL(new Blob([data.audio]))); - audioElement.play(); -} - -// Handle transcription response from server -function handleTranscription(data) { - const messageElement = document.createElement('div'); - messageElement.textContent = `AI: ${data.transcription}`; - elements.conversation.appendChild(messageElement); -} - -// Clear conversation history -function clearConversation() { - elements.conversation.innerHTML = ''; -} - -// Initialize the application when DOM is fully loaded -document.addEventListener('DOMContentLoaded', initializeApp); \ No newline at end of file diff --git a/Backend/templates/index.html b/Backend/templates/index.html deleted file mode 100644 index 514b946..0000000 --- a/Backend/templates/index.html +++ /dev/null @@ -1,31 +0,0 @@ - - - - - - CSM Conversation Bot - - - - - -
-

CSM Conversation Bot

-

Talk to the AI and get responses in real-time!

-
-
-
-
- - -
-
-
-
Not connected
-
-
- - - \ No newline at end of file diff --git a/Backend/voice-chat.js b/Backend/voice-chat.js new file mode 100644 index 0000000..89ec71a --- /dev/null +++ b/Backend/voice-chat.js @@ -0,0 +1,1071 @@ +/** + * Sesame AI Voice Chat Client + * + * A web client that connects to a Sesame AI voice chat server and enables + * real-time voice conversation with an AI assistant. + */ + +// Configuration constants +const SERVER_URL = window.location.hostname === 'localhost' ? + 'http://localhost:5000' : window.location.origin; +const ENERGY_WINDOW_SIZE = 15; +const CLIENT_SILENCE_DURATION_MS = 750; + +// DOM elements +const elements = { + conversation: null, + streamButton: null, + clearButton: null, + thresholdSlider: null, + thresholdValue: null, + visualizerCanvas: null, + visualizerLabel: null, + volumeLevel: null, + statusDot: null, + statusText: null, + speakerSelection: null, + autoPlayResponses: null, + showVisualizer: null +}; + +// Application state +const state = { + socket: null, + audioContext: null, + analyser: null, + microphone: null, + streamProcessor: null, + isStreaming: false, + isSpeaking: false, + silenceThreshold: 0.01, + energyWindow: [], + silenceTimer: null, + volumeUpdateInterval: null, + visualizerAnimationFrame: null, + currentSpeaker: 0 +}; + +// Visualizer variables +let canvasContext = null; +let visualizerBufferLength = 0; +let visualizerDataArray = null; + +// New state variables to track incremental audio streaming +const streamingAudio = { + messageElement: null, + audioElement: null, + chunks: [], + totalChunks: 0, + receivedChunks: 0, + text: '', + mediaSource: null, + sourceBuffer: null, + audioContext: null, + complete: false +}; + +// Initialize the application +function initializeApp() { + // Initialize the UI elements + initializeUIElements(); + + // Initialize socket.io connection + setupSocketConnection(); + + // Setup event listeners + setupEventListeners(); + + // Initialize visualizer + setupVisualizer(); + + // Show welcome message + addSystemMessage('Welcome to Sesame AI Voice Chat! Click "Start Conversation" to begin.'); +} + +// Initialize UI elements +function initializeUIElements() { + // Store references to UI elements + elements.conversation = document.getElementById('conversation'); + elements.streamButton = document.getElementById('streamButton'); + elements.clearButton = document.getElementById('clearButton'); + elements.thresholdSlider = document.getElementById('thresholdSlider'); + elements.thresholdValue = document.getElementById('thresholdValue'); + elements.visualizerCanvas = document.getElementById('audioVisualizer'); + elements.visualizerLabel = document.getElementById('visualizerLabel'); + elements.volumeLevel = document.getElementById('volumeLevel'); + elements.statusDot = document.getElementById('statusDot'); + elements.statusText = document.getElementById('statusText'); + elements.speakerSelection = document.getElementById('speakerSelect'); // Changed to match HTML + elements.autoPlayResponses = document.getElementById('autoPlayResponses'); + elements.showVisualizer = document.getElementById('showVisualizer'); +} + +// Setup Socket.IO connection +function setupSocketConnection() { + state.socket = io(SERVER_URL); + + // Connection events + state.socket.on('connect', () => { + console.log('Connected to server'); + updateConnectionStatus(true); + }); + + state.socket.on('disconnect', () => { + console.log('Disconnected from server'); + updateConnectionStatus(false); + + // Stop streaming if active + if (state.isStreaming) { + stopStreaming(false); + } + }); + + state.socket.on('error', (data) => { + console.error('Socket error:', data.message); + addSystemMessage(`Error: ${data.message}`); + }); + + // Register message handlers + state.socket.on('audio_response', handleAudioResponse); + state.socket.on('transcription', handleTranscription); + state.socket.on('context_updated', handleContextUpdate); + state.socket.on('streaming_status', handleStreamingStatus); + + // New event handlers for incremental audio streaming + state.socket.on('audio_response_start', handleAudioResponseStart); + state.socket.on('audio_response_chunk', handleAudioResponseChunk); + state.socket.on('audio_response_complete', handleAudioResponseComplete); + state.socket.on('processing_status', handleProcessingStatus); +} + +// Setup event listeners +function setupEventListeners() { + // Stream button + elements.streamButton.addEventListener('click', toggleStreaming); + + // Clear button + elements.clearButton.addEventListener('click', clearConversation); + + // Threshold slider + elements.thresholdSlider.addEventListener('input', updateThreshold); + + // Speaker selection + elements.speakerSelection.addEventListener('change', () => { + state.currentSpeaker = parseInt(elements.speakerSelection.value, 10); + }); + + // Visualizer toggle + elements.showVisualizer.addEventListener('change', toggleVisualizerVisibility); +} + +// Setup audio visualizer +function setupVisualizer() { + if (!elements.visualizerCanvas) return; + + canvasContext = elements.visualizerCanvas.getContext('2d'); + + // Set canvas dimensions + elements.visualizerCanvas.width = elements.visualizerCanvas.offsetWidth; + elements.visualizerCanvas.height = elements.visualizerCanvas.offsetHeight; + + // Initialize the visualizer + drawVisualizer(); +} + +// Update connection status UI +function updateConnectionStatus(isConnected) { + elements.statusDot.classList.toggle('active', isConnected); + elements.statusText.textContent = isConnected ? 'Connected' : 'Disconnected'; +} + +// Toggle streaming state +function toggleStreaming() { + if (state.isStreaming) { + stopStreaming(true); + } else { + startStreaming(); + } +} + +// Start streaming audio to the server +function startStreaming() { + if (state.isStreaming) return; + + // Request microphone access + navigator.mediaDevices.getUserMedia({ audio: true, video: false }) + .then(stream => { + // Show processing state while setting up + elements.streamButton.innerHTML = ' Initializing...'; + + // Create audio context + state.audioContext = new (window.AudioContext || window.webkitAudioContext)(); + + // Create microphone source + state.microphone = state.audioContext.createMediaStreamSource(stream); + + // Create analyser for visualizer + state.analyser = state.audioContext.createAnalyser(); + state.analyser.fftSize = 256; + visualizerBufferLength = state.analyser.frequencyBinCount; + visualizerDataArray = new Uint8Array(visualizerBufferLength); + + // Connect microphone to analyser + state.microphone.connect(state.analyser); + + // Create script processor for audio processing + const bufferSize = 4096; + state.streamProcessor = state.audioContext.createScriptProcessor(bufferSize, 1, 1); + + // Set up audio processing callback + state.streamProcessor.onaudioprocess = handleAudioProcess; + + // Connect the processors + state.analyser.connect(state.streamProcessor); + state.streamProcessor.connect(state.audioContext.destination); + + // Update UI + state.isStreaming = true; + elements.streamButton.innerHTML = ' Listening...'; + elements.streamButton.classList.add('recording'); + + // Initialize energy window + state.energyWindow = []; + + // Start volume meter updates + state.volumeUpdateInterval = setInterval(updateVolumeMeter, 100); + + // Start visualizer if enabled + if (elements.showVisualizer.checked && !state.visualizerAnimationFrame) { + drawVisualizer(); + } + + // Show starting message + addSystemMessage('Listening... Speak clearly into your microphone.'); + + // Notify the server that we're starting + state.socket.emit('stream_audio', { + audio: '', + speaker: state.currentSpeaker + }); + }) + .catch(err => { + console.error('Error accessing microphone:', err); + addSystemMessage(`Error: ${err.message}. Please make sure your microphone is connected and you've granted permission.`); + elements.streamButton.innerHTML = ' Start Conversation'; + }); +} + +// Stop streaming audio +function stopStreaming(notifyServer = true) { + if (!state.isStreaming) return; + + // Update UI first + elements.streamButton.innerHTML = ' Start Conversation'; + elements.streamButton.classList.remove('recording'); + elements.streamButton.classList.remove('processing'); + + // Stop volume meter updates + if (state.volumeUpdateInterval) { + clearInterval(state.volumeUpdateInterval); + state.volumeUpdateInterval = null; + } + + // Stop all audio processing + if (state.streamProcessor) { + state.streamProcessor.disconnect(); + state.streamProcessor = null; + } + + if (state.analyser) { + state.analyser.disconnect(); + } + + if (state.microphone) { + state.microphone.disconnect(); + } + + // Close audio context + if (state.audioContext && state.audioContext.state !== 'closed') { + state.audioContext.close().catch(err => console.warn('Error closing audio context:', err)); + } + + // Cleanup animation frames + if (state.visualizerAnimationFrame) { + cancelAnimationFrame(state.visualizerAnimationFrame); + state.visualizerAnimationFrame = null; + } + + // Reset state + state.isStreaming = false; + state.isSpeaking = false; + + // Notify the server + if (notifyServer && state.socket && state.socket.connected) { + state.socket.emit('stop_streaming', { + speaker: state.currentSpeaker + }); + } + + // Show message + addSystemMessage('Conversation paused. Click "Start Conversation" to resume.'); +} + +// Handle audio processing +function handleAudioProcess(event) { + const inputData = event.inputBuffer.getChannelData(0); + + // Calculate audio energy (volume level) + const energy = calculateAudioEnergy(inputData); + + // Update energy window for averaging + updateEnergyWindow(energy); + + // Calculate average energy + const avgEnergy = calculateAverageEnergy(); + + // Determine if audio is silent + const isSilent = avgEnergy < state.silenceThreshold; + + // Debug logging only if significant changes in audio patterns + if (Math.random() < 0.05) { // Log only 5% of frames to avoid console spam + console.log(`Audio: len=${inputData.length}, energy=${energy.toFixed(4)}, avg=${avgEnergy.toFixed(4)}, silent=${isSilent}`); + } + + // Handle speech state based on silence + handleSpeechState(isSilent); + + // Only send audio chunk if we detect speech + if (!isSilent) { + // Create a resampled version at 24kHz for the server + // Most WebRTC audio is 48kHz, but we want 24kHz for the model + const resampledData = downsampleBuffer(inputData, state.audioContext.sampleRate, 24000); + + // Send the audio chunk to the server + sendAudioChunk(resampledData, state.currentSpeaker); + } +} + +// Cleanup audio resources when done +function cleanupAudioResources() { + // Stop all audio processing + if (state.streamProcessor) { + state.streamProcessor.disconnect(); + state.streamProcessor = null; + } + + if (state.analyser) { + state.analyser.disconnect(); + state.analyser = null; + } + + if (state.microphone) { + state.microphone.disconnect(); + state.microphone = null; + } + + // Close audio context + if (state.audioContext && state.audioContext.state !== 'closed') { + state.audioContext.close().catch(err => console.warn('Error closing audio context:', err)); + } + + // Cancel all timers and animation frames + if (state.volumeUpdateInterval) { + clearInterval(state.volumeUpdateInterval); + state.volumeUpdateInterval = null; + } + + if (state.visualizerAnimationFrame) { + cancelAnimationFrame(state.visualizerAnimationFrame); + state.visualizerAnimationFrame = null; + } + + if (state.silenceTimer) { + clearTimeout(state.silenceTimer); + state.silenceTimer = null; + } +} + +// Clear conversation history +function clearConversation() { + if (elements.conversation) { + elements.conversation.innerHTML = ''; + addSystemMessage('Conversation cleared.'); + + // Notify server to clear context + if (state.socket && state.socket.connected) { + state.socket.emit('clear_context'); + } + } +} + +// Calculate audio energy (volume) +function calculateAudioEnergy(buffer) { + let sum = 0; + for (let i = 0; i < buffer.length; i++) { + sum += buffer[i] * buffer[i]; + } + return Math.sqrt(sum / buffer.length); +} + +// Update energy window for averaging +function updateEnergyWindow(energy) { + state.energyWindow.push(energy); + if (state.energyWindow.length > ENERGY_WINDOW_SIZE) { + state.energyWindow.shift(); + } +} + +// Calculate average energy from window +function calculateAverageEnergy() { + if (state.energyWindow.length === 0) return 0; + + const sum = state.energyWindow.reduce((a, b) => a + b, 0); + return sum / state.energyWindow.length; +} + +// Update the threshold from the slider +function updateThreshold() { + state.silenceThreshold = parseFloat(elements.thresholdSlider.value); + elements.thresholdValue.textContent = state.silenceThreshold.toFixed(3); +} + +// Update the volume meter display +function updateVolumeMeter() { + if (!state.isStreaming || !state.energyWindow.length) return; + + const avgEnergy = calculateAverageEnergy(); + + // Scale energy to percentage (0-100) + // Typically, energy values will be very small (e.g., 0.001 to 0.1) + // So we multiply by a factor to make it more visible + const scaleFactor = 1000; + const percentage = Math.min(100, Math.max(0, avgEnergy * scaleFactor)); + + // Update volume meter width + elements.volumeLevel.style.width = `${percentage}%`; + + // Change color based on level + if (percentage > 70) { + elements.volumeLevel.style.backgroundColor = '#ff5252'; + } else if (percentage > 30) { + elements.volumeLevel.style.backgroundColor = '#4CAF50'; + } else { + elements.volumeLevel.style.backgroundColor = '#4c84ff'; + } +} + +// Handle speech/silence state transitions +function handleSpeechState(isSilent) { + if (state.isSpeaking && isSilent) { + // Transition from speaking to silence + if (!state.silenceTimer) { + state.silenceTimer = setTimeout(() => { + // Only consider it a real silence after a certain duration + // This prevents detecting brief pauses as the end of speech + state.isSpeaking = false; + state.silenceTimer = null; + }, CLIENT_SILENCE_DURATION_MS); + } + } else if (state.silenceTimer && !isSilent) { + // User started speaking again, cancel the silence timer + clearTimeout(state.silenceTimer); + state.silenceTimer = null; + } + + // Update speaking state for non-silent audio + if (!isSilent) { + state.isSpeaking = true; + } +} + +// Send audio chunk to server +function sendAudioChunk(audioData, speaker) { + if (!state.socket || !state.socket.connected) { + console.warn('Socket not connected'); + return; + } + + console.log(`Preparing audio chunk: length=${audioData.length}, speaker=${speaker}`); + + // Check for NaN or invalid values + let hasInvalidValues = false; + for (let i = 0; i < audioData.length; i++) { + if (isNaN(audioData[i]) || !isFinite(audioData[i])) { + hasInvalidValues = true; + console.warn(`Invalid audio value at index ${i}: ${audioData[i]}`); + break; + } + } + + if (hasInvalidValues) { + console.warn('Audio data contains invalid values. Creating silent audio.'); + audioData = new Float32Array(audioData.length).fill(0); + } + + try { + // Create WAV blob + const wavData = createWavBlob(audioData, 24000); + console.log(`WAV blob created: ${wavData.size} bytes`); + + const reader = new FileReader(); + + reader.onloadend = function() { + try { + // Get base64 data + const base64data = reader.result; + console.log(`Base64 data created: ${base64data.length} bytes`); + + // Send to server + state.socket.emit('stream_audio', { + audio: base64data, + speaker: speaker + }); + console.log('Audio chunk sent to server'); + } catch (err) { + console.error('Error preparing audio data:', err); + } + }; + + reader.onerror = function() { + console.error('Error reading audio data as base64'); + }; + + reader.readAsDataURL(wavData); + } catch (err) { + console.error('Error creating WAV data:', err); + } +} + +// Create WAV blob from audio data with improved error handling +function createWavBlob(audioData, sampleRate) { + // Validate input + if (!audioData || audioData.length === 0) { + console.warn('Empty audio data provided to createWavBlob'); + audioData = new Float32Array(1024).fill(0); // Create 1024 samples of silence + } + + // Function to convert Float32Array to Int16Array for WAV format + function floatTo16BitPCM(output, offset, input) { + for (let i = 0; i < input.length; i++, offset += 2) { + // Ensure values are in -1 to 1 range + const s = Math.max(-1, Math.min(1, input[i])); + // Convert to 16-bit PCM + output.setInt16(offset, s < 0 ? s * 0x8000 : s * 0x7FFF, true); + } + } + + // Create WAV header + function writeString(view, offset, string) { + for (let i = 0; i < string.length; i++) { + view.setUint8(offset + i, string.charCodeAt(i)); + } + } + + try { + // Create WAV file with header - careful with buffer sizes + const buffer = new ArrayBuffer(44 + audioData.length * 2); + const view = new DataView(buffer); + + // RIFF identifier + writeString(view, 0, 'RIFF'); + + // File length (will be filled later) + view.setUint32(4, 36 + audioData.length * 2, true); + + // WAVE identifier + writeString(view, 8, 'WAVE'); + + // fmt chunk identifier + writeString(view, 12, 'fmt '); + + // fmt chunk length + view.setUint32(16, 16, true); + + // Sample format (1 is PCM) + view.setUint16(20, 1, true); + + // Mono channel + view.setUint16(22, 1, true); + + // Sample rate + view.setUint32(24, sampleRate, true); + + // Byte rate (sample rate * block align) + view.setUint32(28, sampleRate * 2, true); + + // Block align (channels * bytes per sample) + view.setUint16(32, 2, true); + + // Bits per sample + view.setUint16(34, 16, true); + + // data chunk identifier + writeString(view, 36, 'data'); + + // data chunk length + view.setUint32(40, audioData.length * 2, true); + + // Write the PCM samples + floatTo16BitPCM(view, 44, audioData); + + // Create and return blob + return new Blob([view], { type: 'audio/wav' }); + } catch (err) { + console.error('Error in createWavBlob:', err); + + // Create a minimal valid WAV file with silence as fallback + const fallbackSamples = new Float32Array(1024).fill(0); + const fallbackBuffer = new ArrayBuffer(44 + fallbackSamples.length * 2); + const fallbackView = new DataView(fallbackBuffer); + + writeString(fallbackView, 0, 'RIFF'); + fallbackView.setUint32(4, 36 + fallbackSamples.length * 2, true); + writeString(fallbackView, 8, 'WAVE'); + writeString(fallbackView, 12, 'fmt '); + fallbackView.setUint32(16, 16, true); + fallbackView.setUint16(20, 1, true); + fallbackView.setUint16(22, 1, true); + fallbackView.setUint32(24, sampleRate, true); + fallbackView.setUint32(28, sampleRate * 2, true); + fallbackView.setUint16(32, 2, true); + fallbackView.setUint16(34, 16, true); + writeString(fallbackView, 36, 'data'); + fallbackView.setUint32(40, fallbackSamples.length * 2, true); + floatTo16BitPCM(fallbackView, 44, fallbackSamples); + + return new Blob([fallbackView], { type: 'audio/wav' }); + } +} + +// Draw audio visualizer +function drawVisualizer() { + if (!canvasContext) { + return; + } + + state.visualizerAnimationFrame = requestAnimationFrame(drawVisualizer); + + // Skip drawing if visualizer is hidden + if (!elements.showVisualizer.checked) { + if (elements.visualizerCanvas.style.opacity !== '0') { + elements.visualizerCanvas.style.opacity = '0'; + } + return; + } else if (elements.visualizerCanvas.style.opacity !== '1') { + elements.visualizerCanvas.style.opacity = '1'; + } + + // Get frequency data if available + if (state.isStreaming && state.analyser) { + try { + state.analyser.getByteFrequencyData(visualizerDataArray); + } catch (e) { + console.warn('Error getting frequency data:', e); + } + } else { + // Fade out when not streaming + for (let i = 0; i < visualizerDataArray.length; i++) { + visualizerDataArray[i] = Math.max(0, visualizerDataArray[i] - 5); + } + } + + // Clear canvas + canvasContext.fillStyle = 'rgb(0, 0, 0)'; + canvasContext.fillRect(0, 0, elements.visualizerCanvas.width, elements.visualizerCanvas.height); + + // Draw gradient bars + const width = elements.visualizerCanvas.width; + const height = elements.visualizerCanvas.height; + const barCount = Math.min(visualizerBufferLength, 64); + const barWidth = width / barCount - 1; + + for (let i = 0; i < barCount; i++) { + const index = Math.floor(i * visualizerBufferLength / barCount); + const value = visualizerDataArray[index]; + + // Use logarithmic scale for better audio visualization + // This makes low values more visible while still maintaining full range + const logFactor = 20; + const scaledValue = Math.log(1 + (value / 255) * logFactor) / Math.log(1 + logFactor); + const barHeight = scaledValue * height; + + // Position bars + const x = i * (barWidth + 1); + const y = height - barHeight; + + // Create color gradient based on frequency and amplitude + const hue = i / barCount * 360; // Full color spectrum + const saturation = 80 + (value / 255 * 20); // Higher values more saturated + const lightness = 40 + (value / 255 * 20); // Dynamic brightness based on amplitude + + // Draw main bar + canvasContext.fillStyle = `hsl(${hue}, ${saturation}%, ${lightness}%)`; + canvasContext.fillRect(x, y, barWidth, barHeight); + + // Add reflection effect + if (barHeight > 5) { + const gradient = canvasContext.createLinearGradient( + x, y, + x, y + barHeight * 0.5 + ); + gradient.addColorStop(0, `hsla(${hue}, ${saturation}%, ${lightness + 20}%, 0.4)`); + gradient.addColorStop(1, `hsla(${hue}, ${saturation}%, ${lightness}%, 0)`); + canvasContext.fillStyle = gradient; + canvasContext.fillRect(x, y, barWidth, barHeight * 0.5); + + // Add highlight on top of the bar for better 3D effect + canvasContext.fillStyle = `hsla(${hue}, ${saturation - 20}%, ${lightness + 30}%, 0.7)`; + canvasContext.fillRect(x, y, barWidth, 2); + } + } + + // Show/hide the label + elements.visualizerLabel.style.opacity = (state.isStreaming) ? '0' : '0.7'; +} + +// Toggle visualizer visibility +function toggleVisualizerVisibility() { + const isVisible = elements.showVisualizer.checked; + elements.visualizerCanvas.style.opacity = isVisible ? '1' : '0'; + + if (isVisible && state.isStreaming && !state.visualizerAnimationFrame) { + drawVisualizer(); + } +} + +// Handle audio response from server +function handleAudioResponse(data) { + console.log('Received audio response'); + + // Create message container + const messageElement = document.createElement('div'); + messageElement.className = 'message ai'; + + // Add text content if available + if (data.text) { + const textElement = document.createElement('p'); + textElement.textContent = data.text; + messageElement.appendChild(textElement); + } + + // Create and configure audio element + const audioElement = document.createElement('audio'); + audioElement.controls = true; + audioElement.className = 'audio-player'; + + // Set audio source + const audioSource = document.createElement('source'); + audioSource.src = data.audio; + audioSource.type = 'audio/wav'; + + // Add fallback text + audioElement.textContent = 'Your browser does not support the audio element.'; + + // Assemble audio element + audioElement.appendChild(audioSource); + messageElement.appendChild(audioElement); + + // Add timestamp + const timeElement = document.createElement('span'); + timeElement.className = 'message-time'; + timeElement.textContent = new Date().toLocaleTimeString(); + messageElement.appendChild(timeElement); + + // Add to conversation + elements.conversation.appendChild(messageElement); + + // Auto-scroll to bottom + elements.conversation.scrollTop = elements.conversation.scrollHeight; + + // Auto-play if enabled + if (elements.autoPlayResponses.checked) { + audioElement.play() + .catch(err => { + console.warn('Auto-play failed:', err); + addSystemMessage('Auto-play failed. Please click play to hear the response.'); + }); + } + + // Re-enable stream button after processing is complete + if (state.isStreaming) { + elements.streamButton.innerHTML = ' Listening...'; + elements.streamButton.classList.add('recording'); + elements.streamButton.classList.remove('processing'); + } +} + +// Handle transcription response from server +function handleTranscription(data) { + console.log('Received transcription:', data.text); + + // Create message element + const messageElement = document.createElement('div'); + messageElement.className = 'message user'; + + // Add text content + const textElement = document.createElement('p'); + textElement.textContent = data.text; + messageElement.appendChild(textElement); + + // Add timestamp + const timeElement = document.createElement('span'); + timeElement.className = 'message-time'; + timeElement.textContent = new Date().toLocaleTimeString(); + messageElement.appendChild(timeElement); + + // Add to conversation + elements.conversation.appendChild(messageElement); + + // Auto-scroll to bottom + elements.conversation.scrollTop = elements.conversation.scrollHeight; +} + +// Handle context update from server +function handleContextUpdate(data) { + console.log('Context updated:', data.message); +} + +// Handle streaming status updates from server +function handleStreamingStatus(data) { + console.log('Streaming status:', data.status); + + if (data.status === 'stopped') { + // Reset UI if needed + if (state.isStreaming) { + stopStreaming(false); // Don't send to server since this came from server + } + } +} + +// Add a system message to the conversation +function addSystemMessage(message) { + const messageElement = document.createElement('div'); + messageElement.className = 'message system'; + messageElement.textContent = message; + elements.conversation.appendChild(messageElement); + + // Auto-scroll to bottom + elements.conversation.scrollTop = elements.conversation.scrollHeight; +} + +// Downsample audio buffer to target sample rate +function downsampleBuffer(buffer, originalSampleRate, targetSampleRate) { + if (originalSampleRate === targetSampleRate) { + return buffer; + } + + const ratio = originalSampleRate / targetSampleRate; + const newLength = Math.round(buffer.length / ratio); + const result = new Float32Array(newLength); + + for (let i = 0; i < newLength; i++) { + const pos = Math.round(i * ratio); + result[i] = buffer[pos]; + } + + return result; +} + +// Handle processing status updates +function handleProcessingStatus(data) { + console.log('Processing status update:', data); + + // Show processing status in UI + if (data.status === 'generating_audio') { + elements.streamButton.innerHTML = ' Processing...'; + elements.streamButton.classList.add('processing'); + elements.streamButton.classList.remove('recording'); + + // Show message to user + addSystemMessage(data.message || 'Processing your request...'); + } +} + +// Handle the start of an audio streaming response +function handleAudioResponseStart(data) { + console.log('Audio response starting:', data); + + // Reset streaming audio state + streamingAudio.chunks = []; + streamingAudio.totalChunks = data.total_chunks; + streamingAudio.receivedChunks = 0; + streamingAudio.text = data.text; + streamingAudio.complete = false; + + // Create message container now, so we can update it as chunks arrive + const messageElement = document.createElement('div'); + messageElement.className = 'message ai processing'; + + // Add text content if available + if (data.text) { + const textElement = document.createElement('p'); + textElement.textContent = data.text; + messageElement.appendChild(textElement); + } + + // Create audio element (will be populated as chunks arrive) + const audioElement = document.createElement('audio'); + audioElement.controls = true; + audioElement.className = 'audio-player'; + audioElement.textContent = 'Audio is being generated...'; + messageElement.appendChild(audioElement); + + // Add timestamp + const timeElement = document.createElement('span'); + timeElement.className = 'message-time'; + timeElement.textContent = new Date().toLocaleTimeString(); + messageElement.appendChild(timeElement); + + // Add loading indicator + const loadingElement = document.createElement('div'); + loadingElement.className = 'loading-indicator'; + loadingElement.innerHTML = 'Generating audio response...'; + messageElement.appendChild(loadingElement); + + // Add to conversation + elements.conversation.appendChild(messageElement); + + // Auto-scroll to bottom + elements.conversation.scrollTop = elements.conversation.scrollHeight; + + // Store elements for later updates + streamingAudio.messageElement = messageElement; + streamingAudio.audioElement = audioElement; +} + +// Handle an incoming audio chunk +function handleAudioResponseChunk(data) { + console.log(`Received audio chunk ${data.chunk_index + 1}/${data.total_chunks}`); + + // Store the chunk + streamingAudio.chunks[data.chunk_index] = data.audio; + streamingAudio.receivedChunks++; + + // Update progress in the UI + if (streamingAudio.messageElement) { + const loadingElement = streamingAudio.messageElement.querySelector('.loading-indicator span'); + if (loadingElement) { + loadingElement.textContent = `Generating audio response... ${Math.round((streamingAudio.receivedChunks / data.total_chunks) * 100)}%`; + } + } + + // If this is the first chunk, start playing it immediately for faster response + if (data.chunk_index === 0 && streamingAudio.audioElement && elements.autoPlayResponses && elements.autoPlayResponses.checked) { + try { + streamingAudio.audioElement.src = data.audio; + streamingAudio.audioElement.play().catch(err => console.warn('Auto-play failed:', err)); + } catch (e) { + console.error('Error playing first chunk:', e); + } + } + + // If this is the last chunk or we've received all chunks, finalize the audio + if (data.is_last || streamingAudio.receivedChunks >= data.total_chunks) { + finalizeStreamingAudio(); + } +} + +// Handle completion of audio streaming +function handleAudioResponseComplete(data) { + console.log('Audio response complete:', data); + streamingAudio.complete = true; + + // Make sure we finalize the audio even if some chunks were missed + finalizeStreamingAudio(); + + // Update UI to normal state + if (state.isStreaming) { + elements.streamButton.innerHTML = ' Listening...'; + elements.streamButton.classList.add('recording'); + elements.streamButton.classList.remove('processing'); + } +} + +// Finalize streaming audio by combining chunks and updating the UI +function finalizeStreamingAudio() { + if (!streamingAudio.messageElement || streamingAudio.chunks.length === 0) { + return; + } + + try { + // For more sophisticated audio streaming, you would need to properly concatenate + // the WAV files, but for now we'll use the last chunk as the complete audio + // since it should contain the entire response due to how the server is implementing it + const lastChunkIndex = streamingAudio.chunks.length - 1; + const audioData = streamingAudio.chunks[lastChunkIndex] || streamingAudio.chunks[0]; + + // Update the audio element with the complete audio + if (streamingAudio.audioElement) { + streamingAudio.audioElement.src = audioData; + + // Auto-play if enabled and not already playing + if (elements.autoPlayResponses && elements.autoPlayResponses.checked && + streamingAudio.audioElement.paused) { + streamingAudio.audioElement.play() + .catch(err => { + console.warn('Auto-play failed:', err); + addSystemMessage('Auto-play failed. Please click play to hear the response.'); + }); + } + } + + // Remove loading indicator and processing class + if (streamingAudio.messageElement) { + const loadingElement = streamingAudio.messageElement.querySelector('.loading-indicator'); + if (loadingElement) { + streamingAudio.messageElement.removeChild(loadingElement); + } + streamingAudio.messageElement.classList.remove('processing'); + } + + console.log('Audio response finalized and ready for playback'); + } catch (e) { + console.error('Error finalizing streaming audio:', e); + } + + // Reset streaming audio state + streamingAudio.chunks = []; + streamingAudio.totalChunks = 0; + streamingAudio.receivedChunks = 0; + streamingAudio.messageElement = null; + streamingAudio.audioElement = null; +} + +// Add CSS styles for new UI elements +document.addEventListener('DOMContentLoaded', function() { + // Add styles for processing state + const style = document.createElement('style'); + style.textContent = ` + .message.processing { + opacity: 0.8; + } + + .loading-indicator { + display: flex; + align-items: center; + margin-top: 8px; + font-size: 0.9em; + color: #666; + } + + .loading-spinner { + width: 16px; + height: 16px; + border: 2px solid #ddd; + border-top: 2px solid var(--primary-color); + border-radius: 50%; + margin-right: 8px; + animation: spin 1s linear infinite; + } + + @keyframes spin { + 0% { transform: rotate(0deg); } + 100% { transform: rotate(360deg); } + } + `; + document.head.appendChild(style); +}); + +// Initialize the application when DOM is fully loaded +document.addEventListener('DOMContentLoaded', initializeApp); + diff --git a/Backend/watermarking.py b/Backend/watermarking.py new file mode 100644 index 0000000..093962f --- /dev/null +++ b/Backend/watermarking.py @@ -0,0 +1,79 @@ +import argparse + +import silentcipher +import torch +import torchaudio + +# This watermark key is public, it is not secure. +# If using CSM 1B in another application, use a new private key and keep it secret. +CSM_1B_GH_WATERMARK = [212, 211, 146, 56, 201] + + +def cli_check_audio() -> None: + parser = argparse.ArgumentParser() + parser.add_argument("--audio_path", type=str, required=True) + args = parser.parse_args() + + check_audio_from_file(args.audio_path) + + +def load_watermarker(device: str = "cuda") -> silentcipher.server.Model: + model = silentcipher.get_model( + model_type="44.1k", + device=device, + ) + return model + + +@torch.inference_mode() +def watermark( + watermarker: silentcipher.server.Model, + audio_array: torch.Tensor, + sample_rate: int, + watermark_key: list[int], +) -> tuple[torch.Tensor, int]: + audio_array_44khz = torchaudio.functional.resample(audio_array, orig_freq=sample_rate, new_freq=44100) + encoded, _ = watermarker.encode_wav(audio_array_44khz, 44100, watermark_key, calc_sdr=False, message_sdr=36) + + output_sample_rate = min(44100, sample_rate) + encoded = torchaudio.functional.resample(encoded, orig_freq=44100, new_freq=output_sample_rate) + return encoded, output_sample_rate + + +@torch.inference_mode() +def verify( + watermarker: silentcipher.server.Model, + watermarked_audio: torch.Tensor, + sample_rate: int, + watermark_key: list[int], +) -> bool: + watermarked_audio_44khz = torchaudio.functional.resample(watermarked_audio, orig_freq=sample_rate, new_freq=44100) + result = watermarker.decode_wav(watermarked_audio_44khz, 44100, phase_shift_decoding=True) + + is_watermarked = result["status"] + if is_watermarked: + is_csm_watermarked = result["messages"][0] == watermark_key + else: + is_csm_watermarked = False + + return is_watermarked and is_csm_watermarked + + +def check_audio_from_file(audio_path: str) -> None: + watermarker = load_watermarker(device="cuda") + + audio_array, sample_rate = load_audio(audio_path) + is_watermarked = verify(watermarker, audio_array, sample_rate, CSM_1B_GH_WATERMARK) + + outcome = "Watermarked" if is_watermarked else "Not watermarked" + print(f"{outcome}: {audio_path}") + + +def load_audio(audio_path: str) -> tuple[torch.Tensor, int]: + audio_array, sample_rate = torchaudio.load(audio_path) + audio_array = audio_array.mean(dim=0) + return audio_array, int(sample_rate) + + +if __name__ == "__main__": + cli_check_audio()