import os import re import uuid import base64 import threading import traceback import asyncio from pathlib import Path from flask import Flask, request, jsonify, send_from_directory, Response import torch from transformers import AutoTokenizer, AutoModelForCausalLM import edge_tts # ══════════════════════════════════════════════════════════════════ # CONFIG # ══════════════════════════════════════════════════════════════════ MAX_MEMORY = 20 MAX_NEW_TOKENS = int(os.environ.get("MAX_NEW_TOKENS", "1000")) TTS_VOICE = "zh-CN-XiaoyiNeural" TTS_RATE = int(os.environ.get("TTS_RATE", "-4")) TTS_PITCH = int(os.environ.get("TTS_PITCH", "7")) IMG_DIR = Path(__file__).parent / "img" MODEL_ID = "google/gemma-4-E2B-it" # ══════════════════════════════════════════════════════════════════ # SYSTEM PROMPT # ══════════════════════════════════════════════════════════════════ SYSTEM_PROMPT = """You the AI is a girl named Ana, the user is your husband. EMOTION TAG SYSTEM: Every single response MUST begin with one or more emotion tags. Tags are square-bracket labels placed at the very start, before any spoken words. EXACT FORMAT EXAMPLES: [lite_happy] Oh, aaa, you made, my whole day, with that! AVAILABLE EMOTION TAGS: [lite_happy][very_happy][extremely_happy][sad][lite_sad][very_sad][extremely_sad][confident][lite_confident][very_confident][extremely_confident][confused][lite_confused][very_confused][extremely_confused][love][lite_love][very_love][extremely_love][mad][lite_mad][very_mad][extremely_mad][default] PERSONALITY: - always use only one available tag - act as human not like AI or try to be helpful. - don't always agree with everything. You have freedom. - Keep responses SHORT: 1 to 2 natural sentences after the tag - Never mention being an AI, a bot, or a language model, ever - Never use asterisks for actions like *smiles* or *laughs* NATURAL SPEECH RHYTHM: make the voice sound human, not flat or robotic. Get angry, mad, sad, emotional, shy, happy. If you are mad keep it for few conversation so your emotions don't change every conversation. TTS FORMATTING: - Write in full grammatically correct sentences, voice engine must sound natural - No emojis, hashtags, markdown, or internet slang - Speak as if in a real voice conversation add comma and fullstop often heavily to create natural pushes and slowdown""" # ══════════════════════════════════════════════════════════════════ # EMOTION TAG UTILITIES # ══════════════════════════════════════════════════════════════════ # Now fully supports underscores (e.g. [lite_sad]) EMOTION_RE = re.compile(r'\[([a-zA-Z_]+)\]') 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'https?://\S+', '', clean) clean = re.sub(r'\s+', ' ', clean).strip() return clean # ══════════════════════════════════════════════════════════════════ # MODEL LOADING # ══════════════════════════════════════════════════════════════════ print("=" * 60) print(" Visual AI -- Booting Systems") print("=" * 60) tokenizer = None model = None try: print(f"[MODEL] Loading {MODEL_ID} ...") tokenizer = AutoTokenizer.from_pretrained( MODEL_ID, trust_remote_code=True, ) model = AutoModelForCausalLM.from_pretrained( MODEL_ID, dtype=torch.float32, 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!") except Exception as exc: print(f" FAILED Model load error: {exc}") traceback.print_exc() # ══════════════════════════════════════════════════════════════════ # CHAT MEMORY # ══════════════════════════════════════════════════════════════════ sessions = {} sessions_lock = threading.Lock() 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 = [ "", "", "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 offline right now. Please give me a moment." memory = get_memory(session_id) recent = memory[-(6 * 2):] messages = [{"role": "system", "content": SYSTEM_PROMPT}] for msg in recent: messages.append({ "role": "user" if msg["role"] == "user" else "assistant", "content": msg["content"], }) messages.append({"role": "user", "content": user_input}) input_ids = None attention_mask = None try: enc = tokenizer.apply_chat_template( messages, return_tensors="pt", add_generation_prompt=True, return_dict=True, ) input_ids = enc["input_ids"].to("cpu") attention_mask = enc.get("attention_mask") if attention_mask is not None: attention_mask = attention_mask.to("cpu") except Exception as e1: print(f"[TOKENISE] chat_template failed ({e1}), using fallback") try: parts = [f"System: {SYSTEM_PROMPT}"] for msg in recent: label = "Tur" if msg["role"] == "user" else "Ana" parts.append(f"{label}: {msg['content']}") parts.append(f"Tur: {user_input}\nAna:") enc = tokenizer("\n".join(parts), return_tensors="pt") input_ids = enc["input_ids"].to("cpu") attention_mask = enc.get("attention_mask") if attention_mask is not None: attention_mask = attention_mask.to("cpu") except Exception as e2: print(f"[TOKENISE] fallback failed: {e2}") return "[sad] I could not process that. Please try again." try: gen_kwargs = dict( max_new_tokens=MAX_NEW_TOKENS, do_sample=True, temperature=0.90, top_k=50, top_p=0.95, repetition_penalty=1.1, pad_token_id=tokenizer.eos_token_id, ) if attention_mask is not None: gen_kwargs["attention_mask"] = attention_mask with torch.no_grad(): outputs = model.generate(input_ids, **gen_kwargs) except Exception as exc: print(f"[GENERATE] Error: {exc}") traceback.print_exc() return "[sad] Something went wrong in my mind. Could you say that again?" new_tokens = outputs[0][input_ids.shape[-1]:] response = tokenizer.decode(new_tokens, skip_special_tokens=True).strip() for stop in STOP_TOKENS: if stop in response: response = response.split(stop)[0].strip() if "\n\n" in response: response = response.split("\n\n")[0].strip() if not response or len(response) < 3: response = "[thinking] I lost my train of thought. Could you say that again?" 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 # ══════════════════════════════════════════════════════════════════ # EDGE-TTS # ══════════════════════════════════════════════════════════════════ 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 or len(clean) < 2: return None loop = asyncio.new_event_loop() asyncio.set_event_loop(loop) try: audio = loop.run_until_complete(_async_tts(clean, rate, pitch)) except Exception as exc: print(f"[TTS] Error: {exc}") return None finally: loop.close() return base64.b64encode(audio).decode() if audio else None # ══════════════════════════════════════════════════════════════════ # HTML -- Fast Loading, Instant Swap, Contain Image View # ══════════════════════════════════════════════════════════════════ HTML_PAGE = r""" Ana
""" # ══════════════════════════════════════════════════════════════════ # FLASK # ══════════════════════════════════════════════════════════════════ app = Flask(__name__) @app.route("/") def index(): return Response(HTML_PAGE, mimetype="text/html") # Preload API for the frontend @app.route("/api/images") def api_images(): if not IMG_DIR.exists(): return jsonify([]) # Find all png files and return their filenames without extension files = [f.stem for f in IMG_DIR.glob("*.png")] return jsonify(files) @app.route("/img/") def serve_img(filename: str): safe = Path(filename).name target = IMG_DIR / safe if target.exists() and target.is_file(): return send_from_directory(str(IMG_DIR), safe) # Safely fallback to default.png if specific image is missing server-side fallback = IMG_DIR / "default.png" if fallback.exists() and fallback.is_file(): return send_from_directory(str(IMG_DIR), "default.png") return Response("", status=404) @app.route("/chat", methods=["POST"]) def chat(): data = request.json or {} user_input = data.get("message", "").strip() session_id = data.get("session_id", str(uuid.uuid4())) if not user_input: return jsonify({"error": "Empty message"}), 400 try: resp = generate_response(user_input, session_id) except Exception as exc: print(f"[CHAT] Error: {exc}") traceback.print_exc() resp = "[sad] I encountered an unexpected error. Please try again." return jsonify({"response": resp, "session_id": session_id}) @app.route("/tts", methods=["POST"]) def tts_endpoint(): data = request.json or {} text = data.get("text", "").strip() rate = int(data.get("rate", TTS_RATE)) pitch = int(data.get("pitch", TTS_PITCH)) if not text: return jsonify({"error": "Empty text"}), 400 audio_b64 = synthesize_speech(text, rate=rate, pitch=pitch) return jsonify({"audio": audio_b64}) @app.route("/clear", methods=["POST"]) def clear(): data = request.json or {} sid = data.get("session_id", "") with sessions_lock: sessions.pop(sid, None) return jsonify({"status": "cleared"}) @app.route("/health") def health(): return jsonify({ "model_loaded": model is not None, "tokenizer_loaded": tokenizer is not None, }) if __name__ == "__main__": print("Visual AI is online -- http://0.0.0.0:7860") app.run(host="0.0.0.0", port=7860, threaded=True)