Update app.py

app.py

@@ -1,133 +1,257 @@
+# app.py — FastAPI REST API + mounted Gradio UI (Hugging Face Spaces compatible)
 import os
 import shutil
 import asyncio
+import inspect
 from typing import Optional
 
+from fastapi import FastAPI, UploadFile, File, Form
+from fastapi.responses import JSONResponse, PlainTextResponse
+from fastapi.middleware.cors import CORSMiddleware
 import gradio as gr
-from fastapi import FastAPI, UploadFile, Form
 import uvicorn
-import socket
 
+# ---- Your module ----
 from multimodal_module import MultiModalChatModule
 
-# Initialize AI module
+# Instantiate once at import time
 AI = MultiModalChatModule()
 
-# ---------------------------
-# Utility
-# ---------------------------
-class GradioFileWrapper:
-    def __init__(self, file_path):
-        self._path = file_path
+TMP_DIR = "/tmp"
+os.makedirs(TMP_DIR, exist_ok=True)
 
-    async def download_to_drive(self, dst_path: str):
+# ---------------- Helpers ----------------
+class FileWrapper:
+    """Tiny adapter so your module can .download_to_drive(path)."""
+    def __init__(self, path: str):
+        self._path = path
+
+    async def download_to_drive(self, dst_path: str) -> None:
         loop = asyncio.get_event_loop()
         await loop.run_in_executor(None, shutil.copyfile, self._path, dst_path)
 
-def run_async(coro):
-    return asyncio.run(coro)
+async def save_upload_to_tmp(up: UploadFile) -> str:
+    if not up or not up.filename:
+        raise ValueError("No file uploaded")
+    dest = os.path.join(TMP_DIR, up.filename)
+    data = await up.read()
+    with open(dest, "wb") as f:
+        f.write(data)
+    return dest
 
