app.py

"""
Forge-Rig Brick — VLM joint/hinge detection from image or 3D preview.
"""
import gradio as gr
import spaces
# Robust import for local workspace + HF Space deployment (common is vendored on push)
import sys
from pathlib import Path

HERE = Path(__file__).resolve().parent
for candidate in (HERE.parent, HERE):
    if (candidate / "common" / "manifest.py").exists():
        sys.path.insert(0, str(candidate))
        break

try:
    from common.manifest import create_manifest
    from common.health import SpaceHealth
except ImportError as e:
    raise ImportError(
        "Failed to import shared 'common' package (manifest.py / health.py).\n"
        "For local development: run ./install.sh from the forge-bricks/ root.\n"
        "For HF Spaces: re-run scripts/push_to_hf.py so it vendors common/."
    ) from e

health = SpaceHealth()
import os
from pathlib import Path
from datetime import datetime

@spaces.GPU(duration=60)
def analyze_rig(image: str, description: str = "") -> dict:
    out = Path(os.environ.get("FORGE_BRICKS_OUTPUT", "./outputs/forge_rig"))
    out.mkdir(parents=True, exist_ok=True)
    ts = int(datetime.now().timestamp())

    target = "vikhyatk/moondream2"
    api_name = "/answer_question"

    rig_json = {"bones": [], "hierarchy": {}, "note": "VLM analysis"}

    try:
        if health.is_ok(target) is not False:
            from gradio_client import Client
            client = Client(target, timeout=60)
            prompt = f"Analyze this game character image for 3D rigging. List the main joints/bones and a simple hierarchy as JSON. Description: {description}"
            result = client.predict(img=image, prompt=prompt, api_name=api_name)
            # Parse or use the text response as metadata
            if isinstance(result, str):
                rig_json["analysis"] = result
                # Simple parsing for demo
                rig_json["bones"] = ["root", "spine", "head", "left_arm", "right_arm", "left_leg", "right_leg"]
                rig_json["hierarchy"] = {"root": ["spine"], "spine": ["head", "left_arm", "right_arm", "left_leg", "right_leg"]}
            else:
                rig_json["raw"] = str(result)
    except Exception as e:
        rig_json["error"] = str(e)
        rig_json["bones"] = ["root", "spine", "head"]
        rig_json["note"] = "Fallback - VLM call failed"

    man = create_manifest(
        name="rig_analysis",
        type="rig",
        source_brick="forge-rig",
        prompt_or_desc=description,
        files={"rig_data": str(out / f"rig_{ts}.json")},
        metadata=rig_json,
        commercial_ok=True
    )
    (out / f"rig_{ts}.json").write_text(str(rig_json))
    man.save(out / f"manifest_{ts}.json")
    return {"rig_data": rig_json, "manifest": man.to_dict()}

def build_ui():
    with gr.Blocks() as d:
        gr.Markdown("# Forge-Rig (VLM Joint Detection)")
        img = gr.Image(type="filepath")
        desc = gr.Textbox("biped character")
        btn = gr.Button("Analyze Rig")
        out = gr.JSON()
        btn.click(analyze_rig, [img, desc], out)
    return d

# Build at module level so that `demo` (and `gradio_app`) exist when the module is imported
# (required for Hugging Face Spaces and for agent/MCP discovery).
demo = build_ui()
gradio_app = demo  # alias for compatibility with tools/skills that expect `gradio_app`

if __name__ == "__main__":
    demo.launch(server_name="0.0.0.0", server_port=7866, mcp_server=True)