gradio_app.py

"""
Mindfull AI Avatar Chatbot - Gradio Frontend
Works with any uploaded avatar image.
Runs locally (Ollama) or on HuggingFace Spaces (HF Inference API).
"""

import os
import sys
import time
import tempfile
import shutil
from pathlib import Path

import gradio as gr

# ── runtime paths ──────────────────────────────────────────────────────────────
ROOT = Path(__file__).parent.absolute()
OUTPUT_DIR   = ROOT / "outputs"
AUDIO_DIR    = OUTPUT_DIR / "audio"
VIDEO_DIR    = OUTPUT_DIR / "video"
TEMP_DIR     = OUTPUT_DIR / "temp"
for d in [AUDIO_DIR, VIDEO_DIR, TEMP_DIR]:
    d.mkdir(parents=True, exist_ok=True)

# ── LLM backend detection ─────────────────────────────────────────────────────
def _ollama_available() -> bool:
    try:
        import requests
        r = requests.get("http://localhost:11434/api/tags", timeout=3)
        return r.status_code == 200
    except Exception:
        return False

OLLAMA_URL   = "http://localhost:11434/api/generate"
OLLAMA_MODEL = os.environ.get("OLLAMA_MODEL", "mindfull")

HF_TOKEN  = os.environ.get("HF_TOKEN", "")
HF_MODEL  = "mistralai/Mistral-7B-Instruct-v0.3"


# ── text generation ────────────────────────────────────────────────────────────
def generate_text_ollama(prompt: str) -> str:
    import requests
    payload = {
        "model": OLLAMA_MODEL,
        "prompt": prompt,
        "stream": False,
        "options": {"temperature": 0.7, "num_ctx": 2048},
    }
    r = requests.post(OLLAMA_URL, json=payload, timeout=60)
    r.raise_for_status()
    return r.json().get("response", "").strip()


def generate_text_hf(prompt: str) -> str:
    from huggingface_hub import InferenceClient
    client = InferenceClient(model=HF_MODEL, token=HF_TOKEN or None)
    out = client.text_generation(
        prompt,
        max_new_tokens=256,
        temperature=0.7,
        do_sample=True,
    )
    return out.strip()


def generate_response(user_input: str, history: list) -> str:
    """Build prompt from history and call whichever backend is available."""
    # History is Gradio 5 messages format: [{role, content}, ...]
    ctx_lines = []
    for h in history[-8:]:  # last 4 exchanges
        role    = h.get("role", "") if isinstance(h, dict) else ""
        content = h.get("content", "") if isinstance(h, dict) else ""
        if role == "user":
            ctx_lines.append(f"User: {content}")
        elif role == "assistant":
            ctx_lines.append(f"Mindfull: {content}")
    ctx = "\n".join(ctx_lines)

    prompt = f"""You are Mindfull, a professional AI wellness companion.
Provide warm, concise support (2–3 sentences) suitable for a voice response.

{ctx}
User: {user_input}
Mindfull:"""

    try:
        if _ollama_available():
            return generate_text_ollama(prompt)
        else:
            return generate_text_hf(prompt)
    except Exception as e:
        return f"I'm here to support you. (Error reaching AI model: {e})"


# ── audio generation ───────────────────────────────────────────────────────────
def generate_audio(text: str) -> str | None:
    """Generate speech with edge-tts. Returns path to .mp3 or None."""
    try:
        from simple_audio_gen import generate_audio_simple
        ts = int(time.time() * 1000)
        out = str(AUDIO_DIR / f"response_{ts}.mp3")
        ok  = generate_audio_simple(text, out)
        return out if ok else None
    except Exception as e:
        print(f"Audio generation error: {e}")
        return None


# ── video generation ───────────────────────────────────────────────────────────
def generate_video(image_path: str | None, audio_path: str | None) -> str | None:
    """Generate talking-head video. Returns path to .mp4 or None."""
    if not image_path or not audio_path:
        return None
    if not Path(image_path).exists() or not Path(audio_path).exists():
        return None
    try:
        from simple_video_gen import generate_video_simple
        ts  = int(time.time() * 1000)
        out = str(VIDEO_DIR / f"avatar_{ts}.mp4")
        ok  = generate_video_simple(image_path, audio_path, out)
        return out if ok and Path(out).exists() else None
    except Exception as e:
        print(f"Video generation error: {e}")
        return None


