app.py

#!/usr/bin/env python3
# rebuild: 2026-03-29
"""
Splat Explorer — HF Spaces version
"""
import urllib.request
import json
import os
import tempfile
import base64
import threading
import uuid

from fastapi import FastAPI, Request
from fastapi.responses import JSONResponse, HTMLResponse, Response
from fastapi.middleware.gzip import GZipMiddleware
import uvicorn

PORT              = int(os.environ.get("PORT", 7860))
HF_SPACE          = "frogleo/Image-to-3D"
TEXT2IMG_MODEL    = "black-forest-labs/FLUX.1-schnell"
VISION_MODEL      = "Qwen/Qwen2.5-VL-7B-Instruct"
VISION_API        = "https://router.huggingface.co/v1/chat/completions"
OCTREE_RESOLUTION = int(os.environ.get("OCTREE_RESOLUTION", 256))   # upstream UI default
TARGET_FACE_NUM   = int(os.environ.get("TARGET_FACE_NUM", 10000))   # upstream UI default
_app_dir          = os.path.dirname(os.path.abspath(__file__))

_client = None
_client_token = None
_client_lock = threading.Lock()

def get_client(hf_token=None):
    global _client, _client_token
    with _client_lock:
        if _client is None or hf_token != _client_token:
            from gradio_client import Client
            if hf_token:
                os.environ["HF_TOKEN"] = hf_token
            _client = Client(HF_SPACE)
            _client_token = hf_token
    return _client

def vision_analyze(images_data, hf_token):
    content = [{"type": "image_url", "image_url": {"url": u}} for u in images_data[:6]]
    content.append({"type": "text", "text": (
        "You are a 3D modeling expert. These are reference photos of the same object "
        "from different angles. Produce a precise description optimized for 3D generation. "
        "Cover: shape, proportions, materials, textures, colors. 2-4 sentences only."
    )})
    payload = json.dumps({"model": VISION_MODEL,
                          "messages": [{"role": "user", "content": content}],
                          "max_tokens": 300, "temperature": 0.3}).encode()
    req = urllib.request.Request(VISION_API, data=payload,
                                 headers={"Authorization": f"Bearer {hf_token}",
                                          "Content-Type": "application/json"})
    with urllib.request.urlopen(req, timeout=180) as r:
        return json.loads(r.read())["choices"][0]["message"]["content"].strip()

def flux_generate(prompt, hf_token):
    import io
    from huggingface_hub import InferenceClient
    client = InferenceClient(token=hf_token)
    image = client.text_to_image(prompt, model=TEXT2IMG_MODEL)
    buf = io.BytesIO()
    image.save(buf, format="PNG")
    return buf.getvalue()

def image_to_3d(img_bytes, ext="png", hf_token=None):
    with tempfile.NamedTemporaryFile(suffix=f".{ext}", delete=False) as f:
        f.write(img_bytes)
        tmp = f.name

    last_err = None
    for attempt in range(2):  # 1 retry — ZeroGPU occasionally fails on first allocation
        try:
            from gradio_client import handle_file
            result = get_client(hf_token).predict(
                image=handle_file(tmp), steps=5, guidance_scale=5.5,
                seed=1234, octree_resolution=OCTREE_RESOLUTION, num_chunks=8000,
                target_face_num=TARGET_FACE_NUM, randomize_seed=True, api_name="/gen_shape"
            )
            os.unlink(tmp)
            return result[2] if isinstance(result, (list, tuple)) and len(result) > 2 else None
        except Exception as e:
            last_err = e
            err_str = str(e).lower()
            if "gpu task" in err_str or "aborted" in err_str or "timeout" in err_str:
                # Reset client so next attempt gets a fresh connection
                with _client_lock:
                    global _client, _client_token
                    _client = None
                    _client_token = None
                if attempt == 0:
                    continue  # retry once
            break  # non-GPU error — don't retry

    try:
        os.unlink(tmp)
    except OSError:
        pass
    raise last_err

def fetch_file(path):
    if path and os.path.exists(path):
        with open(path, "rb") as f:
            return f.read()
    url = ("https://frogleo-image-to-3d.hf.space" + path) if path.startswith("/") else path
    with urllib.request.urlopen(url, timeout=60) as r:
        return r.read()


# ── Model store (binary, avoids base64 +33% overhead) ────────────────────────
_models: dict = {}
_model_order: list = []
_MAX_MODELS = 20
_model_lock = threading.Lock()

