Spaces:
Build error
Build error
| import json, subprocess, tempfile, uuid, os, shutil | |
| import numpy as np, torch, spaces, gradio as gr | |
| from PIL import Image | |
| from diffusers import QwenImageEditPipeline | |
| # ---------- GPU models: load at module level, ZeroGPU moves them on demand ---------- | |
| edit_pipe = QwenImageEditPipeline.from_pretrained( | |
| "Qwen/Qwen-Image-Edit-2509", torch_dtype=torch.bfloat16).to("cuda") | |
| # 4-8 step Lightning LoRA keeps the call inside the time cap (check the repo for the exact weight file) | |
| edit_pipe.load_lora_weights("lightx2v/Qwen-Image-Lightning", weight_name="<edit lightning .safetensors>") | |
| from hy3dgen.shapegen import Hunyuan3DDiTFlowMatchingPipeline # or your TRELLIS wrapper | |
| shape_pipe = Hunyuan3DDiTFlowMatchingPipeline.from_pretrained("tencent/Hunyuan3D-2mini").to("cuda") | |
| # ---------- CPU pose model ---------- | |
| from rtmlib import Body | |
| pose_model = Body(mode="balanced", backend="onnxruntime", device="cpu") # COCO-17 keypoints | |
| PROMPT = ("Same character, same outfit and colors, full body, T-pose with arms straight out " | |
| "to the sides, facing the camera, plain white background") | |
| def make_tpose(img): | |
| return edit_pipe(image=img, prompt=PROMPT, num_inference_steps=8, | |
| true_cfg_scale=1.0).images[0] | |
| def make_mesh(tpose_img, job): | |
| mesh = shape_pipe(image=tpose_img)[0] # trimesh | |
| path = f"{job}/raw.glb"; mesh.export(path) | |
| return path | |
| def keypoints(img): | |
| kp, sc = pose_model(np.array(img.convert("RGB"))[:, :, ::-1]) | |
| return kp[0].tolist(), sc[0].tolist() # first person only | |
| def fg_bbox(img, thr=240): | |
| a = np.array(img.convert("L")); ys, xs = np.where(a < thr) | |
| return [int(xs.min()), int(ys.min()), int(xs.max()), int(ys.max())] | |
| def blender(job, cmd, **kw): | |
| kw.update(cmd=cmd, job=job) | |
| subprocess.run(["python", "blender_worker.py", json.dumps(kw)], check=True) | |
| # ---------- Gradio event handlers ---------- | |
| def step_generate(img): | |
| job = tempfile.mkdtemp(); img = img.convert("RGB") | |
| tpose = make_tpose(img); tpose.save(f"{job}/tpose.png"); img.save(f"{job}/orig.png") | |
| raw = make_mesh(tpose, job) | |
| kp_t, sc_t = keypoints(tpose) | |
| json.dump(dict(kp=kp_t, sc=sc_t, bbox=fg_bbox(tpose), size=tpose.size), open(f"{job}/tpose_kp.json", "w")) | |
| blender(job, "rig") # auto-rig from T-pose joints | |
| return job, tpose, f"{job}/rigged.glb" | |
| def step_pose(job): | |
| orig = Image.open(f"{job}/orig.png") | |
| kp, sc = keypoints(orig) | |
| json.dump(dict(kp=kp, sc=sc, size=orig.size), open(f"{job}/target_kp.json", "w")) | |
| blender(job, "pose") # writes posed.glb/.obj/.fbx | |
| return f"{job}/posed.glb", f"{job}/posed.glb", f"{job}/posed.obj", f"{job}/posed.fbx" | |
| def step_uv(job): | |
| blender(job, "uv"); return f"{job}/uv_layout.png", f"{job}/texture.png" | |
| def step_texture(job, tex): | |
| shutil.copy(tex, f"{job}/new_texture.png"); blender(job, "retexture") | |
| return f"{job}/posed.glb", f"{job}/posed.glb", f"{job}/posed.obj", f"{job}/posed.fbx" | |
| with gr.Blocks() as demo: | |
| job = gr.State() | |
| with gr.Row(): | |
| inp = gr.Image(type="pil", label="Upload") | |
| tp = gr.Image(label="T-pose") | |
| view = gr.Model3D(label="3D") | |
| go = gr.Button("1. Generate T-pose + 3D + rig") | |
| pose = gr.Button("2. Pose like the uploaded image") | |
| with gr.Row(): | |
| d_glb, d_obj, d_fbx = (gr.DownloadButton(f"Download {x}") for x in ("GLB", "OBJ", "FBX")) | |
| with gr.Accordion("UV / texture", open=False): | |
| uv_btn = gr.Button("Export UV layout + texture") | |
| uv_img, tex_img = gr.Image(label="UV layout"), gr.Image(label="Texture") | |
| new_tex = gr.File(label="Upload new texture (PNG, same UV layout)", type="filepath") | |
| apply_btn = gr.Button("Apply texture") | |
| go.click(step_generate, inp, [job, tp, view]) | |
| pose.click(step_pose, job, [view, d_glb, d_obj, d_fbx]) | |
| uv_btn.click(step_uv, job, [uv_img, tex_img]) | |
| apply_btn.click(step_texture, [job, new_tex], [view, d_glb, d_obj, d_fbx]) | |
| demo.queue().launch() |