-def get_free_port(default=7860):
-    """Find a free port if default is busy."""
-    with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
-        try:
-            s.bind(("0.0.0.0", default))
-            return default
-        except OSError:
-            s.bind(("0.0.0.0", 0))
-            return s.getsockname()[1]
-
-# ---------------------------
-# FastAPI API for external apps
-# ---------------------------
-api = FastAPI()
-
-@api.post("/api/text_chat")
-async def api_text_chat(
-    user_id: Optional[int] = Form(0),
+async def call_ai(fn, *args, **kwargs):
+    """Call AI methods whether they are sync or async."""
+    if fn is None:
+        raise AttributeError("Requested AI method is not implemented in multimodal_module")
+    if inspect.iscoroutinefunction(fn):
+        return await fn(*args, **kwargs)
+    return await asyncio.to_thread(lambda: fn(*args, **kwargs))
+
+# ---------------- FastAPI app ----------------
+app = FastAPI(title="Multimodal Module API", version="1.0.0")
+
+# CORS so external apps can call it
+app.add_middleware(
+    CORSMiddleware,
+    allow_origins=["*"],          # tighten for production
+    allow_credentials=True,
+    allow_methods=["*"],
+    allow_headers=["*"],
+)
+
+# ---- Health / root ----
+@app.get("/health", response_class=PlainTextResponse)
+async def health():
+    return "ok"
+
+@app.get("/")
+async def root():
+    return {
+        "name": "Multimodal Module API",
+        "status": "ready",
+        "docs": "/docs",
+        "gradio_ui": "/ui"
+    }
+
+# ---------------- REST Endpoints ----------------
+# Text chat
+@app.post("/api/text")
+async def api_text(
     text: str = Form(...),
-    lang: str = Form("en")
+    user_id: Optional[int] = Form(0),
+    lang: str = Form("en"),
+):
+    try:
+        fn = getattr(AI, "generate_response", getattr(AI, "process_text", None))
+        reply = await call_ai(fn, text, int(user_id), lang)
+        return {"status": "ok", "reply": reply}
+    except Exception as e:
+        return JSONResponse({"error": str(e)}, status_code=500)
+
+# Hugging Face-style predict (optional)
+@app.post("/api/predict")
+async def api_predict(
+    inputs: str = Form(...),
+    user_id: Optional[int] = Form(0),
+    lang: str = Form("en"),
+):
+    try:
+        fn = getattr(AI, "generate_response", getattr(AI, "process_text", None))
+        reply = await call_ai(fn, inputs, int(user_id), lang)
+        return {"data": [reply]}
+    except Exception as e:
+        return JSONResponse({"error": str(e)}, status_code=500)
+
+# Voice -> ASR / emotion
+@app.post("/api/voice")
+async def api_voice(
+    user_id: Optional[int] = Form(0),
+    audio_file: UploadFile = File(...),
+):
+    try:
+        path = await save_upload_to_tmp(audio_file)
+        fn = getattr(AI, "process_voice_message", None)
+        result = await call_ai(fn, FileWrapper(path), int(user_id))
+        return JSONResponse(result)
+    except Exception as e:
+        return JSONResponse({"error": str(e)}, status_code=500)
+
+# TTS
+@app.post("/api/voice_reply")
+async def api_voice_reply(
+    user_id: Optional[int] = Form(0),
+    reply_text: str = Form(...),
+    fmt: str = Form("ogg"),
 ):
     try:
-        reply = await AI.generate_response(text, int(user_id), lang)
-        return {"reply": reply}
+        fn = getattr(AI, "generate_voice_reply", None)
+        out_path = await call_ai(fn, reply_text, int(user_id), fmt)
+        return {"status": "ok", "file": out_path}
     except Exception as e:
-        return {"error": str(e)}
+        return JSONResponse({"error": str(e)}, status_code=500)
 
-@api.post("/api/image_caption")
-async def api_image_caption(user_id: Optional[int] = Form(0), image: UploadFile = None):
+# Image caption
+@app.post("/api/image_caption")
+async def api_image_caption(
+    user_id: Optional[int] = Form(0),
+    image_file: UploadFile = File(...),
+):
     try:
-        temp_path = f"/tmp/{image.filename}"
-        with open(temp_path, "wb") as f:
-            f.write(await image.read())
-        wrapper = GradioFileWrapper(temp_path)
-        caption = await AI.process_image_message(wrapper, int(user_id))
-        return {"caption": caption}
+        path = await save_upload_to_tmp(image_file)
+        fn = getattr(AI, "process_image_message", None)
+        caption = await call_ai(fn, FileWrapper(path), int(user_id))
+        return {"status": "ok", "caption": caption}
     except Exception as e:
-        return {"error": str(e)}
+        return JSONResponse({"error": str(e)}, status_code=500)
 
-@api.post("/api/voice_process")
-async def api_voice_process(user_id: Optional[int] = Form(0), audio: UploadFile = None):
+# Text-to-image
+@app.post("/api/generate_image")
+async def api_generate_image(
+    user_id: Optional[int] = Form(0),
+    prompt: str = Form(...),
+    width: int = Form(512),
+    height: int = Form(512),
+    steps: int = Form(30),
+):
     try:
-        temp_path = f"/tmp/{audio.filename}"
-        with open(temp_path, "wb") as f:
-            f.write(await audio.read())
-        wrapper = GradioFileWrapper(temp_path)
-        reply = await AI.process_voice_message(wrapper, int(user_id))
-        return {"reply": reply}
+        fn = getattr(AI, "generate_image_from_text", None)
+        out_path = await call_ai(fn, prompt, int(user_id), width, height, steps)
+        return {"status": "ok", "file": out_path}
     except Exception as e:
-        return {"error": str(e)}
+        return JSONResponse({"error": str(e)}, status_code=500)
 
-@api.post("/api/video_process")
-async def api_video_process(user_id: Optional[int] = Form(0), video: UploadFile = None):
+# Image edit / inpaint
+@app.post("/api/edit_image")
+async def api_edit_image(
+    user_id: Optional[int] = Form(0),
+    image_file: UploadFile = File(...),
+    mask_file: Optional[UploadFile] = File(None),
+    prompt: str = Form(""),
+):
     try:
-        temp_path = f"/tmp/{video.filename}"
-        with open(temp_path, "wb") as f:
-            f.write(await video.read())
-        wrapper = GradioFileWrapper(temp_path)
-        reply = await AI.process_video_message(wrapper, int(user_id))
-        return {"reply": reply}
+        img_path = await save_upload_to_tmp(image_file)
+        mask_path = None
+        if mask_file:
+            mask_path = await save_upload_to_tmp(mask_file)
+        fn = getattr(AI, "edit_image_inpaint", None)
+        out_path = await call_ai(
+            fn,
+            FileWrapper(img_path),
+            FileWrapper(mask_path) if mask_path else None,
+            prompt,
+            int(user_id),
+        )
+        return {"status": "ok", "file": out_path}
     except Exception as e:
-        return {"error": str(e)}
+        return JSONResponse({"error": str(e)}, status_code=500)
 
-@api.post("/api/file_process")
-async def api_file_process(user_id: Optional[int] = Form(0), file: UploadFile = None):
+# Video
+@app.post("/api/video")
+async def api_video(
+    user_id: Optional[int] = Form(0),
+    video_file: UploadFile = File(...),
+):
     try:
-        temp_path = f"/tmp/{file.filename}"
-        with open(temp_path, "wb") as f:
-            f.write(await file.read())
-        wrapper = GradioFileWrapper(temp_path)
-        reply = await AI.process_file_message(wrapper, int(user_id))
-        return {"reply": reply}
+        path = await save_upload_to_tmp(video_file)
+        fn = getattr(AI, "process_video", None)
+        result = await call_ai(fn, FileWrapper(path), int(user_id))
+        return JSONResponse(result)
     except Exception as e:
-        return {"error": str(e)}
-
-# ---------------------------
-# Gradio UI
-# ---------------------------
-with gr.Blocks(title="Multimodal Bot") as demo:
-    gr.Markdown("# 🧠 Multimodal Bot\nInteract via text, voice, images, video, or files.")
-
-    with gr.Tab("💬 Text Chat"):
-        user_id_txt = gr.Textbox(label="User ID", placeholder="0")
-        lang_sel = gr.Dropdown(choices=["en","zh","ja","ko","es","fr","de","it"], value="en", label="Language")
-        txt_in = gr.Textbox(label="Your message", lines=4)
-        txt_out = gr.Textbox(label="Bot reply", lines=6)
-        gr.Button("Send").click(lambda uid, txt, lang: run_async(AI.generate_response(txt, int(uid or 0), lang)),
-                                [user_id_txt, txt_in, lang_sel], txt_out)
-
-    with gr.Tab("🖼 Image Captioning"):
-        user_id_img = gr.Textbox(label="User ID", placeholder="0")
-        img_in = gr.Image(type="filepath", label="Upload an image")
-        img_out = gr.Textbox(label="Caption")
-        gr.Button("Caption").click(lambda uid, img: run_async(AI.process_image_message(GradioFileWrapper(img), int(uid or 0))),
-                                   [user_id_img, img_in], img_out)
-
-# ---------------------------
-# Mount Gradio UI to FastAPI
-# ---------------------------
-api = gr.mount_gradio_app(api, demo, path="/")
+        return JSONResponse({"error": str(e)}, status_code=500)
+
+# Files (PDF/DOCX/TXT)
+@app.post("/api/file")
+async def api_file(
+    user_id: Optional[int] = Form(0),
+    file_obj: UploadFile = File(...),
+):
+    try:
+        path = await save_upload_to_tmp(file_obj)
+        fn = getattr(AI, "process_file", None)
+        result = await call_ai(fn, FileWrapper(path), int(user_id))
+        return JSONResponse(result)
+    except Exception as e:
+        return JSONResponse({"error": str(e)}, status_code=500)
+
+# Code completion
+@app.post("/api/code")
+async def api_code(
+    user_id: Optional[int] = Form(0),
+    prompt: str = Form(...),
+    max_tokens: int = Form(512),
+):
+    try:
+        fn = getattr(AI, "code_complete", None)
+        try:
+            result = await call_ai(fn, int(user_id), prompt, max_tokens)
+        except TypeError:
+            result = await call_ai(fn, prompt, max_tokens=max_tokens)
+        return {"status": "ok", "code": result}
+    except Exception as e:
+        return JSONResponse({"error": str(e)}, status_code=500)
+
+# ---------------- Gradio UI (mounted at /ui) ----------------
+def _gradio_text_fn(text, user_id, lang):
+    fn = getattr(AI, "generate_response", getattr(AI, "process_text", None))
+    if fn is None:
+        return "Error: text handler not implemented in multimodal_module"
+    # Gradio callbacks run in a worker thread, safe to create/own an event loop
+    return asyncio.run(call_ai(fn, text, int(user_id or 0), lang))
+
+with gr.Blocks(title="Multimodal Bot — UI") as demo:
+    gr.Markdown("# 🧠 Multimodal Bot — UI\nThis is a helper UI. Use the REST API for external apps.")
+    with gr.Row():
+        g_uid = gr.Textbox(label="User ID", value="0")
+        g_lang = gr.Dropdown(["en", "zh", "ja", "ko", "es", "fr", "de", "it"], value="en", label="Language")
+    g_in = gr.Textbox(lines=3, label="Message")
+    g_out = gr.Textbox(lines=6, label="Reply")
+    gr.Button("Send").click(_gradio_text_fn, [g_in, g_uid, g_lang], g_out)
+
+# Mount Gradio *into* FastAPI at /ui (does not open another port)
+app = gr.mount_gradio_app(app, demo, path="/ui")
 
+# ---------------- Entrypoint ----------------
 if __name__ == "__main__":
-    port = get_free_port()
-    uvicorn.run(api, host="0.0.0.0", port=port)
+    # Hugging Face Spaces (FastAPI template) sets PORT; bind exactly to it.
+    port = int(os.environ.get("PORT", "7860"))
+    uvicorn.run("app:app", host="0.0.0.0", port=port)