Motion-Rig / docker_app.py
STATICAI's picture
Restore Motion Rig project identity
807bba4 verified
Raw History Blame
6.93 kB
import json
import logging
import shutil
import time
import uuid
from pathlib import Path
from runtime import ROOT, MODEL_ROOT, configure_storage, startup_checks
configure_storage()
logging.basicConfig(filename=ROOT/"logs/workflow.log",level=logging.INFO,
format="%(asctime)s %(levelname)s %(message)s")
runtime_report = startup_checks()
import gradio as gr
import torch
import trimesh
from preflight import inspect_glb
import upstream.app as mia
mia.init_models()
# Initialize existing upstream component references used by its workflow.
# The original UI is not launched.
upstream_blocks = mia.init_blocks()
print("MODELS_READY dedicated CUDA inference",flush=True)
def preserve_and_inspect(path):
if not path:
raise gr.Error("Upload a GLB first.")
folder = ROOT/"inputs"/uuid.uuid4().hex
folder.mkdir()
original = folder/"original.glb"
shutil.copyfile(path,original)
try:
report = inspect_glb(original)
except (ValueError,KeyError,IndexError,TypeError) as exc:
report = {"inference_allowed":False,"warning":"Invalid or unsupported GLB: "+str(exc)}
report.update(original=str(original),cuda_available=torch.cuda.is_available(),gpu=runtime_report["gpu"])
(folder/"preflight.json").write_text(json.dumps(report,indent=2))
return str(original),report,str(original)
def execute(original,report):
if not original or not report or not report.get("inference_allowed"):
raise gr.Error((report or {}).get("warning") or "Inspect an eligible GLB before inference.")
if not torch.cuda.is_available():
raise gr.Error("CUDA is unavailable; CPU inference is disabled.")
# Recheck the preserved file before expensive processing.
checked=inspect_glb(original)
if not checked["inference_allowed"]:
raise gr.Error(checked["warning"])
job=ROOT/"outputs"/uuid.uuid4().hex
job.mkdir()
working=job/"character.glb"
shutil.copyfile(original,working)
db=mia.DB()
started=time.time()
stages=("Loading mesh","Preparing coordinates","Predicting skeleton and weights",
"Preparing rest pose","Exporting rigged asset","Finishing")
try:
for stage,_ in zip(stages,mia._pipeline(input_path=str(working),db=db,
reset_to_rest=True,rest_pose_type=None)):
yield stage,None,None,None
for p in (db.anim_path,db.anim_vis_path):
if not p or not Path(p).is_file() or Path(p).stat().st_size==0:
raise RuntimeError("Upstream export did not produce a nonempty asset")
# Validate export parsing without changing it.
result=trimesh.load(db.anim_vis_path,force="scene",process=False)
if not result.geometry:
raise RuntimeError("Exported GLB contains no geometry")
with open(db.anim_vis_path,"rb") as stream:
import struct
stream.seek(12)
length,kind=struct.unpack("<II",stream.read(8))
doc=json.loads(stream.read(length))
if not doc.get("skins") or not any("skin" in n for n in doc.get("nodes",[])):
raise RuntimeError("Exported GLB has no skin binding")
evidence=dict(status="PASS",original=original,glb=db.anim_vis_path,fbx=db.anim_path,
elapsed_seconds=time.time()-started,export_meshes=len(result.geometry),
skin_count=len(doc["skins"]),runtime=runtime_report["startup_id"])
(job/"result.json").write_text(json.dumps(evidence,indent=2))
logging.info("WORKFLOW_PASS %s",json.dumps(evidence))
print("WORKFLOW_PASS "+json.dumps(evidence),flush=True)
yield "Complete: rigged GLB and FBX exported.",db.anim_vis_path,db.anim_vis_path,db.anim_path
except Exception as exc:
logging.exception("Workflow failed in %s",job)
(job/"error.txt").write_text(str(exc))
raise gr.Error("Workflow failed: "+str(exc)+". Original preserved; details in /data/logs/workflow.log.")
finally:
mia.clear(db)
def example():
return preserve_and_inspect(MODEL_ROOT/"data/examples/ninja.glb")
def dense_character_test():
# Densify an existing character only for preflight; never modify its source.
source=MODEL_ROOT/"data/examples/ninja.glb"
mesh=trimesh.load(source,force="mesh",process=False)
while len(mesh.faces)<=500_000:
mesh=mesh.subdivide()
folder=ROOT/"inputs"/("dense-test-"+uuid.uuid4().hex)
folder.mkdir()
path=folder/"dense-character.glb"
mesh.export(path)
report=inspect_glb(path)
assert not report["inference_allowed"],"Dense character must be blocked before inference"
report["test"]="PASS: dense character preprocessing warning"
report["source"]=str(source)
report["original"]=str(path)
(folder/"preflight.json").write_text(json.dumps(report,indent=2))
print("DENSE_CHARACTER_TEST "+json.dumps(report),flush=True)
return report
with gr.Blocks(title="Make-It-Animatable") as demo:
gr.Markdown("# Motion Rig\nUpload a character GLB, inspect it, then prepare and export its rig.")
original=gr.State()
report_state=gr.State()
upload=gr.File(label="Upload GLB",file_types=[".glb"],type="filepath")
inspect_button=gr.Button("Inspect / preflight")
report_view=gr.JSON(label="Mesh preflight")
original_download=gr.File(label="Preserved original",interactive=False)
run=gr.Button("Prepare and export rig",variant="primary")
status=gr.Textbox(label="Status",interactive=False)
preview=gr.Model3D(label="Export preview",interactive=False)
with gr.Row():
glb=gr.File(label="Download rigged GLB",interactive=False)
fbx=gr.File(label="Download rigged FBX",interactive=False)
with gr.Accordion("Runtime and acceptance checks",open=False):
gr.JSON(value=runtime_report,label="Startup GPU and dependency checks")
normal_test=gr.Button("Load bundled normal character for testing")
dense_test=gr.Button("Test dense character preflight")
dense_result=gr.JSON(label="Dense test result")
inspect_button.click(preserve_and_inspect,inputs=upload,
outputs=[original,report_view,original_download]).success(
lambda x:x,inputs=report_view,outputs=report_state)
normal_test.click(example,outputs=[original,report_view,original_download]).success(
lambda x:x,inputs=report_view,outputs=report_state)
run.click(execute,inputs=[original,report_state],outputs=[status,preview,glb,fbx],
concurrency_limit=1,concurrency_id="workflow")
dense_test.click(dense_character_test,outputs=dense_result,
concurrency_limit=1,concurrency_id="workflow")
demo.queue(max_size=4,default_concurrency_limit=1).launch(
server_name="0.0.0.0",server_port=7860,show_error=True,ssr_mode=False,
allowed_paths=[str(ROOT/"inputs"),str(ROOT/"outputs")])