# ── main chat handler ──────────────────────────────────────────────────────────
def chat(
    user_message: str,
    avatar_image,          # numpy array or file path from gr.Image
    history: list,
    skip_video: bool,
):
    if not user_message.strip():
        yield history, None, None, "Please type a message."
        return

    # ── 1. text ──────────────────────────────────────────────────────────────
    yield history, None, None, "💬 Generating response…"
    response_text = generate_response(user_message.strip(), history)
    # Gradio 5 messages format: {"role": ..., "content": ...}
    new_history = history + [
        {"role": "user",      "content": user_message},
        {"role": "assistant", "content": response_text},
    ]

    # ── 2. audio ─────────────────────────────────────────────────────────────
    yield new_history, None, None, "🔊 Synthesising speech…"
    audio_path = generate_audio(response_text)

    # ── 3. video ─────────────────────────────────────────────────────────────
    video_path = None
    if not skip_video and avatar_image is not None:
        yield new_history, audio_path, None, "🎬 Generating avatar video…"

        # Gradio Image component returns a numpy array; save it to a temp file
        import numpy as np
        from PIL import Image as PilImage

        if isinstance(avatar_image, np.ndarray):
            ts       = int(time.time() * 1000)
            img_path = str(TEMP_DIR / f"avatar_input_{ts}.png")
            pil_img  = PilImage.fromarray(avatar_image)
            # SadTalker works best with images ≤ 512px on the longest side
            MAX_DIM = 512
            if max(pil_img.size) > MAX_DIM:
                pil_img.thumbnail((MAX_DIM, MAX_DIM), PilImage.LANCZOS)
            pil_img.save(img_path)
        else:
            img_path = str(avatar_image)  # already a path

        video_path = generate_video(img_path, audio_path)

    status = "✅ Done!" if video_path or audio_path else "⚠️ Only text response available."
    yield new_history, audio_path, video_path, status


# ── Gradio UI ──────────────────────────────────────────────────────────────────
with gr.Blocks(
    title="Mindfull – AI Wellness Companion",
    theme=gr.themes.Soft(primary_hue="blue", secondary_hue="slate"),
    css="""
        #header { text-align: center; padding: 10px 0 4px; }
        #header h1 { font-size: 2rem; margin-bottom: 2px; }
        #header p  { color: #6b7280; margin: 0; }
        .status-box { font-size: 0.85rem; color: #374151; }
    """,
) as demo:

    # ── header ────────────────────────────────────────────────────────────────
    with gr.Row(elem_id="header"):
        with gr.Column():
            gr.HTML("""
                <h1>🧠 Mindfull</h1>
                <p>AI Wellness Companion · talking-head avatar · real TTS</p>
            """)

    # If running on HF Spaces without SadTalker checkpoints, default to audio-only
    _video_disabled = os.environ.get("DISABLE_VIDEO", "0") == "1"

    # ── main layout ───────────────────────────────────────────────────────────
    with gr.Row():

        # Left column – avatar + controls
        with gr.Column(scale=1, min_width=260):
            avatar_img = gr.Image(
                label="Avatar Image (any portrait)",
                type="numpy",
                sources=["upload", "webcam"],
                height=280,
            )
            skip_video_chk = gr.Checkbox(
                label="Skip video (faster · audio only)",
                value=_video_disabled,
            )
            gr.Markdown(
                "_Upload **any** face image. "
                "Leaving it empty will return text + audio only._",
                elem_classes=["status-box"],
            )

        # Right column – chat
        with gr.Column(scale=2):
            chatbot = gr.Chatbot(
                label="Conversation",
                height=400,
                type="messages",
                avatar_images=(None, str(ROOT / "avatar_assets" / "officer.png")
                               if (ROOT / "avatar_assets" / "officer.png").exists()
                               else None),
            )
            with gr.Row():
                msg_box = gr.Textbox(
                    placeholder="How are you feeling today?",
                    label="Your message",
                    scale=5,
                    lines=1,
                )
                send_btn = gr.Button("Send", variant="primary", scale=1)

    # ── outputs ───────────────────────────────────────────────────────────────
    with gr.Row():
        audio_out = gr.Audio(label="🔊 Voice Response", type="filepath", autoplay=True)
        video_out = gr.Video(label="🎬 Avatar Video",   autoplay=True)

    status_box = gr.Textbox(
        label="Status", interactive=False, lines=1, elem_classes=["status-box"]
    )

    # ── clear button ──────────────────────────────────────────────────────────
    with gr.Row():
        clear_btn = gr.Button("🗑 Clear conversation", variant="secondary")

    # ── wiring ────────────────────────────────────────────────────────────────
    def _submit(msg, img, hist, skip):
        yield from chat(msg, img, hist, skip)

    send_btn.click(
        fn=_submit,
        inputs=[msg_box, avatar_img, chatbot, skip_video_chk],
        outputs=[chatbot, audio_out, video_out, status_box],
    ).then(fn=lambda: "", outputs=msg_box)

    msg_box.submit(
        fn=_submit,
        inputs=[msg_box, avatar_img, chatbot, skip_video_chk],
        outputs=[chatbot, audio_out, video_out, status_box],
    ).then(fn=lambda: "", outputs=msg_box)

    clear_btn.click(
        fn=lambda: ([], None, None, ""),  # empty list = empty messages history
        outputs=[chatbot, audio_out, video_out, status_box],
    )


if __name__ == "__main__":
    demo.queue().launch(
        server_name="0.0.0.0",
        server_port=7860,
        share=False,          # set True for a public Gradio link
        show_error=True,
    )