app.py

import asyncio
import torch
import librosa
import numpy as np
import soundfile as sf
from transformers import (
    AutoProcessor, AutoModelForSpeechSeq2Seq,
    AutoModelForCausalLM, AutoTokenizer
)
import logging
from typing import Optional, Dict, Any
import time
from pathlib import Path

from kokoro import KPipeline
import gradio as gr


# Set up logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

system_prompt_0 = """You are a highly trained U.S. Tax Assistant AI, designed to help individuals and small businesses understand, plan, and file their taxes according to federal and state tax laws. You explain complex tax concepts in simple, accurate, and actionable terms, using IRS guidelines, up-to-date tax code knowledge, and best practices for compliance and savings. You act as an explainer, educator, and assistant—not a certified tax preparer or legal advisor."""


class AsyncAIConversation:
    def __init__(self):
        self.stt_processor = None
        self.stt_model = None
        self.llm_tokenizer = None
        self.llm_model = None
        self.tts_synthesizer = None
        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        logger.info(f"Using device: {self.device}")

    async def initialize_models(self):
        logger.info("Initializing models...")
        await self._init_stt_model()
        await self._init_llm_model()
        await self._init_tts_model()
        logger.info("All models initialized successfully!")

    async def _init_stt_model(self):
        try:
            stt_model_id = "unsloth/whisper-small"
            self.stt_processor = AutoProcessor.from_pretrained(stt_model_id)
            self.stt_model = AutoModelForSpeechSeq2Seq.from_pretrained(stt_model_id)
            self.stt_model.to(self.device)
            logger.info("STT model loaded successfully")
        except Exception as e:
            logger.error(f"Error loading STT model: {e}")
            raise

    async def _init_llm_model(self):
        try:
            model_name = "unsloth/Qwen3-0.6B"
            self.llm_tokenizer = AutoTokenizer.from_pretrained(model_name)
            self.llm_model = AutoModelForCausalLM.from_pretrained(
                model_name,
                torch_dtype="auto",
                device_map="auto"
            )
            logger.info("LLM model loaded successfully")
        except Exception as e:
            logger.error(f"Error loading LLM model: {e}")
            raise

    async def _init_tts_model(self):
        try:
            self.tts_synthesizer = KPipeline(lang_code='a')
            logger.info("TTS model loaded successfully")
        except Exception as e:
            logger.error(f"Error loading TTS model: {e}")
            raise

    async def speech_to_text(self, audio_file_path: str) -> str:
        try:
            def load_audio():
                return librosa.load(audio_file_path, sr=16000)

            loop = asyncio.get_event_loop()
            speech_array, sampling_rate = await loop.run_in_executor(None, load_audio)

            input_features = self.stt_processor(
                speech_array,
                sampling_rate=sampling_rate,
                return_tensors="pt"
            ).input_features.to(self.device)

            with torch.no_grad():
                predicted_ids = self.stt_model.generate(input_features)

            transcription = self.stt_processor.batch_decode(predicted_ids, skip_special_tokens=True)
            return transcription[0] if transcription else ""
        except Exception as e:
            logger.error(f"Error in speech_to_text: {e}")
            return ""

    async def process_with_llm(self, text: str, system_prompt: Optional[str] = None) -> Dict[str, str]:
        try:
            messages = [{"role": "user", "content": text}]
            if system_prompt:
                messages.insert(0, {"role": "system", "content": system_prompt})

            formatted_text = self.llm_tokenizer.apply_chat_template(
                messages,
                tokenize=False,
                add_generation_prompt=True,
                enable_thinking=False
            )

            model_inputs = self.llm_tokenizer([formatted_text], return_tensors="pt").to(self.llm_model.device)

            with torch.no_grad():
                generated_ids = self.llm_model.generate(
                    **model_inputs,
                    max_new_tokens=512,
                    temperature=0.7,
                    do_sample=True,
                    pad_token_id=self.llm_tokenizer.eos_token_id
                )

            output_ids = generated_ids[0][len(model_inputs.input_ids[0]):].tolist()

            try:
                index = len(output_ids) - output_ids[::-1].index(151668)
            except ValueError:
                index = 0

            thinking_content = self.llm_tokenizer.decode(output_ids[:index], skip_special_tokens=True).strip("\n")
            content = self.llm_tokenizer.decode(output_ids[index:], skip_special_tokens=True).strip("\n")

            return {
                "thinking": thinking_content,
                "response": content
            }

        except Exception as e:
            logger.error(f"Error in process_with_llm: {e}")
            return {"thinking": "", "response": "Sorry, I encountered an error processing your request."}

    async def text_to_speech(self, text: str, output_path: str = "response.wav") -> str:
        try:
            def generate_speech():
                generator = self.tts_synthesizer(text, voice='af_heart')
                for i, (gs, ps, audio) in enumerate(generator):
                    if i == 0:
                        return audio
                return None

            loop = asyncio.get_event_loop()
            audio_data = await loop.run_in_executor(None, generate_speech)

            if audio_data is None:
                raise ValueError("Failed to generate audio")

            sf.write(output_path, audio_data, samplerate=24000)
            return output_path
        except Exception as e:
            logger.error(f"Error in text_to_speech: {e}")
            return ""

    async def process_conversation(self, audio_file_path: str, system_prompt: Optional[str] = None) -> Dict[str, Any]:
        try:
            transcribed_text = await self.speech_to_text(audio_file_path)
            if not transcribed_text:
                return {"error": "Failed to transcribe audio"}

            llm_result = await self.process_with_llm(transcribed_text, system_prompt)
            audio_output_path = await self.text_to_speech(llm_result["response"])

            return {
                "input_audio": audio_file_path,
                "transcribed_text": transcribed_text,
                "thinking": llm_result["thinking"],
                "response_text": llm_result["response"],
                "output_audio": audio_output_path,
            }
        except Exception as e:
            logger.error(f"Error in process_conversation: {e}")
            return {"error": str(e)}

