import os
import gradio as gr
import torch
import torchaudio
from huggingface_hub import login
from transformers import AutoModel, AutoProcessor, AutoModelForCausalLM

# Step 1: Secret se token nikal kar Hugging Face me login karna
hf_token = os.environ.get("HF_TOKEN")
if hf_token:
    login(token=hf_token)
    print("Successfully logged in with HF_TOKEN.")
else:
    print("Warning: HF_TOKEN secret not found. Model might fail to load.")

model_id = "LiquidAI/LFM2.5-Audio-1.5B"

# Step 2: Model aur Processor load karna
print("Loading processor...")
processor = AutoProcessor.from_pretrained(model_id, trust_remote_code=True)

print("Loading model...")
try:
    # Pehle normal tarike se load karne ki koshish karega
    model = AutoModel.from_pretrained(model_id, trust_remote_code=True, torch_dtype=torch.float32)
except Exception as e:
    print(f"AutoModel failed: {e}. Trying AutoModelForCausalLM...")
    # Agar model_type error aaya, toh CausalLM class se load karega (Liquid models aksar isme load hote hain)
    model = AutoModelForCausalLM.from_pretrained(model_id, trust_remote_code=True, torch_dtype=torch.float32)

print("Model loaded successfully!")

# Step 3: Inference Function
def process_audio(audio_path):
    if audio_path is None:
        return "Please upload an audio file."
    
    try:
        waveform, sample_rate = torchaudio.load(audio_path)
        
        # Audio preprocessing
        inputs = processor(audio=waveform, sampling_rate=sample_rate, return_tensors="pt")
        
        # Output generate karna
        with torch.no_grad():
            outputs = model(**inputs)
            
        result = str(outputs)
        return "Process Complete!\n\n" + result[:800]
        
    except Exception as e:
        return f"Error during processing: {str(e)}"

# Step 4: Gradio Interface
interface = gr.Interface(
    fn=process_audio,
    inputs=gr.Audio(type="filepath", label="Upload Audio"),
    outputs=gr.Textbox(label="Model Output"),
    title="LiquidAI Audio App 🚀",
    description="Testing Liquid LFM2.5 Audio Model on Free Tier."
)

if __name__ == "__main__":
    interface.launch()