import gradio as gr
import transformers
import librosa
import torch
import numpy as np

# Load the Shuka model pipeline.
pipe = transformers.pipeline(
    model="sarvamai/shuka_v1",
    trust_remote_code=True,
    device=0 if torch.cuda.is_available() else -1,
    torch_dtype=torch.bfloat16 if torch.cuda.is_available() else None
)

def process_audio(audio):
    """
    Processes the input audio and returns a text response generated by the Shuka model.
    """
    if audio is None:
        return "No audio provided. Please upload or record an audio file."
    
    try:
        # Gradio returns a tuple: (sample_rate, numpy_array)
        sample_rate, audio_data = audio
    except Exception as e:
        return f"Error processing audio input: {e}"
    
    if audio_data is None or len(audio_data) == 0:
        return "Audio data is empty. Please try again with a valid audio file."
    
    # Convert audio data to float if not already floating-point.
    if not np.issubdtype(audio_data.dtype, np.floating):
        audio_data = audio_data.astype(np.float32)
    
    # Resample to 16000 Hz if necessary.
    if sample_rate != 16000:
        try:
            audio_data = librosa.resample(audio_data, orig_sr=sample_rate, target_sr=16000)
            sample_rate = 16000
        except Exception as e:
            return f"Error during resampling: {e}"
    
    # Define conversation turns for the model.
    turns = [
        {'role': 'system', 'content': 'Respond naturally and informatively.'},
        {'role': 'user', 'content': '<|audio|>'}
    ]
    
    try:
        result = pipe({'audio': audio_data, 'turns': turns, 'sampling_rate': sample_rate}, max_new_tokens=512)
    except Exception as e:
        return f"Error during model processing: {e}"
    
    # Extract the generated text response.
    if isinstance(result, list) and len(result) > 0:
        response = result[0].get('generated_text', '')
    else:
        response = str(result)
    
    return response

# Create the Gradio interface.
iface = gr.Interface(
    fn=process_audio,
    inputs=gr.Audio(type="numpy"),  # File upload for audio.
    outputs="text",
    title="Sarvam AI Shuka Voice Demo",
    description="Upload an audio file and get a response using Sarvam AI's Shuka model."
)

if __name__ == "__main__":
    # Set share=True to create a public link, and specify a server port.
    iface.launch(share=True, server_port=7861)