import gradio as gr
import spaces
import torch
from transformers import AutoModelForSpeechSeq2Seq, AutoProcessor, pipeline

MODEL_ID = "collabora/whisper-large-v2-hindi"

print("काव्यात्मक सृजन के लिए मॉडल और प्रोसेसर को रैम में लोड किया जा रहा है...")

# १. मॉडल और प्रोसेसर को पहले ही CPU रैम में लोड कर लें (इससे ६.१७ GB डाउनलोड और लोड होने का समय बच जाएगा)
model = AutoModelForSpeechSeq2Seq.from_pretrained(
    MODEL_ID, 
    torch_dtype=torch.float16, 
    low_cpu_mem_usage=True, 
    use_safetensors=True
)
processor = AutoProcessor.from_pretrained(MODEL_ID)

print("मॉडल और प्रोसेसर सफलतापूर्वक रैम में आ चुके हैं!")

# २. अब केवल इस गणना वाले हिस्से को ZeroGPU सौंपेंगे
@spaces.GPU(duration=120)
def transcribe_hindi(audio):
    if audio is None:
        return "कृपया ऑडियो प्रदान करें।", None
    
    try:
        # मॉडल को तुरंत GPU पर भेजें (यह १ सेकंड से कम समय लेगा क्योंकि यह पहले से रैम में है)
        model.to("cuda")
        
        # फ़ंक्शन के अंदर ही पाइपलाइन का निर्माण (ZeroGPU के नियमों के अनुसार अनिवार्य)
        asr_pipe = pipeline(
            "automatic-speech-recognition",
            model=model,
            tokenizer=processor.tokenizer,
            feature_extractor=processor.feature_extractor,
            chunk_length_s=30,
            device="cuda",
            torch_dtype=torch.float16
        )
        
        # ट्रांसक्रिप्शन प्रक्रिया
        result = asr_pipe(audio, batch_size=8, generate_kwargs={"language": "hindi"})
        text_output = result["text"].strip()
        
        # काम पूरा होते ही मॉडल को वापस CPU पर भेजें ताकि ZeroGPU का कंटेनर मुक्त हो सके
        model.to("cpu")
        
        file_path = "transcription.txt"
        with open(file_path, "w", encoding="utf-8") as f:
            f.write(text_output)
            
        return text_output, file_path

    except Exception as e:
        # किसी भी त्रुटि की स्थिति में मॉडल को वापस CPU पर लाना सुरक्षित है
        try:
            model.to("cpu")
        except:
            pass
        return f"प्रक्रिया में कुछ व्यवधान आया: {str(e)}। कृपया पुनः प्रयास करें।", None

custom_css = """
footer {visibility: hidden}
.gradio-container {background-color: #fcfcfc}
#header {text-align: center; margin-bottom: 20px}
"""

with gr.Blocks(title="IndicWhisper Collabora GPU") as demo:
    gr.HTML("<div id='header'><h1>🎙️ Hindi Whisper (Collabora Pipeline - ZeroGPU)</h1></div>")
    
    with gr.Row():
        with gr.Column():
            audio_input = gr.Audio(
                sources=["microphone", "upload"], 
                type="filepath", 
                label="ऑडियो रिकॉर्ड करें या अपलोड करें"
            )
            submit_btn = gr.Button("अनुवाद करें (Transcribe)", variant="primary")
            
        with gr.Column():
            output_text = gr.Textbox(
                label="Transcription Output", 
                lines=10, 
                placeholder="आपका टेक्स्ट यहाँ दिखाई देगा..."
            )
            download_file = gr.File(
                label="टैक्स्ट फ़ाइल डाउनलोड करें",
                visible=True
            )
            
    gr.Markdown("""
    ---
    **सुझाव:** यहाँ से प्राप्त आउटपुट को कॉपी करें या फ़ाइल डाउनलोड करके अपने **Gemini Gem** में पेस्ट करें ताकि **पञ्चमाक्षर नियमों** (ङ्, ञ्, ण्, न्, म्) के अनुसार शुद्धिकरण किया जा सके。
    """)

    submit_btn.click(
        fn=transcribe_hindi, 
        inputs=audio_input, 
        outputs=[output_text, download_file]
    )

if __name__ == "__main__":
    demo.launch(css=custom_css)