Update app.py

app.py

@@ -1,253 +1,114 @@
-# app.py
-import os
-import uuid
-import tempfile
-from pathlib import Path
-from io import BytesIO
-import numpy as np
-from PIL import Image, ImageDraw, ImageOps
 import gradio as gr
-import moviepy.editor as mpy
-from pydub import AudioSegment
-import soundfile as sf
-import math
-import random
-
-# Try faster_whisper first (optional on Spaces), fallback to whisper
-try:
-    from faster_whisper import WhisperModel
-    WHISPER_AVAILABLE = True
-except Exception:
-    WHISPER_AVAILABLE = False
-    import whisper
-
-# Transformers LLM (Flan-T5 small)
-from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, pipeline
-
-# gTTS for multilingual TTS
+import torch
+import cv2
+import numpy as np
 from gtts import gTTS
+from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, WhisperProcessor, WhisperForConditionalGeneration
+from PIL import Image
+import ffmpeg
+import tempfile
+import os
+
+# -----------------------
+# Load Models
+# -----------------------
+
+device = "cpu"
 
-# ---------- Model setup (load once) ----------
-MODEL_NAME = "google/flan-t5-small"
-tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
-model = AutoModelForSeq2SeqLM.from_pretrained(MODEL_NAME)
-llm = pipeline("text2text-generation", model=model, tokenizer=tokenizer)
-
-# Whisper
-if WHISPER_AVAILABLE:
-    whisper_model = WhisperModel("small", device="cpu", compute_type="int8")
-else:
-    whisper_model = whisper.load_model("small")
-
-# ---------- Helpers ----------
-def transcribe_audio(fp, lang=None):
-    """Return text from audio file path fp."""
-    try:
-        if WHISPER_AVAILABLE:
-            segments, info = whisper_model.transcribe(fp, language=lang) if lang else whisper_model.transcribe(fp)
-            text = " ".join([s.text for s in segments])
-            return text
-        else:
-            res = whisper_model.transcribe(fp, language=lang) if lang else whisper_model.transcribe(fp)
-            return res["text"]
-    except Exception as e:
+# Speech-to-text (Whisper small)
+whisper_processor = WhisperProcessor.from_pretrained("openai/whisper-small")
+whisper_model = WhisperForConditionalGeneration.from_pretrained("openai/whisper-small").to(device)
+
+# Text generation (Flan-T5 small)
+tokenizer = AutoTokenizer.from_pretrained("google/flan-t5-small")
+t5_model = AutoModelForSeq2SeqLM.from_pretrained("google/flan-t5-small").to(device)
+
+# -----------------------
+# Helper Functions
+# -----------------------
+
+def transcribe(audio):
+    if audio is None:
         return ""
