Spaces:
Build error
Build error
File size: 4,074 Bytes
653077f | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 | 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")
@spaces.GPU(duration=90)
def make_tpose(img):
return edit_pipe(image=img, prompt=PROMPT, num_inference_steps=8,
true_cfg_scale=1.0).images[0]
@spaces.GPU(duration=120)
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() |