backend

from flask import Flask, render_template, request, jsonify
import requests
import os
import soundfile as sf
import numpy as np
from kokoro import KPipeline
import uuid
import base64

app = Flask(__name__)

# Ollama endpoint configuration
OLLAMA_API_URL = "http://localhost:11434/api/generate"
OLLAMA_TAGS_URL = "http://localhost:11434/api/tags"

# Kokoro TTS configuration (latest version 1.0)
KOKORO_MODEL_PATH = "models/kokoro-v1.0.onnx"
KOKORO_VOICES_PATH = "models/voices-v1.0.bin"
KOKORO_LANG_CODE = "a"  # American English
VOICE = "af_bella"  # Default voice

# Initialize Kokoro TTS pipeline
try:
    tts_pipeline = KPipeline(lang_code=KOKORO_LANG_CODE)
except Exception as e:
    print(f"Error initializing Kokoro TTS: {str(e)}")
    tts_pipeline = None

# Ensure model and voice files exist
MODEL_DIR = "models"
os.makedirs(MODEL_DIR, exist_ok=True)

def download_kokoro_model():
    """Download the latest Kokoro model and voices if not present."""
    if not os.path.exists(KOKORO_MODEL_PATH):
        try:
            model_url = "https://github.com/nazdridoy/kokoro-tts/releases/download/v1.0.0/kokoro-v1.0.onnx"
            voices_url = "https://github.com/nazdridoy/kokoro-tts/releases/download/v1.0.0/voices-v1.0.bin"
            os.makedirs(MODEL_DIR, exist_ok=True)
            for url, path in [(model_url, KOKORO_MODEL_PATH), (voices_url, KOKORO_VOICES_PATH)]:
                print(f"Downloading {url} to {path}")
                response = requests.get(url)
                response.raise_for_status()
                with open(path, "wb") as f:
                    f.write(response.content)
        except Exception as e:
            print(f"Error downloading Kokoro model: {str(e)}")

download_kokoro_model()

def get_ollama_models():
    """Fetch available models from Ollama API."""
    try:
        response = requests.get(OLLAMA_TAGS_URL)
        response.raise_for_status()
        models_data = response.json().get("models", [])
        return [model["name"] for model in models_data]
    except requests.RequestException as e:
        print(f"Error fetching Ollama models: {str(e)}")
        return ["llama2"]  # Fallback model

@app.route('/')
def index():
    models = get_ollama_models()
    return render_template('index2.html', models=models)

@app.route('/process_text', methods=['POST'])
def process_text():
    data = request.json
    user_input = data.get('text')
    selected_model = data.get('model')
    system_prompt = data.get('system_prompt', "You are a helpful, friendly AI assistant.")

    if not user_input:
        return jsonify({"error": "No input provided"}), 400

    # Validate selected model
    available_models = get_ollama_models()
    if selected_model not in available_models:
        selected_model = available_models[0] if available_models else "llama2"

    # Send to Ollama with the system prompt
    try:
        ollama_response = requests.post(
            OLLAMA_API_URL,
            json={
                "model": selected_model,
                "prompt": user_input,
                "system": system_prompt,  # Add the system prompt here
                "stream": False
            }
        )
        ollama_response.raise_for_status()
        response_text = ollama_response.json().get("response", "")
    except requests.RequestException as e:
        return jsonify({"error": f"Ollama API error: {str(e)}"}), 500

    # Generate audio with Kokoro TTS
    audio_data = None
    if tts_pipeline:
        try:
            generator = tts_pipeline(response_text, voice=VOICE)
            audio_chunks = []
            for _, _, audio in generator:
                audio_chunks.append(audio)
            final_audio = np.concatenate(audio_chunks)

            # Save audio to temporary file
            temp_audio_path = f"static/output_{uuid.uuid4()}.wav"
            sf.write(temp_audio_path, final_audio, 24000)

            # Convert audio to base64 for web playback
            with open(temp_audio_path, "rb") as audio_file:
                audio_data = base64.b64encode(audio_file.read()).decode('utf-8')

            # Clean up
            os.remove(temp_audio_path)
        except Exception as e:
            return jsonify({"error": f"TTS generation error: {str(e)}"}), 500

    return jsonify({
        "text": response_text,
        "audio": audio_data
    })

if __name__ == '__main__':
    app.run(host='0.0.0.0', port=5000)