+    audio = whisper_processor(audio["array"], sampling_rate=16000, return_tensors="pt")
+    result = whisper_model.generate(audio["input_features"])
+    return whisper_processor.batch_decode(result, skip_special_tokens=True)[0]
+
+def reply(text):
+    inp = tokenizer(text, return_tensors="pt")
+    out = t5_model.generate(**inp, max_length=120)
+    return tokenizer.decode(out[0], skip_special_tokens=True)
+
+def synth_voice(text, path):
+    tts = gTTS(text=text, lang="en", tld="com", slow=False)
+    tts.save(path)
+    return path
+
+def animate_avatar(image, audio_path):
+    avatar = Image.open(image).convert("RGBA")
+    w, h = avatar.size
+    avatar_np = np.array(avatar)
+
+    # Extract audio amplitude → fake lip motion
+    import wave
+    with wave.open(audio_path, "rb") as wav:
+        frames = wav.readframes(-1)
+        audio_np = np.frombuffer(frames, dtype=np.int16)
+        amp = np.abs(audio_np)[::2000]  # downsample amplitude curve
+
+    frames_list = []
+    for a in amp:
+        frame = avatar_np.copy()
+        intensity = min(8, int(a / 3000))
+        frame[h - 40 : h - 20, w//2 - 20 : w//2 + 20, 3] = 255 - intensity * 20
+        frames_list.append(frame)
+
+    # Export to video
+    temp_video = tempfile.NamedTemporaryFile(suffix=".mp4", delete=False).name
+    out = cv2.VideoWriter(temp_video, cv2.VideoWriter_fourcc(*"mp4v"), 20, (w, h))
+
+    for f in frames_list:
+        out.write(cv2.cvtColor(f, cv2.COLOR_RGBA2BGR))
+    out.release()
+
+    return temp_video
+
+# -----------------------
+# Main Chat Logic
+# -----------------------
+
+def chat(image, audio, text):
+    user_input = text if text else transcribe(audio)
+    if not user_input:
+        return "Say something!", None
+
+    ai_answer = reply(user_input)
+
+    # TTS
+    temp_audio = tempfile.NamedTemporaryFile(suffix=".mp3", delete=False).name
+    synth_voice(ai_answer, temp_audio)
+
+    # Talking avatar
+    video = animate_avatar(image, temp_audio)
+
+    return ai_answer, video
+
+# -----------------------
+# Gradio UI
+# -----------------------
+
+with gr.Blocks() as interface:
+    gr.Markdown("## 🧚‍♀️ AI Avatar Companion — Free & No-Install")
+
+    avatar = gr.Image(type="filepath", label="Upload Avatar PNG")
+    audio = gr.Audio(source="microphone", type="numpy", label="Speak")
+    txt = gr.Textbox(label="Or type your message")
+
+    out_text = gr.Textbox(label="AI Response")
+    out_video = gr.Video(label="Talking Avatar")
+
+    submit = gr.Button("Talk")
+
+    submit.click(chat, inputs=[avatar, audio, txt], outputs=[out_text, out_video])
 
-def ask_llm(user_text, history=None):
-    """Simple LLM wrapper: keep replies short and tutor style."""
-    prompt = f"You are a friendly English tutor and helpful assistant. Reply concisely, give an example if appropriate, and ask one follow-up question. User said: {user_text}"
-    out = llm(prompt, max_length=256, do_sample=False)
-    return out[0]["generated_text"]
-
-def tts_gtts(text, out_path, lang="en"):
-    """gTTS synthesis (multilingual)."""
-    if not text or len(text.strip())==0:
-        # fallback silence
-        silent = AudioSegment.silent(duration=500)
-        silent.export(out_path, format="mp3")
-        return out_path
-    tts = gTTS(text=text, lang=lang, slow=False)
-    tts.save(out_path)
-    return out_path
-
-def make_envelope(audio_path, fps=25):
-    """Compute amplitude envelope per frame from audio file."""
-    seg = AudioSegment.from_file(audio_path)
-    samples = np.array(seg.get_array_of_samples()).astype(np.float32)
-    if seg.channels > 1:
-        samples = samples.reshape((-1, seg.channels)).mean(axis=1)
-    if samples.size == 0:
-        return np.zeros(1)
-    samples = samples / (np.max(np.abs(samples)) + 1e-9)
-    duration_s = seg.duration_seconds
-    n_frames = max(1, int(duration_s * fps))
-    parts = np.array_split(samples, n_frames)
-    env = np.array([np.sqrt(np.mean(p**2)) if p.size>0 else 0.0 for p in parts])
-    env = (env - env.min()) / (env.max()-env.min()+1e-9)
-    return env
-
-def create_talking_video(image_pil, audio_path, out_video_path, fps=25, emotion="neutral"):
-    """
-    Create an MP4 by overlaying animated 'jaw' effect on the image, synced to audio envelope.
-    Also adds simple blink and head-tilt micro gestures based on 'emotion'.
-    """
-    envelope = make_envelope(audio_path, fps=fps)
-    duration = max(0.5, len(envelope)/fps)
-    w,h = image_pil.size
-
-    # jaw box (estimate lower center)
-    jw = int(w * 0.26); jh = int(h * 0.08)
-    jx = int(w*0.5 - jw/2); jy = int(h*0.68 - jh/2)
-
-    # emotion-driven head tilt/scale patterns
-    if emotion == "happy":
-        head_tilt = lambda t: math.sin(2*math.pi*t/duration)*2.2
-    elif emotion == "thinking":
-        head_tilt = lambda t: math.sin(2*math.pi*t/duration)*-2.5
-    elif emotion == "surprised":
-        head_tilt = lambda t: math.sin(2*math.pi*t/duration)*1.8
-    else:
-        head_tilt = lambda t: math.sin(2*math.pi*t/duration)*0.8
-
-    def make_frame(t):
-        i = min(int(t*fps), len(envelope)-1)
-        level = float(envelope[i])
-        frame = image_pil.copy().convert("RGBA")
-        draw = ImageDraw.Draw(frame, 'RGBA')
-
-        # subtle head tilt via shear or rotate (simple)
-        angle = head_tilt(t)
-        frame = frame.rotate(angle, resample=Image.BICUBIC, center=(w//2, h//3), expand=False)
-
-        # mouth ellipse overlay (simulate opening)
-        mouth_h = int(jh * (1.0 + level*1.2))
-        mouth_y = int(jy + jh - mouth_h/2)
-        alpha = int(20 + level*100)
-        draw.ellipse([jx, mouth_y, jx+jw, mouth_y+mouth_h], fill=(20,20,20, alpha))
-
-        # blink: occasionally draw eyelid rectangles based on time
-        # simple periodic blink
-        if (int(t*2) % 7) == 0 and random.random()>0.65:
-            # top lid
-            draw.rectangle([0, 0, w, int(h*0.23)], fill=(245,245,255,230))
-
-        return np.asarray(frame)
-
-    clip = mpy.VideoClip(make_frame, duration=duration)
-    audio = mpy.AudioFileClip(audio_path)
-    clip = clip.set_audio(audio)
-    clip.write_videofile(out_video_path, fps=fps, codec="libx264", audio_codec="aac", verbose=False, logger=None)
-    return out_video_path
-
-def detect_language_hint(text):
-    # crude hint: check for non-ascii characters to switch language; default en
-    if any(ord(ch) > 127 for ch in text):
-        return "auto"
-    return "en"
-
-# Emotion detection: tiny heuristic
-def detect_emotion_from_text(text):
-    t = text.lower()
-    if any(w in t for w in ["love","like","happy","great","awesome","good"]):
-        return "happy"
-    if any(w in t for w in ["why","how","think","wonder","question","confused"]):
-        return "thinking"
-    if any(w in t for w in ["wow","surprise","omg","amazed","shocked"]):
-        return "surprised"
-    if any(w in t for w in ["sorry","shy","nervous","awkward"]):
-        return "shy"
-    return "neutral"
-
-# ---------- Gradio app function ----------
-def process(image, upload_audio, mic_audio, typed_text):
-    # Save working folder
-    uid = str(uuid.uuid4())[:8]
-    tmp = Path(tempfile.gettempdir()) / f"space_{uid}"
-    tmp.mkdir(parents=True, exist_ok=True)
-
-    # Ensure image
-    if image is None:
-        return None, "Please upload an avatar image (head & shoulders).", None
-    if isinstance(image, np.ndarray):
-        image_pil = Image.fromarray(image)
-    else:
-        image_pil = Image.open(image).convert("RGBA")
-
-    # Determine input audio or typed text
-    user_text = ""
-    lang_hint = "en"
-    audio_in = None
-
-    if typed_text and typed_text.strip():
-        user_text = typed_text.strip()
-        lang_hint = detect_language_hint(user_text)
-    else:
-        # prioritize mic_audio, then upload_audio
-        audio_file = None
-        if mic_audio:
-            audio_file = mic_audio
-        elif upload_audio:
-            audio_file = upload_audio
-
-        if audio_file:
-            # save
-            audio_path = tmp / "user_in.wav"
-            with open(audio_path, "wb") as f:
-                f.write(audio_file.read())
-            # transcribe
-            try:
-                user_text = transcribe_audio(str(audio_path), lang=None)
-            except Exception as e:
-                user_text = ""
-            lang_hint = detect_language_hint(user_text)
-            audio_in = str(audio_path)
-
-    if not user_text:
-        return None, "Couldn't capture text. Try typing or recording again.", None
-
-    # Ask LLM
-    reply = ask_llm(user_text)
-
-    # Emotion
-    emotion = detect_emotion_from_text(user_text + " " + reply)
-
-    # TTS: generate mp3
-    tts_path = tmp / "reply.mp3"
-    try:
-        tts_gtts(reply, str(tts_path), lang=lang_hint)
-    except Exception:
-        # fallback to english
-        tts_gtts(reply, str(tts_path), lang="en")
-
-    # Create talking clip
-    out_video = tmp / "talking.mp4"
-    create_talking_video(image_pil, str(tts_path), str(out_video), fps=25, emotion=emotion)
-
-    # Return video file, text, audio
-    return str(out_video), reply, str(tts_path)
-
-# ---------- Gradio UI ----------
-title = "No-install Avatar Companion — Hugging Face Space (Free)"
-desc = "Upload a single avatar image (head+shoulders). Speak or type. The app transcribes, replies, synthesizes voice, and auto-generates a talking clip (MP4) with mouth animation and simple gestures."
-
-demo = gr.Interface(
-    fn=process,
-    inputs=[
-        gr.Image(type="pil", label="Upload avatar (head & shoulders PNG)"),
-        gr.Audio(source="upload", type="file", label="Upload audio (optional)"),
-        gr.Audio(source="microphone", type="file", label="Record via mic (optional)"),
-        gr.Textbox(lines=2, placeholder="Or type your message (optional)", label="Type message")
-    ],
-    outputs=[
-        gr.Video(label="Talking clip (MP4)"),
-        gr.Textbox(label="Assistant reply"),
-        gr.Audio(label="Reply audio (mp3)", type="file")
-    ],
-    title=title,
-    description=desc,
-    allow_flagging="never",
-    examples=[]
-)
-
-if __name__ == "__main__":
-    demo.launch()
+interface.launch()