#!/usr/bin/env python3
# ------------------------------------------------------------
# Whisper Small LS-100  →  Text  →  Image (CPU optimized)
# Adds ⚡ Fast / 🎨 Quality mode for Hugging Face Spaces
# Author: Murali Chandran (codeninja3d)
# ------------------------------------------------------------

import os, gc, torch, gradio as gr
try:
    import torchaudio
except Exception:
    torchaudio = None
import soundfile as sf

from transformers import WhisperProcessor, WhisperForConditionalGeneration
from diffusers import AutoPipelineForText2Image

# ------------------------------------------------------------
# Config
# ------------------------------------------------------------
DEVICE = "cpu"
ASR_MODEL_ID = "codeninja3d/whisper-small-ls100"

# ------------------------------------------------------------
# Helpers
# ------------------------------------------------------------
def load_audio(audio, target_sr=16000):
    """Load and normalize mono waveform (CPU-safe)."""
    if isinstance(audio, tuple):
        sr, arr = audio
        wav = torch.tensor(arr, dtype=torch.float32)
        if wav.ndim > 1:
            wav = wav.mean(dim=-1, keepdim=True).T
    else:
        if torchaudio:
            try:
                wav, sr = torchaudio.load(audio)
            except Exception:
                data, sr = sf.read(audio)
                wav = torch.tensor(data, dtype=torch.float32)
        else:
            data, sr = sf.read(audio)
            wav = torch.tensor(data, dtype=torch.float32)
    if wav.ndim > 1:
        wav = wav.mean(dim=0, keepdim=True)
    if sr != target_sr and torchaudio:
        wav = torchaudio.functional.resample(wav, sr, target_sr)
    wav = wav / wav.abs().max()
    return wav.squeeze(0)

# ------------------------------------------------------------
# Load Whisper
# ------------------------------------------------------------
processor = WhisperProcessor.from_pretrained(ASR_MODEL_ID)
asr_model = WhisperForConditionalGeneration.from_pretrained(ASR_MODEL_ID).to(DEVICE).eval()

# ------------------------------------------------------------
# Load both diffusion models lazily (to save RAM)
# ------------------------------------------------------------
pipe_cache = {}

def get_pipeline(model_id):
    if model_id not in pipe_cache:
        pipe = AutoPipelineForText2Image.from_pretrained(model_id, torch_dtype=torch.float32).to(DEVICE)
        pipe.set_progress_bar_config(disable=True)
        pipe_cache[model_id] = pipe
    return pipe_cache[model_id]

# ------------------------------------------------------------
# Main logic
# ------------------------------------------------------------
def transcribe_and_generate(audio, mode):
    if audio is None:
        return "⚠️ Please record or upload an audio clip.", None

    # --- Transcribe ---
    wav = load_audio(audio)
    inputs = processor.feature_extractor(wav.numpy(), sampling_rate=16000, return_tensors="pt").to(DEVICE)
    forced_ids = processor.get_decoder_prompt_ids(language="en", task="transcribe")

    with torch.no_grad():
        ids = asr_model.generate(**inputs, forced_decoder_ids=forced_ids, max_new_tokens=192, num_beams=3)
    text = processor.tokenizer.batch_decode(ids, skip_special_tokens=True)[0].strip()

    # --- Image generation ---
    if mode == "Fast":
        model_id = "stabilityai/sd-turbo"
        width = height = 384
        steps, cfg = 4, 1.0
    else:
        model_id = "runwayml/stable-diffusion-v1-5"
        width = height = 512
        steps, cfg = 15, 7.5

    prompt = f"photorealistic, soft light, {text}, detailed"
    gc.collect()
    pipe = get_pipeline(model_id)

    try:
        with torch.inference_mode():
            image = pipe(prompt=prompt, num_inference_steps=steps,
                         guidance_scale=cfg, width=width, height=height).images[0]
        return text, image
    except Exception as e:
        return f"{text}\\n\\n❗ Image generation failed: {e}", None

# ------------------------------------------------------------
# Gradio UI
# ------------------------------------------------------------
with gr.Blocks(theme="soft") as demo:
    gr.Markdown("""
    ## 🎙️ Whisper → 🖼️ Image Generator  
    Choose **Fast** (⚡ quick demo) or **Quality** (🎨 detailed, slower).
    """)

    with gr.Row():
        audio_in = gr.Audio(type="numpy", label="🎧 Record or Upload Audio")
        mode = gr.Radio(["Fast", "Quality"], value="Fast", label="Mode")

    text_out = gr.Textbox(label="📄 Transcription")
    img_out = gr.Image(label="🖼️ Generated Image")
    run_btn = gr.Button("Transcribe + Generate", variant="primary")

    run_btn.click(
        fn=transcribe_and_generate,
        inputs=[audio_in, mode],
        outputs=[text_out, img_out],
    )

demo.launch(server_name="0.0.0.0", server_port=7860)