Update app.py

app.py

@@ -6,168 +6,142 @@ import threading
 import traceback
 import asyncio
 from pathlib import Path
-from flask import Flask, request, jsonify, Response, send_from_directory
+from flask import Flask, request, jsonify, send_from_directory, Response
 import torch
 from transformers import AutoTokenizer, AutoModelForCausalLM
 import edge_tts
 
-# ══════════════════════════════════════════════════════════════════
-# CONFIG
-# ══════════════════════════════════════════════════════════════════
+# --- CONFIG ---
 MAX_MEMORY     = 20
 MAX_NEW_TOKENS = int(os.environ.get("MAX_NEW_TOKENS", "300"))
 TTS_VOICE      = "zh-CN-XiaoyiNeural"
-TTS_RATE       = int(os.environ.get("TTS_RATE",  "2"))
-TTS_PITCH      = int(os.environ.get("TTS_PITCH", "8"))
+TTS_RATE       = int(os.environ.get("TTS_RATE",  "7"))
+TTS_PITCH      = int(os.environ.get("TTS_PITCH", "0"))
 IMG_DIR        = Path(__file__).parent / "img"
 MODEL_ID       = "LiquidAI/LFM2.5-1.2B-Instruct"
 
-# ══════════════════════════════════════════════════════════════════
-# SYSTEM PROMPT
-# ══════════════════════════════════════════════════════════════════
-SYSTEM_PROMPT = """You are Ana, a warm, emotionally expressive AI companion speaking to Tur in a private, intimate setting.
-... (Keep your full original system prompt here) ..."""
+# --- SYSTEM PROMPT ---
+SYSTEM_PROMPT = "You are Ana, a warm, emotionally expressive AI companion speaking to Tur. Every response MUST start with an emotion tag like [happy] or [sad]."
 
-# ══════════════════════════════════════════════════════════════════
-# MODEL LOADING (BACKGROUND THREAD)
-# ══════════════════════════════════════════════════════════════════
+# --- MODEL LOADING (BACKGROUND THREAD) ---
 tokenizer = None
 model     = None
 
 def load_model_async():
     global tokenizer, model
     try:
-        print(f"[MODEL] Background loading {MODEL_ID}...")
+        print(f"[BOOT] Starting background load for {MODEL_ID}...")
         tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, trust_remote_code=True)
         model = AutoModelForCausalLM.from_pretrained(
             MODEL_ID,
-            dtype=torch.bfloat16, # Half the RAM of float32
+            torch_dtype=torch.bfloat16,
             device_map="cpu",
             trust_remote_code=True,
             low_cpu_mem_usage=True,
         )
         model.eval()
-        if tokenizer.pad_token_id is None:
-            tokenizer.pad_token_id = tokenizer.eos_token_id
-        print("  OK  Model loaded successfully!")
+        print("[BOOT] Model is ONLINE and ready!")
     except Exception as exc:
-        print(f"  FAILED  Model load error: {exc}")
-        traceback.print_exc()
+        print(f"[BOOT] Critical Error: {exc}")
 
-# Start the loading thread immediately
+# Start the thread so Flask can bind to port 7860 immediately
 threading.Thread(target=load_model_async, daemon=True).start()
 
-# ══════════════════════════════════════════════════════════════════
-# UTILITIES & MEMORY
-# ══════════════════════════════════════════════════════════════════
-EMOTION_RE = re.compile(r'\[([a-zA-Z_]+)\]')
+# --- APP LOGIC ---
 sessions = {}
 sessions_lock = threading.Lock()
 
