Deploy SatQuery TEOChat Bi-Temporal Service

app.py

@@ -1,284 +1,296 @@
+"""Hugging Face ZeroGPU Space: SatQuery TEOChat Official Bi-Temporal Reasoning Service.
+
+Official Model: jirvin16/TEOChat
+Official Repository: https://github.com/ermongroup/TEOChat
+Responsibilities:
+1. Bi-temporal satellite image reasoning (T1 -> T2)
+2. Temporal change description
+3. Change-based Visual Question Answering (VQA)
+4. Temporal Earth Observation question answering
+
+NO mock fallbacks. NO NumPy image differencing. NO hardcoded percentages or categories.
+Uses real Video-LLaVA / TEOChat official inference code only.
 """
-SatQuery — Real EarthMind-4B Hugging Face Space
-
-This app uses the official sy1998/EarthMind-4B checkpoint and its
-custom `predict_forward(...)` inference interface.
-
-Supported:
-- Single-image VQA
-- Scene captioning / description
-- Text-guided segmentation (pixel mask)
-- Bounding-box extraction derived from the predicted mask
-
-Important:
-- No fake/deterministic fallback.
-- The uploaded image is passed to EarthMind itself.
-- Store HF_TOKEN as a Hugging Face Space Secret if the Space needs it.
-"""
-
-# IMPORTANT for Hugging Face ZeroGPU:
-# Import `spaces` before importing torch / transformers.
-import spaces
-
 import os
-from typing import Any, Optional
-
-import gradio as gr
-import numpy as np
+import io
+import time
+import base64
+import tempfile
+import logging
+from typing import Optional, Dict, Any, List
 from PIL import Image
 import torch
-from transformers import AutoModelForCausalLM, AutoTokenizer
-
-
-# ---------------------------------------------------------------------
-# Configuration
-# ---------------------------------------------------------------------
-
-MODEL_ID = os.getenv("EARTHMIND_MODEL_ID", "sy1998/EarthMind-4B")
-HF_TOKEN = os.getenv("HF_TOKEN") or os.getenv("HUGGINGFACE_HUB_TOKEN")
-
-print(f"[SatQuery] Loading official EarthMind checkpoint: {MODEL_ID}")
-
-# EarthMind's official Hugging Face quick-start uses:
-# AutoModelForCausalLM.from_pretrained(..., trust_remote_code=True,
-#                                      device_map="cuda:0")
-MODEL = AutoModelForCausalLM.from_pretrained(
-    MODEL_ID,
-    torch_dtype="auto",
-    device_map="cuda:0",
-    trust_remote_code=True,
-    token=HF_TOKEN if HF_TOKEN else None,
-)
-
-TOKENIZER = AutoTokenizer.from_pretrained(
-    MODEL_ID,
-    trust_remote_code=True,
-    token=HF_TOKEN if HF_TOKEN else None,
-)
-
-MODEL.eval()
-
-print("[SatQuery] EarthMind loaded successfully.")
-
-
-# ---------------------------------------------------------------------
-# Helpers
-# ---------------------------------------------------------------------
-
-def _build_prompt(query: str, task: str) -> str:
-    query = (query or "").strip()
-
-    if task == "vqa":
-        if not query:
-            query = "What is visible in this remote-sensing image?"
-        return f"<image>{query}"
-
-    if task == "captioning":
-        if query:
-            return f"<image>Please provide a detailed remote-sensing description of the image, with emphasis on: {query}"
-        return (
-            "<image>Please provide a detailed description of the remote-sensing "
-            "image, including major land-cover types, objects, infrastructure, "
-            "and spatial relationships."
-        )
-
-    if task == "grounding":
-        if not query:
-            query = "the most prominent water body"
-        return (
-            "<image>Please segment the region/object described by the following "
-            f"instruction: {query}"
-        )
-
-    raise ValueError(f"Unsupported task: {task}")
-
-
-def _to_numpy_mask(mask: Any) -> Optional[np.ndarray]:
-    """Normalize EarthMind's returned mask to H x W boolean NumPy array."""
-    if mask is None:
-        return None
-
-    if torch.is_tensor(mask):
-        mask = mask.detach().cpu().numpy()
-
-    arr = np.asarray(mask)
-
-    # Official EarthMind returns masks shaped like (1, H, W)
-    # for image inference. Be defensive about extra singleton dims.
-    while arr.ndim > 2:
-        arr = arr[0]
-
-    if arr.ndim != 2:
-        return None
-
-    return arr.astype(bool)
-
-
-def _extract_first_mask(result: dict) -> Optional[np.ndarray]:
-    masks = result.get("prediction_masks", [])
-    if not masks:
-        return None
-
-    # Usually masks = [np.ndarray(1, H, W)] for one [SEG] prediction.
-    first = masks[0]
-
-    # If several masks are returned, the first corresponds to the first
-    # [SEG] token/output region.
-    return _to_numpy_mask(first)
-
-
-def _mask_to_bbox(mask: Optional[np.ndarray]) -> Optional[list]:
-    """
-    Return normalized [ymin, xmin, ymax, xmax] in [0, 1].
-    """
-    if mask is None or not mask.any():
-        return None
-
-    ys, xs = np.where(mask)
-    h, w = mask.shape
-
-    ymin = float(ys.min() / h)
-    xmin = float(xs.min() / w)
-    ymax = float((ys.max() + 1) / h)
-    xmax = float((xs.max() + 1) / w)
-
-    return [
-        round(ymin, 6),
-        round(xmin, 6),
-        round(ymax, 6),
-        round(xmax, 6),
-    ]
-
-
-def _overlay_mask(image: Image.Image, mask: Optional[np.ndarray]) -> Optional[Image.Image]:
-    if mask is None:
-        return None
-
-    base = image.convert("RGBA")
-    mask_img = Image.fromarray((mask.astype(np.uint8) * 140), mode="L")
-
-    # Red overlay only for visualization.
-    overlay = Image.new("RGBA", base.size, (255, 0, 0, 0))
-    overlay.putalpha(mask_img)
-
-    return Image.alpha_composite(base, overlay).convert("RGB")
-
-
-# ---------------------------------------------------------------------
-# Real EarthMind inference
-# ---------------------------------------------------------------------
-
-@spaces.GPU(duration=120)
-def run_earthmind(
-    image: Image.Image,
-    query: str,
-    task: str,
-):
-    if image is None:
-        raise gr.Error("Please upload an image.")
-
-    if not isinstance(image, Image.Image):
-        image = Image.fromarray(np.asarray(image))
-
-    image = image.convert("RGB")
-
-    prompt = _build_prompt(query, task)
-
-    print(f"[SatQuery] Task: {task}")
-    print(f"[SatQuery] Prompt: {prompt}")
-
-    # THIS IS THE OFFICIAL EARTHMIND INFERENCE INTERFACE.
-    result = MODEL.predict_forward(
-        image=image,
-        text=prompt,
-        tokenizer=TOKENIZER,
-    )
-
-    prediction = str(result.get("prediction", "")).strip()
-    mask = _extract_first_mask(result)
-    bbox = _mask_to_bbox(mask)
-    overlay = _overlay_mask(image, mask)
-
-    trace = {
-        "model": MODEL_ID,
-        "task": task,
-        "inference_backend": "EarthMind.predict_forward",
-        "image_used": True,
-        "segmentation_generated": bool(mask is not None),
-        "normalized_bbox_ymin_xmin_ymax_xmax": bbox,
-    }
-
-    return prediction, overlay, trace
-
-
-# ---------------------------------------------------------------------
-# Gradio UI
-# ---------------------------------------------------------------------
-
-with gr.Blocks(title="SatQuery — EarthMind-4B") as demo:
-    gr.Markdown(
-        """
-# 🛰️ SatQuery — EarthMind-4B
-**Real Hugging Face EarthMind-4B inference**
-
-Use one remote-sensing image with a natural-language instruction.
-
-**VQA:** Ask a question about the image.  
-**Captioning:** Generate a remote-sensing scene description.  
-**Grounding:** Ask EarthMind to segment a specific object/region.
-"""
-    )
-
-    with gr.Row():
-        with gr.Column():
-            image = gr.Image(
-                type="pil",
-                label="Remote-Sensing Image",
-            )
-
-            task = gr.Dropdown(
-                choices=["vqa", "captioning", "grounding"],
-                value="vqa",
-                label="Task",
+import gradio as gr
+from fastapi import FastAPI, Request, Response, status
+from pydantic import BaseModel, Field
+
+# Setup logging
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger("satquery.teochat")
+
+try:
+    import spaces
+    has_spaces = True
+except ImportError:
+    has_spaces = False
+
+MODEL_ID = os.getenv("MODEL_ID", "jirvin16/TEOChat")
+DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
+
+_model = None
+_tokenizer = None
+_processor = None
+_load_error = None
+
+
+def ensure_videollava_installed():
+    """Ensure videollava package from official TEOChat repo is importable."""
+    try:
+        import videollava
+        return
+    except ImportError:
+        logger.info("[TEOChat] videollava not found. Attempting to install from official repo...")
+        import subprocess
+        import sys
+        try:
+            subprocess.run(
+                [sys.executable, "-m", "pip", "install", "--no-deps", "git+https://github.com/ermongroup/TEOChat.git"],
+                check=True,
+                capture_output=True
             )
+            logger.info("[TEOChat] Successfully installed TEOChat videollava package!")
+        except Exception as err:
+            logger.error(f"[TEOChat] Failed to pip install git+https://github.com/ermongroup/TEOChat.git: {err}")
+            raise ImportError(
+                "Official TEOChat 'videollava' library is required. "
+                "Please run: pip install git+https://github.com/ermongroup/TEOChat.git"
+            ) from err
+
+
+def load_teochat_model():
+    """Load the official TEOChat model using official videollava.eval.eval.load_model."""
+    global _model, _tokenizer, _processor, _load_error
+    if _model is not None and _tokenizer is not None and _processor is not None:
+        return _tokenizer, _model, _processor
+
+    logger.info(f"[TEOChat] Loading official model from {MODEL_ID} on {DEVICE}...")
+    try:
+        ensure_videollava_installed()
+        from videollava.eval.eval import load_model
+
+        # Load 8-bit when CUDA is available to fit ZeroGPU / T4 / A10G memory requirements
+        use_8bit = torch.cuda.is_available()
+        tokenizer, model, processor = load_model(
+            model_path=MODEL_ID,
+            model_base=None,
+            load_8bit=use_8bit,
+            device=DEVICE
+        )
 
-            query = gr.Textbox(
-                lines=3,
-                label="Natural-Language Query / Instruction",
-                placeholder=(
-                    "Examples:\n"
-                    "• What land-cover types are visible?\n"
-                    "• Describe the major objects and infrastructure.\n"
-                    "• Segment the water body."
-                ),
+        _tokenizer = tokenizer
+        _model = model
+        _processor = processor
+        _load_error = None
+        logger.info("[TEOChat] Successfully loaded official checkpoint!")
+        return _tokenizer, _model, _processor
+    except Exception as exc:
+        _load_error = str(exc)
+        logger.error(f"[TEOChat] Failed to load official model: {exc}", exc_info=True)
+        raise exc
+
+
+def decode_and_save_temp(img_input: Any, suffix: str = ".png") -> str:
+    """Decode image payload and persist to a temporary file path for the official processor."""
+    if img_input is None:
+        raise ValueError("Missing required temporal image input.")
+
+    # If it's already a valid existing file path on disk, return it directly
+    if isinstance(img_input, str) and os.path.exists(img_input):
+        return os.path.abspath(img_input)
+
+    pil_img = None
+    if isinstance(img_input, Image.Image):
+        pil_img = img_input.convert("RGB")
+    elif isinstance(img_input, str):
+        if img_input.startswith("data:image"):
+            _, b64 = img_input.split(",", 1)
+            raw = base64.b64decode(b64)
+            pil_img = Image.open(io.BytesIO(raw)).convert("RGB")
+        else:
+            try:
+                raw = base64.b64decode(img_input)
+                pil_img = Image.open(io.BytesIO(raw)).convert("RGB")
+            except Exception as e:
+                raise ValueError(f"Failed to decode base64 image: {e}")
+
+    if pil_img is None:
+        raise ValueError(f"Unrecognized image input format: {type(img_input)}")
+
+    tmp_file = tempfile.NamedTemporaryFile(suffix=suffix, delete=False)
+    pil_img.save(tmp_file.name, format="PNG")
+    tmp_file.close()
+    return tmp_file.name
+
+
+def _run_teochat_inference(
+    image_t1_input: Any,
+    image_t2_input: Any,
+    query: str
+) -> Dict[str, Any]:
+    """Execute official TEOChat forward pass on bi-temporal sequence [T1, T2]."""
+    t0 = time.time()
+    tokenizer, model, processor = load_teochat_model()
+    from videollava.eval.inference import run_inference_single
+
+    if not query or not query.strip():
+        query = "What changed between T1 and T2?"
+
+    # 1. Save or expose valid local file paths
+    t1_path = decode_and_save_temp(image_t1_input, suffix="_t1.png")
+    t2_path = decode_and_save_temp(image_t2_input, suffix="_t2.png")
+
+    try:
+        # 2. Preserve strict chronological order [T1, T2]
+        image_paths = [t1_path, t2_path]
+
+        # 3. Build official prompt containing <video>
+        inp = f"These are two satellite images in chronological order: <video> {query.strip()}"
+
+        logger.info(f"[TEOChat] Executing official run_inference_single with {len(image_paths)} temporal images...")
+        with torch.inference_mode():
+            response = run_inference_single(
+                model=model,
+                processor=processor,
+                tokenizer=tokenizer,
+                inp=inp,
+                image_paths=image_paths,
+                conv_mode="v1",
+                timestamps=[],
+                prompt_strategy="interleave",
+                chronological_prefix=True,
+                temperature=0.2,
+                max_new_tokens=256
             )
+    finally:
+        # Clean up temporary files
+        for p in [t1_path, t2_path]:
+            try:
+                if os.path.exists(p):
+                    os.remove(p)
+            except Exception:
+                pass
+
+    elapsed_sec = round(time.time() - t0, 3)
+
+    # Return standard normalized payload
+    return {
+        "status": "success",
+        "model": "TEOChat",
+        "task": "temporal_change_qa",
+        "answer": response,
+        "confidence": None,  # No fabricated confidence score
+        "metadata": {
+            "model_id": MODEL_ID,
+            "checkpoint": "jirvin16/TEOChat",
+            "num_temporal_images": 2,
+            "inference_time_sec": elapsed_sec,
+            "device": str(DEVICE)
+        }
+    }
 
-            run_btn = gr.Button(
-                "Run EarthMind",
-                variant="primary",
-            )
 
-        with gr.Column():
-            answer = gr.Textbox(
-                label="EarthMind Answer",
-                lines=10,
-            )
+# Wrap in ZeroGPU decorator if running on Hugging Face Spaces with ZeroGPU
+if has_spaces:
+    execute_temporal_inference = spaces.GPU(duration=60)(_run_teochat_inference)
+else:
+    execute_temporal_inference = _run_teochat_inference
+
+
+# ==============================================================================
+# Gradio Web Interface
+# ==============================================================================
+def gradio_fn(t1_img, t2_img, query):
+    try:
+        res = execute_temporal_inference(t1_img, t2_img, query)
+        return res["answer"]
+    except Exception as err:
+        return f"Error: {str(err)}"
+
+
+demo = gr.Interface(
+    fn=gradio_fn,
+    inputs=[
+        gr.Image(type="pil", label="Baseline Observation (T1 / Pre-Event)"),
+        gr.Image(type="pil", label="Target Observation (T2 / Post-Event)"),
+        gr.Textbox(lines=2, placeholder="e.g. What changed between T1 and T2?", label="Temporal Reasoning Query", value="What changed between T1 and T2?")
+    ],
+    outputs=[
+        gr.Textbox(label="TEOChat Official Delta Narrative")
+    ],
+    title="🛰️ SatQuery: Official TEOChat Bi-Temporal Specialist Service",
+    description="Official jirvin16/TEOChat Bi-Temporal Remote Sensing Specialist. Temporal change description and change-based VQA."
+)
 
-            result_image = gr.Image(
-                label="Predicted Spatial Evidence / Mask",
-                type="pil",
-            )
+app = demo.app
+
+
+# ==============================================================================
+# Direct REST API for SatQuery Backend Adapter
+# ==============================================================================
+class ChangeVqaRequest(BaseModel):
+    image_t1: Optional[str] = Field(default=None, description="T1 image (base64 or path)")
+    image_t2: Optional[str] = Field(default=None, description="T2 image (base64 or path)")
+    query: str = Field(default="", description="Temporal change query")
+    metadata: Dict[str, Any] = Field(default_factory=dict)
+    parameters: Dict[str, Any] = Field(default_factory=dict)
+
+
+@app.get("/health")
+def api_health(response: Response):
+    """Healthcheck endpoint. Returns healthy only if model is successfully loaded or can be loaded."""
+    global _model, _load_error
+    if _load_error:
+        response.status_code = status.HTTP_503_SERVICE_UNAVAILABLE
+        return {
+            "status": "unhealthy",
+            "model": "jirvin16/TEOChat",
+            "error": _load_error
+        }
+    return {
+        "status": "healthy",
+        "model": "jirvin16/TEOChat",
+        "device": str(DEVICE),
+        "is_loaded": _model is not None
+    }
 
-            trace = gr.JSON(
-                label="Execution Trace",
-            )
 
-    run_btn.click(
-        fn=run_earthmind,
-        inputs=[image, query, task],
-        outputs=[answer, result_image, trace],
-    )
+@app.post("/change_vqa")
+async def api_change_vqa(req: ChangeVqaRequest, response: Response):
+    """Direct REST endpoint called by SatQuery ChangeAnalysisAdapter."""
+    try:
+        if not req.image_t1 or not req.image_t2:
+            response.status_code = status.HTTP_400_BAD_REQUEST
+            return {
+                "status": "error",
+                "model": "jirvin16/TEOChat",
+                "error": "Both T1 and T2 images are required for bi-temporal reasoning."
+            }
+
+        res = execute_temporal_inference(
+            image_t1_input=req.image_t1,
+            image_t2_input=req.image_t2,
+            query=req.query
+        )
+        return res
+    except Exception as exc:
+        logger.error(f"[TEOChat] API Execution error: {exc}", exc_info=True)
+        response.status_code = status.HTTP_500_INTERNAL_SERVER_ERROR
+        return {
+            "status": "error",
+            "model": "jirvin16/TEOChat",
+            "error": str(exc)
+        }
 
 
 if __name__ == "__main__":
-    demo.launch()
+    demo.launch(server_name="0.0.0.0", server_port=7860, ssr_mode=False)

requirements.txt

@@ -1,7 +1,21 @@
 gradio>=4.20.0
 fastapi>=0.110.0
 uvicorn>=0.28.0
-torch>=2.0.0
-torchvision>=0.15.0
+pydantic>=2.0.0
+torch==2.2.1
+torchvision==0.17.1
+transformers==4.31.0
+tokenizers>=0.12.1,<0.14
+sentencepiece==0.1.99
+accelerate==0.26.1
+peft==0.7.1
+bitsandbytes==0.41.0
 pillow>=10.0.0
 numpy>=1.24.0
+shortuuid
+einops==0.6.1
+einops-exts==0.0.4
+timm==0.6.12
+opencv-python-headless>=4.7.0.72
+decord==0.6.0
+git+https://github.com/ermongroup/TEOChat.git