def store_model(data: bytes) -> str:
    mid = str(uuid.uuid4())[:12]  # 48-bit collision resistance
    with _model_lock:
        _models[mid] = data
        _model_order.append(mid)
        while len(_model_order) > _MAX_MODELS:
            old = _model_order.pop(0)
            _models.pop(old, None)
    return mid


# ── FastAPI ──────────────────────────────────────────────────────────────────
app = FastAPI()
app.add_middleware(GZipMiddleware, minimum_size=1000)

@app.get("/health")
def health():
    return {"status": "ok"}

_HTML_URL = "https://huggingface.co/spaces/ArtelTaleb/splat-explorer/resolve/main/splat-explorer.html"
_html_cache = None

def _get_html():
    global _html_cache
    try:
        req = urllib.request.Request(_HTML_URL, headers={"Cache-Control": "no-cache"})
        with urllib.request.urlopen(req, timeout=5) as r:
            _html_cache = r.read()
    except Exception:
        if _html_cache is None:
            with open(os.path.join(_app_dir, "splat-explorer.html"), "rb") as f:
                _html_cache = f.read()
    return _html_cache

@app.get("/")
def root():
    return HTMLResponse(_get_html())

@app.get("/splat-explorer.html")
def splat():
    return HTMLResponse(_get_html())

@app.get("/model/{model_id}")
def get_model(model_id: str):
    data = _models.get(model_id)
    if not data:
        return JSONResponse({"error": "Model not found"}, status_code=404)
    return Response(
        content=data,
        media_type="model/gltf-binary",
        headers={
            "Cache-Control": "public, max-age=3600",
            "Content-Encoding": "identity",  # skip GZip middleware — GLB already compact
        }
    )

@app.post("/generate")
async def generate(request: Request):
    try:
        body       = await request.json()
        images     = body.get("images")
        image_data = body.get("image")
        prompt     = body.get("prompt")
        hf_token   = body.get("hf_token", "")

        if not images and not image_data and not prompt:
            return JSONResponse({"error": "Provide images, image, or prompt."}, status_code=400)

        generated_prompt = None
        vision_warning   = None

        if images and len(images) > 1:
            if hf_token:
                try:
                    generated_prompt = vision_analyze(images, hf_token)
                except Exception as e:
                    vision_warning = f"Vision skipped ({e.__class__.__name__})"
                if generated_prompt:
                    try:
                        img_bytes = flux_generate(generated_prompt, hf_token)
                        img_ext = "png"
                    except Exception as e:
                        vision_warning = (vision_warning or "") + f" FLUX skipped ({e.__class__.__name__})"
                        generated_prompt = None
            if not generated_prompt:
                best = max(images, key=len)
                header, b64 = best.split(",", 1)
                img_ext = "png" if "png" in header else "jpg"
                img_bytes = base64.b64decode(b64)

        elif prompt and not image_data:
            if not hf_token:
                return JSONResponse({"error": "HF token required."}, status_code=400)
            try:
                img_bytes = flux_generate(prompt, hf_token)
            except Exception as e:
                return JSONResponse({"error": f"[FLUX] {e}"}, status_code=500)
            img_ext = "png"
        else:
            data = image_data or images[0]
            header, b64 = data.split(",", 1)
            img_ext = "png" if "png" in header else "jpg"
            img_bytes = base64.b64decode(b64)

        try:
            glb_path = image_to_3d(img_bytes, img_ext, hf_token=hf_token)
        except Exception as e:
            return JSONResponse({"error": f"[Image-to-3D] {e}"}, status_code=500)

        if not glb_path:
            return JSONResponse({"error": "No GLB output."}, status_code=500)

        glb_bytes = fetch_file(glb_path)
        model_id  = store_model(glb_bytes)

        resp = {"glb": f"/model/{model_id}"}
        if generated_prompt: resp["prompt"] = generated_prompt
        if vision_warning:   resp["warning"] = vision_warning
        return JSONResponse(resp)

    except Exception as e:
        import traceback; traceback.print_exc()
        return JSONResponse({"error": str(e)}, status_code=500)


if __name__ == "__main__":
    print(f"Splat Explorer → http://0.0.0.0:{PORT}/")
    uvicorn.run(app, host="0.0.0.0", port=PORT)