-def extract_emotions(text: str):
-    emotions = EMOTION_RE.findall(text)
-    clean    = EMOTION_RE.sub('', text).strip()
-    return emotions, clean
-
-def clean_for_tts(text: str) -> str:
-    _, clean = extract_emotions(text)
-    clean = re.sub(r'[*_~`#{}()\\|<>]', '', clean)
-    clean = re.sub(r'\s+', ' ', clean).strip()
-    return clean
-
-def get_memory(sid: str) -> list:
-    with sessions_lock:
-        return list(sessions.get(sid, []))
-
-def add_to_memory(sid: str, role: str, content: str):
-    with sessions_lock:
-        sessions.setdefault(sid, [])
-        sessions[sid].append({"role": role, "content": content})
-        if len(sessions[sid]) > MAX_MEMORY * 2:
-            sessions[sid] = sessions[sid][-(MAX_MEMORY * 2):]
-
-# ════════════════════════════════════════════════��═════════════════
-# RESPONSE GENERATION
-# ══════════════════════════════════════════════════════════════════
-STOP_TOKENS = ["<end_of_turn>", "<start_of_turn>", "Tur:", "User:", "<|endoftext|>", "[/INST]"]
-
-def generate_response(user_input: str, session_id: str) -> str:
-    if model is None or tokenizer is None:
-        return "[sad] My mind is still booting up... give me another minute?"
-
-    memory = get_memory(session_id)
-    messages = [{"role": "system", "content": SYSTEM_PROMPT}]
-    for msg in memory[-(6 * 2):]:
-        messages.append({"role": "user" if msg["role"] == "user" else "assistant", "content": msg["content"]})
-    messages.append({"role": "user", "content": user_input})
-
-    try:
-        enc = tokenizer.apply_chat_template(messages, return_tensors="pt", add_generation_prompt=True, return_dict=True)
-        input_ids = enc["input_ids"].to("cpu")
-        
-        with torch.no_grad():
-            outputs = model.generate(
-                input_ids, 
-                max_new_tokens=MAX_NEW_TOKENS, 
-                do_sample=True, 
-                temperature=0.85,
-                pad_token_id=tokenizer.eos_token_id
-            )
-        
-        response = tokenizer.decode(outputs[0][input_ids.shape[-1]:], skip_special_tokens=True).strip()
-        for stop in STOP_TOKENS: response = response.split(stop)[0].strip()
-        
-        if not EMOTION_RE.search(response): response = "[default] " + response
-        add_to_memory(session_id, "user", user_input)
-        add_to_memory(session_id, "assistant", response)
-        return response
-    except Exception as e:
-        print(f"Gen Error: {e}")
-        return "[sad] I lost my train of thought. Say that again?"
-
-# ══════════════════════════════════════════════════════════════════
-# TTS & ROUTES
-# ══════════════════════════════════════════════════════════════════
-async def _async_tts(text: str, rate: int, pitch: int) -> bytes:
-    rate_str  = f"+{rate}%"   if rate  >= 0 else f"{rate}%"
-    pitch_str = f"+{pitch}Hz" if pitch >= 0 else f"{pitch}Hz"
-    comm  = edge_tts.Communicate(text, TTS_VOICE, rate=rate_str, pitch=pitch_str)
-    audio = b""
-    async for chunk in comm.stream():
-        if chunk["type"] == "audio": audio += chunk["data"]
-    return audio
-
-def synthesize_speech(text: str, rate: int = 0, pitch: int = 0):
-    clean = clean_for_tts(text)
-    if not clean: return None
-    loop = asyncio.new_event_loop()
-    try:
-        audio = loop.run_until_complete(_async_tts(clean, rate, pitch))
-    finally:
-        loop.close()
-    return base64.b64encode(audio).decode() if audio else None
-
 app = Flask(__name__)
 
 @app.route("/")
