Update app.py

app.py

@@ -1,39 +1,94 @@
 import gradio as gr
 import os
 import uuid
+import asyncio
 
 from ultralytics import YOLO
 from openai import OpenAI
-from gtts import gTTS
-from PIL import Image
-
 from gradio_client import Client, handle_file
+import edge_tts
+
 
 # =========================
-# YOLO
+# YOLO MODEL
 # =========================
-
 model = YOLO("best_egypt.pt")
 
+
 # =========================
-# OpenRouter
+# OPENROUTER
 # =========================
-
 client = OpenAI(
     base_url="https://openrouter.ai/api/v1",
     api_key=os.environ["OPENROUTER_API_KEY"]
 )
 
+
 # =========================
-# Wav2Lip API
+# WAV2LIP API
 # =========================
-
 wav2lip_client = Client("fatma812/Wav2lip-ZeroGPU2")
 
+
+# =========================
+# PERSONA STYLE (ROYAL VOICES)
+# =========================
+PERSONA_PREFIX = {
+    "bust-of-ramesses-ii": "I speak with royal authority and unmatched power.",
+    "Ramses-II-Red-Granite-Statue": "I am the great king of strength and war.",
+    "Colossal-Statue-of-Ramesses-II": "I am eternal, carved in stone and glory.",
+
+    "Mask-of-Tutankhamun": "I am young, mysterious, and full of hidden legacy.",
+    "Tutankhamun": "I am the boy king of Egypt.",
+
+    "Nefertiti": "I speak with elegance, beauty, and royal grace.",
+    "Hatshepsut": "I rule with wisdom, divine authority, and strength."
+}
+
+
 # =========================
-# Generate Story + Voice
+# VOICE MAP (EDGE TTS)
 # =========================
+VOICE_MAP = {
+    # Ramses II (strong male Arabic)
+    "bust-of-ramesses-ii": "ar-DZ-IsmaelNeural",
+
+    # Tutankhamun (young male Arabic)
+    "Mask-of-Tutankhamun": "ar-EG-ShakirNeural",
+    "Tutankhamun": "ar-EG-ShakirNeural",
+
+    # Queens (elegant English voices)
+    "Nefertiti": "en-GB-SoniaNeural",
+    "Hatshepsut": "en-US-AriaNeural",
+}
+
+
+def get_voice(name):
+    return VOICE_MAP.get(name, "en-US-GuyNeural")
+
+
+def get_persona(name):
+    return PERSONA_PREFIX.get(
+        name,
+        "I am an ancient Egyptian artifact speaking from history."
+    )
+
+
+# =========================
+# EDGE TTS
+# =========================
+async def _speak(text, voice, path):
+    communicate = edge_tts.Communicate(text=text, voice=voice)
+    await communicate.save(path)
+
 
+def generate_audio(text, voice, path):
+    asyncio.run(_speak(text, voice, path))
+
+
+# =========================
+# GENERATE STORY + VOICE
+# =========================
 def generate(image_path, language):
 
     results = model(image_path)
@@ -42,52 +97,44 @@ def generate(image_path, language):
         raise gr.Error("No artifact detected")
 
     label_id = int(results[0].boxes.cls[0])
-
     artifact = results[0].names[label_id]
 
-    if language == "Arabic":
+    persona = get_persona(artifact)
 
+    if language == "Arabic":
         prompt = f"""
-أنت {artifact} أثر مصري قديم.
-تحدث بصيغة المتكلم في 2-3 جمل قصيرة تعرف بنفسك وتحكي حقيقة تاريخية واحدة.
+أنت {artifact} من آثار مصر القديمة.
+{persona}
+تحدث بصوت ملكي قوي في 2-3 جمل قصيرة.
 """
-
-        tts_lang = "ar"
-
+        voice = get_voice(artifact)
     else:
-
         prompt = f"""
 You are {artifact}, an ancient Egyptian artifact.
-Speak in first person in 2-3 short sentences introducing yourself and sharing one historical fact.
+{persona}
+Speak in a powerful royal tone in 2-3 short sentences.
 """
-
-        tts_lang = "en"
+        voice = get_voice(artifact)
 
     response = client.chat.completions.create(
         model="openai/gpt-4o-mini",
-        messages=[
-            {
-                "role": "user",
-                "content": prompt
-            }
-        ]
+        messages=[{"role": "user", "content": prompt}]
     )
 
     story = response.choices[0].message.content
 
-    audio_path = f"audio_{uuid.uuid4().hex}.mp3"
+    # add royal boost
+    story = persona + " " + story
 
-    gTTS(
-        text=story,
-        lang=tts_lang
-    ).save(audio_path)
+    audio_path = f"audio_{uuid.uuid4().hex}.mp3"
+    generate_audio(story, voice, audio_path)
 
     return artifact, story, audio_path
 
+
 # =========================
-# Generate Talking Video
+# MAKE VIDEO (WAV2LIP API)
 # =========================
-
 def make_video(image, audio):
 
     result = wav2lip_client.predict(
@@ -101,70 +148,41 @@ def make_video(image, audio):
 
     return result
 
+
 # =========================
 # UI
 # =========================
-
 with gr.Blocks() as demo:
 
-    gr.Markdown("# 🏛 Talking Egyptian Artifact AI")
+    gr.Markdown("# 🏛 Egyptian Talking AI (Final Version)")
 
-    image = gr.Image(
-        type="filepath",
-        label="Artifact Image"
-    )
+    image = gr.Image(type="filepath")
 
     language = gr.Radio(
         ["Arabic", "English"],
-        value="Arabic",
-        label="Language"
+        value="Arabic"
     )
 
-    generate_btn = gr.Button(
-        "Generate Story + Voice"
-    )
+    btn1 = gr.Button("Generate Story + Voice")
 
-    artifact_box = gr.Textbox(
-        label="Artifact"
-    )
+    artifact = gr.Textbox(label="Artifact")
+    story = gr.Textbox(label="Story")
+    audio = gr.Audio(type="filepath")
 
-    story_box = gr.Textbox(
-        label="Generated Story"
-    )
-
-    audio_box = gr.Audio(
-        type="filepath",
-        label="Generated Voice"
-    )
+    btn2 = gr.Button("Generate Video")
 
-    video_btn = gr.Button(
-        "Generate Talking Video"
-    )
-
-    video_box = gr.Video(
-        label="Talking Artifact"
-    )
+    video = gr.Video()
 
-    generate_btn.click(
+    btn1.click(
         generate,
-        inputs=[
-            image,
-            language
-        ],
-        outputs=[
-            artifact_box,
-            story_box,
-            audio_box
-        ]
+        inputs=[image, language],
+        outputs=[artifact, story, audio]
     )
 
-    video_btn.click(
+    btn2.click(
         make_video,
-        inputs=[
-            image,
-            audio_box
-        ],
-        outputs=video_box
+        inputs=[image, audio],
+        outputs=video
     )
 
 demo.queue()