kimodo-mrp.ccc/demo/main.go
2026-08-26 15:59:58 +01:00

547 lines
20 KiB
Go

// Local Kimodo text-to-motion demo. Generation is serialized so one native
// process owns Vulkan at a time, while the persistent gallery stays readable.
package main
import (
"crypto/rand"
"embed"
"encoding/binary"
"encoding/hex"
"encoding/json"
"flag"
"fmt"
"log"
"math"
"net/http"
"os"
"os/exec"
"path/filepath"
"sort"
"strings"
"sync"
"time"
)
//go:embed index.html
var files embed.FS
//go:embed models.js
var modelUI []byte
//go:embed assets/localai.png
var localAILogo []byte
type animation struct {
ID string `json:"id"`
Prompt string `json:"prompt"`
Frames int `json:"frames"`
DiffusionSteps int `json:"diffusion_steps"`
Seed uint64 `json:"seed"`
CreatedAt string `json:"created_at"`
Status string `json:"status"`
Error string `json:"error,omitempty"`
Kind string `json:"kind"`
Model string `json:"model"`
Segments []promptSegment `json:"segments,omitempty"`
TransitionFrames int `json:"transition_frames,omitempty"`
Progress string `json:"progress,omitempty"`
}
type promptSegment struct {
Prompt string `json:"prompt"`
Frames int `json:"frames"`
}
type motionModel struct {
ID string `json:"id"`
Label string `json:"label"`
Skeleton string `json:"skeleton"`
SkeletonKey string `json:"skeleton_key"`
Upstream string `json:"upstream"`
License string `json:"license"`
LicenseURL string `json:"license_url"`
Commercial bool `json:"commercial"`
Available bool `json:"available"`
Reason string `json:"reason,omitempty"`
Parents []int `json:"parents"`
Offsets [][3]float32 `json:"offsets"`
Motion string `json:"-"`
}
type gallery struct {
mu sync.RWMutex
items map[string]*animation
output string
queue chan string
generator, motion, text string
models map[string]motionModel
}
func token() string {
b := make([]byte, 8)
if _, err := rand.Read(b); err != nil {
panic(err)
}
return hex.EncodeToString(b)
}
func (g *gallery) save(a *animation) error {
b, err := json.MarshalIndent(a, "", " ")
if err != nil {
return err
}
return os.WriteFile(filepath.Join(g.output, a.ID+".json"), b, 0644)
}
func (g *gallery) list() []*animation {
g.mu.RLock()
defer g.mu.RUnlock()
result := make([]*animation, 0, len(g.items))
for _, item := range g.items {
copy := *item
result = append(result, &copy)
}
sort.Slice(result, func(i, j int) bool { return result[i].CreatedAt > result[j].CreatedAt })
return result
}
func readF32(path string) ([]float32, error) {
b, err := os.ReadFile(path)
if err != nil {
return nil, err
}
if len(b)%4 != 0 {
return nil, fmt.Errorf("invalid F32 file: %s", path)
}
values := make([]float32, len(b)/4)
for i := range values {
values[i] = math.Float32frombits(binary.LittleEndian.Uint32(b[i*4:]))
}
return values, nil
}
func writeF32(path string, values []float32) error {
b := make([]byte, len(values)*4)
for i, value := range values {
binary.LittleEndian.PutUint32(b[i*4:], math.Float32bits(value))
}
return os.WriteFile(path, b, 0600)
}
// Each motion is exported as a node-only GLB: it deliberately has no mesh or
// skin, so consumers can attach their own Three.js geometry to the named
// joints. Kimodo stores root translations and local XYZW rotations.
var smplx22Parents = [...]int{-1, 0, 0, 0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 9, 9, 12, 13, 14, 16, 17, 18, 19}
var smplx22Names = [...]string{"pelvis", "left_hip", "right_hip", "spine1", "left_knee", "right_knee", "spine2", "left_ankle", "right_ankle", "spine3", "left_foot", "right_foot", "neck", "left_collar", "right_collar", "head", "left_shoulder", "right_shoulder", "left_elbow", "right_elbow", "left_wrist", "right_wrist"}
var smplx22Offsets = [...][3]float32{{}, {.052299, -.093936, -.027607}, {-.057193, -.106548, -.022218}, {-.001496, .11293, -.024981}, {.058867, -.416442, -.006557}, {-.048074, -.39756, -.014061}, {.0069, .145636, -.006859}, {-.041738, -.437584, -.029512}, {.014489, -.446853, -.01803}, {-.010334, .056082, .021116}, {.049294, -.065279, .126259}, {-.040575, -.065287, .127076}, {-.011026, .171365, -.028827}, {.047725, .087643, -.008375}, {-.046636, .086612, -.014864}, {.024654, .175391, .024463}, {.126285, .05768, -.013885}, {-.109342, .053674, -.009118}, {.272907, -.069853, -.039094}, {-.292029, -.03544, -.024565}, {.276174, .021254, -.002478}, {-.271878, -.004835, -.016445}}
type skeletonDefinition struct {
key string
names []string
parents []int
offsets [][3]float32
}
var skeletonDefinitions = map[string]skeletonDefinition{
"smplx22": {key: "smplx22", names: smplx22Names[:], parents: smplx22Parents[:], offsets: smplx22Offsets[:]},
}
type gltfBufferView struct {
Buffer int `json:"buffer"`
ByteOffset int `json:"byteOffset,omitempty"`
ByteLength int `json:"byteLength"`
}
type gltfAccessor struct {
BufferView int `json:"bufferView"`
ComponentType int `json:"componentType"`
Count int `json:"count"`
Type string `json:"type"`
}
func appendF32(dst []byte, values []float32) []byte {
for _, value := range values {
var b [4]byte
binary.LittleEndian.PutUint32(b[:], math.Float32bits(value))
dst = append(dst, b[:]...)
}
return dst
}
func writeSkeletonGLB(path string, roots, rotations []float32, skeleton skeletonDefinition) error {
frames := len(roots) / 3
joints := len(skeleton.parents)
if frames < 1 || joints < 1 || len(skeleton.names) != joints || len(skeleton.offsets) != joints || len(roots) != frames*3 || len(rotations) != frames*joints*4 {
return fmt.Errorf("invalid %s motion for GLB export", skeleton.key)
}
times := make([]float32, frames)
for i := range times {
times[i] = float32(i) / 30
}
bin := make([]byte, 0, (frames+frames*3+frames*22*4)*4)
views := make([]gltfBufferView, 0, 24)
addView := func(values []float32) int {
offset := len(bin)
bin = appendF32(bin, values)
views = append(views, gltfBufferView{Buffer: 0, ByteOffset: offset, ByteLength: len(bin) - offset})
return len(views) - 1
}
timeView, rootView := addView(times), addView(roots)
rotationViews := make([]int, joints)
for joint := range rotationViews {
track := make([]float32, frames*4)
for frame := 0; frame < frames; frame++ {
copy(track[frame*4:], rotations[(frame*joints+joint)*4:(frame*joints+joint+1)*4])
}
rotationViews[joint] = addView(track)
}
accessors := []gltfAccessor{{BufferView: timeView, ComponentType: 5126, Count: frames, Type: "SCALAR"}, {BufferView: rootView, ComponentType: 5126, Count: frames, Type: "VEC3"}}
for _, view := range rotationViews {
accessors = append(accessors, gltfAccessor{BufferView: view, ComponentType: 5126, Count: frames, Type: "VEC4"})
}
nodes := make([]map[string]any, joints)
for joint := range nodes {
node := map[string]any{"name": skeleton.names[joint]}
if joint != 0 {
node["translation"] = skeleton.offsets[joint]
}
children := make([]int, 0, 3)
for child, parent := range skeleton.parents {
if parent == joint {
children = append(children, child)
}
}
if len(children) != 0 {
node["children"] = children
}
nodes[joint] = node
}
samplers := make([]map[string]any, 0, 23)
channels := make([]map[string]any, 0, 23)
addChannel := func(node, output int, path string) {
samplers = append(samplers, map[string]any{"input": 0, "output": output, "interpolation": "LINEAR"})
channels = append(channels, map[string]any{"sampler": len(samplers) - 1, "target": map[string]any{"node": node, "path": path}})
}
addChannel(0, 1, "translation")
for joint := 0; joint < joints; joint++ {
addChannel(joint, joint+2, "rotation")
}
document := map[string]any{
"asset": map[string]string{"version": "2.0", "generator": "kimodo.cpp skeleton exporter"},
"scene": 0,
"scenes": []map[string]any{{"nodes": []int{0}}},
"nodes": nodes,
"buffers": []map[string]int{{"byteLength": len(bin)}},
"bufferViews": views,
"accessors": accessors,
"animations": []map[string]any{{"name": "KimodoMotion", "samplers": samplers, "channels": channels}},
"extras": map[string]any{"skeleton": skeleton.key, "fps": 30, "rotation_order": "xyzw"},
}
jsonChunk, err := json.Marshal(document)
if err != nil {
return err
}
for len(jsonChunk)%4 != 0 {
jsonChunk = append(jsonChunk, ' ')
}
for len(bin)%4 != 0 {
bin = append(bin, 0)
}
total := 12 + 8 + len(jsonChunk) + 8 + len(bin)
out := make([]byte, 0, total)
putU32 := func(value uint32) {
var b [4]byte
binary.LittleEndian.PutUint32(b[:], value)
out = append(out, b[:]...)
}
putU32(0x46546c67)
putU32(2)
putU32(uint32(total))
putU32(uint32(len(jsonChunk)))
putU32(0x4e4f534a)
out = append(out, jsonChunk...)
putU32(uint32(len(bin)))
putU32(0x004e4942)
out = append(out, bin...)
return os.WriteFile(path, out, 0600)
}
func exportSkeletonGLB(dir, skeletonKey string) error {
skeleton, ok := skeletonDefinitions[skeletonKey]
if !ok {
return fmt.Errorf("unsupported skeleton %q", skeletonKey)
}
roots, err := readF32(filepath.Join(dir, "root_positions.f32"))
if err != nil {
return err
}
rotations, err := readF32(filepath.Join(dir, "local_rotations_xyzw.f32"))
if err != nil {
return err
}
return writeSkeletonGLB(filepath.Join(dir, "animation.glb"), roots, rotations, skeleton)
}
func (g *gallery) worker() {
for id := range g.queue {
g.mu.Lock()
item := g.items[id]
item.Status = "running"
_ = g.save(item)
g.mu.Unlock()
dir := filepath.Join(g.output, id)
err := os.MkdirAll(dir, 0755)
if err == nil {
err = os.WriteFile(filepath.Join(dir, "prompt.txt"), []byte(item.Prompt), 0600)
}
if err == nil {
model, ok := g.models[item.Model]
if !ok || !model.Available {
err = fmt.Errorf("model %q is not available", item.Model)
} else {
segments := item.Segments
if len(segments) == 0 {
segments = []promptSegment{{Prompt: item.Prompt, Frames: item.Frames}}
}
args := []string{model.Motion, g.text, "--sequence", fmt.Sprint(item.TransitionFrames), fmt.Sprint(item.DiffusionSteps), fmt.Sprint(item.Seed), dir}
for index, segment := range segments {
promptPath := filepath.Join(dir, fmt.Sprintf("segment-%02d.txt", index+1))
if err = os.WriteFile(promptPath, []byte(segment.Prompt), 0600); err != nil {
break
}
args = append(args, fmt.Sprint(segment.Frames), promptPath)
}
if err == nil {
g.mu.Lock()
item.Progress = fmt.Sprintf("Generating %d conditioned segments", len(segments))
_ = g.save(item)
g.mu.Unlock()
cmd := exec.Command(g.generator, args...)
cmd.Env = append(os.Environ(), "KIMODO_BACKEND=vulkan")
output, runErr := cmd.CombinedOutput()
if runErr != nil {
err = fmt.Errorf("sequence: %w: %s", runErr, strings.TrimSpace(string(output)))
}
}
if err == nil {
err = exportSkeletonGLB(dir, model.SkeletonKey)
}
}
}
g.mu.Lock()
if err != nil {
item.Status = "failed"
item.Error = err.Error()
} else {
item.Status = "ready"
item.Progress = ""
}
if saveErr := g.save(item); saveErr != nil {
log.Printf("save %s: %v", item.ID, saveErr)
}
g.mu.Unlock()
}
}
func main() {
addr := flag.String("addr", "127.0.0.1:8090", "listen address")
motion := flag.String("motion-model", "models/kimodo-smplx-rp-v1-f32.gguf", "motion GGUF")
somaRP := flag.String("soma-rp-model", "models/kimodo-soma-rp-v1.1-f32.gguf", "SOMA RP v1.1 motion GGUF")
somaSEED := flag.String("soma-seed-model", "models/kimodo-soma-seed-v1.1-f32.gguf", "SOMA SEED v1.1 motion GGUF")
g1RP := flag.String("g1-rp-model", "models/kimodo-g1-rp-v1-f32.gguf", "G1 RP v1 motion GGUF")
g1SEED := flag.String("g1-seed-model", "models/kimodo-g1-seed-v1-f32.gguf", "G1 SEED v1 motion GGUF")
text := flag.String("text-bundle", "generated/llm2vec-text-bundle", "native LLM2Vec component directory")
generator := flag.String("generator", "build/debug/kmd-generate", "native text-to-motion command")
output := flag.String("output", "demo-output", "persistent gallery directory")
flag.Parse()
if err := os.MkdirAll(*output, 0755); err != nil {
log.Fatal(err)
}
makeModel := func(id, label, skeletonLabel, skeletonKey, upstream, license, licenseURL, path string, commercial bool) motionModel {
definition := skeletonDefinitions[skeletonKey]
model := motionModel{ID: id, Label: label, Skeleton: skeletonLabel, SkeletonKey: skeletonKey, Upstream: upstream, License: license, LicenseURL: licenseURL, Commercial: commercial, Parents: definition.parents, Offsets: definition.offsets, Motion: path}
if info, err := os.Stat(path); err == nil && info.Mode().IsRegular() {
model.Available = true
} else {
model.Reason = "GGUF not found at " + path
}
return model
}
const internalLicense = "https://www.nvidia.com/en-us/agreements/enterprise-software/nvidia-internal-scientific-research-and-development-model-license/"
const openLicense = "https://www.nvidia.com/en-us/agreements/enterprise-software/nvidia-open-model-license/"
models := map[string]motionModel{
"smplx-rp-v1": makeModel("smplx-rp-v1", "SMPL-X RP v1", "SMPL-X 22 joints", "smplx22", "nvidia/Kimodo-SMPLX-RP-v1", "NVIDIA Internal Scientific R&D (non-commercial)", internalLicense, *motion, false),
"soma-rp-v1.1": makeModel("soma-rp-v1.1", "SOMA RP v1.1", "SOMA compact 30-joint control skeleton", "soma30", "nvidia/Kimodo-SOMA-RP-v1.1", "NVIDIA Open Model License", openLicense, *somaRP, true),
"soma-seed-v1.1": makeModel("soma-seed-v1.1", "SOMA SEED v1.1", "SOMA compact 30-joint control skeleton", "soma30", "nvidia/Kimodo-SOMA-SEED-v1.1", "NVIDIA Open Model License", openLicense, *somaSEED, true),
"g1-rp-v1": makeModel("g1-rp-v1", "G1 RP v1", "Unitree G1 34 joints", "g1skel34", "nvidia/Kimodo-G1-RP-v1", "NVIDIA Open Model License", openLicense, *g1RP, true),
"g1-seed-v1": makeModel("g1-seed-v1", "G1 SEED v1", "Unitree G1 34 joints", "g1skel34", "nvidia/Kimodo-G1-SEED-v1", "NVIDIA Open Model License", openLicense, *g1SEED, true),
}
g := &gallery{items: map[string]*animation{}, output: *output, queue: make(chan string, 32), generator: *generator, motion: *motion, text: *text, models: models}
entries, _ := filepath.Glob(filepath.Join(*output, "*.json"))
for _, path := range entries {
b, err := os.ReadFile(path)
if err != nil {
continue
}
var a animation
if json.Unmarshal(b, &a) == nil {
g.items[a.ID] = &a
if a.Status == "ready" {
model, ok := models[a.Model]
if !ok {
model = models["smplx-rp-v1"]
}
if err := exportSkeletonGLB(filepath.Join(*output, a.ID), model.SkeletonKey); err != nil && !os.IsNotExist(err) {
log.Printf("export existing animation %s: %v", a.ID, err)
}
}
}
}
go g.worker()
index, err := files.ReadFile("index.html")
if err != nil {
log.Fatal(err)
}
mux := http.NewServeMux()
mux.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) {
if r.URL.Path != "/" {
http.NotFound(w, r)
return
}
w.Header().Set("Content-Type", "text/html; charset=utf-8")
_, _ = w.Write(index)
})
mux.HandleFunc("/localai.png", func(w http.ResponseWriter, r *http.Request) {
w.Header().Set("Content-Type", "image/png")
w.Header().Set("Cache-Control", "public, max-age=86400")
_, _ = w.Write(localAILogo)
})
mux.HandleFunc("/models.js", func(w http.ResponseWriter, r *http.Request) {
w.Header().Set("Content-Type", "application/javascript; charset=utf-8")
w.Header().Set("Cache-Control", "no-store")
_, _ = w.Write(modelUI)
})
mux.HandleFunc("/api/animations", func(w http.ResponseWriter, r *http.Request) {
w.Header().Set("Content-Type", "application/json")
_ = json.NewEncoder(w).Encode(g.list())
})
mux.HandleFunc("/api/models", func(w http.ResponseWriter, r *http.Request) {
result := make([]motionModel, 0, len(g.models))
for _, model := range g.models {
result = append(result, model)
}
sort.Slice(result, func(i, j int) bool { return result[i].ID < result[j].ID })
w.Header().Set("Content-Type", "application/json")
_ = json.NewEncoder(w).Encode(result)
})
mux.HandleFunc("/api/generate", func(w http.ResponseWriter, r *http.Request) {
if r.Method != http.MethodPost {
w.Header().Set("Allow", http.MethodPost)
http.Error(w, "POST required", http.StatusMethodNotAllowed)
return
}
var request struct {
Prompt string `json:"prompt"`
Segments []promptSegment `json:"segments"`
TransitionFrames int `json:"transition_frames"`
Frames int `json:"frames"`
Steps int `json:"steps"`
Seed uint64 `json:"seed"`
Model string `json:"model"`
}
if err := json.NewDecoder(http.MaxBytesReader(w, r.Body, 32<<10)).Decode(&request); err != nil {
http.Error(w, "invalid JSON", 400)
return
}
request.Prompt = strings.TrimSpace(request.Prompt)
if len(request.Segments) == 0 {
request.Segments = []promptSegment{{Prompt: request.Prompt, Frames: request.Frames}}
}
if len(request.Segments) > 16 {
http.Error(w, "at most 16 prompt segments", 400)
return
}
if request.Frames == 0 {
request.Frames = 150
}
if request.Steps == 0 {
request.Steps = 100
}
for index := range request.Segments {
request.Segments[index].Prompt = strings.TrimSpace(request.Segments[index].Prompt)
if request.Segments[index].Frames == 0 {
request.Segments[index].Frames = 150
}
if request.Segments[index].Prompt == "" || len(request.Segments[index].Prompt) > 4096 || request.Segments[index].Frames < 60 || request.Segments[index].Frames > 150 {
http.Error(w, "each prompt segment must be 60..150 frames and 1..4096 bytes", 400)
return
}
}
if request.TransitionFrames == 0 {
request.TransitionFrames = 5
}
if request.TransitionFrames < 1 || request.TransitionFrames > 60 || request.Steps < 1 || request.Steps > 1000 {
http.Error(w, "transition frames must be 1..60 and steps 1..1000", 400)
return
}
if request.Model == "" {
request.Model = "smplx-rp-v1"
}
model, ok := g.models[request.Model]
if !ok || !model.Available {
http.Error(w, "selected motion model is not available: "+model.Reason, http.StatusConflict)
return
}
totalFrames := 0
for _, segment := range request.Segments {
totalFrames += segment.Frames
}
a := &animation{ID: token(), Prompt: request.Segments[0].Prompt, Frames: totalFrames, DiffusionSteps: request.Steps, Seed: request.Seed, CreatedAt: time.Now().UTC().Format(time.RFC3339), Status: "queued", Kind: "generated", Model: request.Model, Segments: request.Segments, TransitionFrames: request.TransitionFrames}
g.mu.Lock()
g.items[a.ID] = a
err := g.save(a)
g.mu.Unlock()
if err != nil {
http.Error(w, err.Error(), 500)
return
}
g.queue <- a.ID
w.Header().Set("Content-Type", "application/json")
w.WriteHeader(http.StatusAccepted)
_ = json.NewEncoder(w).Encode(a)
})
mux.HandleFunc("/api/animations/", func(w http.ResponseWriter, r *http.Request) {
parts := strings.Split(strings.TrimPrefix(r.URL.Path, "/api/animations/"), "/")
if len(parts) != 2 || (parts[1] != "root.f32" && parts[1] != "rotations.f32" && parts[1] != "animation.glb") {
http.NotFound(w, r)
return
}
g.mu.RLock()
a := g.items[parts[0]]
g.mu.RUnlock()
if a == nil || a.Status != "ready" {
http.NotFound(w, r)
return
}
name := "root_positions.f32"
if parts[1] == "rotations.f32" {
name = "local_rotations_xyzw.f32"
}
if parts[1] == "animation.glb" {
name = "animation.glb"
w.Header().Set("Content-Type", "model/gltf-binary")
w.Header().Set("Content-Disposition", "attachment; filename=kimodo-"+a.ID+".glb")
// A GLB is a compact asset; read it directly so browsers always receive
// it as a download rather than invoking any path-cleaning redirects.
data, err := os.ReadFile(filepath.Join(g.output, a.ID, name))
if err != nil {
http.NotFound(w, r)
return
}
w.Header().Set("Content-Length", fmt.Sprint(len(data)))
_, _ = w.Write(data)
return
} else {
w.Header().Set("Content-Type", "application/octet-stream")
}
w.Header().Set("Cache-Control", "no-store")
http.ServeFile(w, r, filepath.Join(g.output, a.ID, name))
})
log.Printf("Kimodo text-to-motion demo listening at http://%s", *addr)
log.Fatal(http.ListenAndServe(*addr, mux))
}