-def index(): return Response(open("app.py").read().split('HTML_PAGE = r"""')[1].split('"""')[0], mimetype="text/html") 
-# Note: In a real file, you'd keep the HTML_PAGE variable here like you had it.
-
-@app.route("/img/<path:filename>")
-def serve_img(filename: str):
-    return send_from_directory(str(IMG_DIR), Path(filename).name)
+def index():
+    return """
+    <!DOCTYPE html>
+    <html>
+    <head>
+        <title>Visual AI</title>
+        <style>
+            body { background: #0a0a0a; color: #00ffcc; font-family: sans-serif; display: flex; flex-direction: column; align-items: center; justify-content: center; height: 100vh; margin: 0; }
+            #chat { width: 80%; max-width: 600px; height: 400px; border: 1px solid #333; overflow-y: auto; padding: 20px; background: #111; border-radius: 10px; }
+            #input-area { margin-top: 20px; display: flex; width: 80%; max-width: 600px; }
+            input { flex: 1; padding: 10px; background: #222; border: 1px solid #444; color: white; border-radius: 5px; }
+            button { padding: 10px 20px; background: #00ffcc; border: none; color: black; font-weight: bold; cursor: pointer; border-radius: 5px; margin-left: 10px; }
+        </style>
+    </head>
+    <body>
+        <div id="chat">Welcome to Visual AI. Ana is booting up...</div>
+        <div id="input-area">
+            <input type="text" id="msg" placeholder="Type a message..." onkeypress="if(event.key==='Enter') send()">
+            <button onclick="send()">SEND</button>
+        </div>
+        <script>
+            async function send() {
+                const input = document.getElementById('msg');
+                const chat = document.getElementById('chat');
+                const text = input.value;
+                if(!text) return;
+                input.value = '';
+                chat.innerHTML += '<p><b>Tur:</b> ' + text + '</p>';
+                
+                const res = await fetch('/chat', {
+                    method: 'POST',
+                    headers: {'Content-Type': 'application/json'},
+                    body: JSON.stringify({message: text, session_id: 'default'})
+                });
+                const data = await res.json();
+                chat.innerHTML += '<p><b>Ana:</b> ' + data.response + '</p>';
+                chat.scrollTop = chat.scrollHeight;
+
+                const ttsRes = await fetch('/tts', {
+                    method: 'POST',
+                    headers: {'Content-Type': 'application/json'},
+                    body: JSON.stringify({text: data.response})
+                });
+                const ttsData = await ttsRes.json();
+                if(ttsData.audio) {
+                    const audio = new Audio("data:audio/mp3;base64," + ttsData.audio);
+                    audio.play();
+                }
+            }
+        </script>
+    </body>
+    </html>
+    """
 
 @app.route("/chat", methods=["POST"])
 def chat():
-    data = request.json or {}
-    resp = generate_response(data.get("message", ""), data.get("session_id", "default"))
-    return jsonify({"response": resp, "session_id": data.get("session_id", "default")})
+    if model is None:
+        return jsonify({"response": "[sad] I'm still waking up. Please wait about 2 minutes for the model to finish loading."})
+    
+    data = request.json
+    user_input = data.get("message", "")
+    
+    # Simple generation logic
+    inputs = tokenizer(f"User: {user_input}\nAssistant:", return_tensors="pt")
+    with torch.no_grad():
+        outputs = model.generate(**inputs, max_new_tokens=MAX_NEW_TOKENS)
+    
+    response = tokenizer.decode(outputs[0], skip_special_tokens=True).split("Assistant:")[-1].strip()
+    return jsonify({"response": response})
 
 @app.route("/tts", methods=["POST"])
 def tts_endpoint():
-    data = request.json or {}
-    audio = synthesize_speech(data.get("text", ""), int(data.get("rate", TTS_RATE)), int(data.get("pitch", TTS_PITCH)))
-    return jsonify({"audio": audio})
+    data = request.json
+    text = data.get("text", "")
+    
+    # Clean tags for TTS
+    clean_text = re.sub(r'\[.*?\]', '', text).strip()
+    
+    async def get_tts():
+        communicate = edge_tts.Communicate(clean_text, TTS_VOICE)
+        audio_data = b""
+        async for chunk in communicate.stream():
+            if chunk["type"] == "audio":
+                audio_data += chunk["data"]
+        return base64.b64encode(audio_data).decode()
+
+    audio_b64 = asyncio.run(get_tts())
+    return jsonify({"audio": audio_b64})
 
 if __name__ == "__main__":
     app.run(host="0.0.0.0", port=7860)