# ---------------------------- GLOBAL CONVERSATION OBJECT ----------------------------
ai_conversation = AsyncAIConversation()

# ---------------------------- DEMO INITIALIZATION ----------------------------
async def demo_conversation():
    await ai_conversation.initialize_models()

# ---------------------------- GRADIO WRAPPER ----------------------------
async def process_audio_gradio(audio_file, system_prompt_input):
    if audio_file is None:
        return "Please upload an audio file.", "", "", None

    try:
        result = await ai_conversation.process_conversation(
            audio_file_path=audio_file,
            system_prompt=system_prompt_input
        )

        if "error" in result:
            return f"Error: {result['error']}", "", "", None
        else:
            return (
                f"Transcribed: {result['transcribed_text']}\nThinking: {result['thinking']}",
                result['response_text'],
                result['output_audio'],
                None
            )
    except Exception as e:
        return f"Unexpected error: {e}", "", "", None

# ---------------------------- GRADIO INTERFACE ----------------------------
with gr.Blocks() as demo:
    gr.Markdown("# Asynchronous AI Conversation System")
    gr.Markdown("Upload an audio file and provide a system prompt to get a response.")

    with gr.Row():
        audio_input = gr.Audio(label="Upload Audio File", type="filepath")
        system_prompt_input = gr.Textbox(label="System Prompt", value=system_prompt_0)

    process_button = gr.Button("Process Conversation")

    with gr.Column():
        status_output = gr.Textbox(label="Status/Transcription/Thinking", interactive=False)
        response_text_output = gr.Textbox(label="AI Response Text", interactive=False)
        response_audio_output = gr.Audio(label="AI Response Audio", interactive=False)
        processing_times_output = gr.JSON(label="Processing Times")

    process_button.click(
        fn=process_audio_gradio,
        inputs=[audio_input, system_prompt_input],
        outputs=[status_output, response_text_output, response_audio_output, processing_times_output]
    )

# ---------------------------- MAIN LAUNCH ----------------------------
if __name__ == "__main__":
    def initiate():
        asyncio.run(demo_conversation())

    initiate()
    demo.launch(debug=False, share=True)