426 lines
14 KiB
Go
426 lines
14 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"`
|
|
Upstream string `json:"upstream"`
|
|
Available bool `json:"available"`
|
|
Reason string `json:"reason,omitempty"`
|
|
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, ©)
|
|
}
|
|
sort.Slice(result, func(i, j int) bool { return result[i].CreatedAt > result[j].CreatedAt })
|
|
return result
|
|
}
|
|
|
|
func copyFile(dst, src string) error {
|
|
b, err := os.ReadFile(src)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
return os.WriteFile(dst, b, 0600)
|
|
}
|
|
|
|
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)
|
|
}
|
|
|
|
func blendQuaternion(a, b []float32, alpha float32) {
|
|
dot := a[0]*b[0] + a[1]*b[1] + a[2]*b[2] + a[3]*b[3]
|
|
if dot < 0 {
|
|
for i := range b {
|
|
b[i] = -b[i]
|
|
}
|
|
}
|
|
length := float32(0)
|
|
for i := range a {
|
|
a[i] = alpha*a[i] + (1-alpha)*b[i]
|
|
length += a[i] * a[i]
|
|
}
|
|
if length > 0 {
|
|
length = 1 / float32(math.Sqrt(float64(length)))
|
|
for i := range a {
|
|
a[i] *= length
|
|
}
|
|
}
|
|
}
|
|
|
|
// stitchSegments joins independently sampled demo segments. The overlap is
|
|
// blended in root space and by normalized linear interpolation for quaternions.
|
|
// Native observed-motion conditioning is deliberately a later parity step.
|
|
func stitchSegments(output string, dirs []string, overlap int) error {
|
|
var roots, rotations []float32
|
|
for index, dir := range dirs {
|
|
root, err := readF32(filepath.Join(dir, "root_positions.f32"))
|
|
if err != nil {
|
|
return err
|
|
}
|
|
rot, err := readF32(filepath.Join(dir, "local_rotations_xyzw.f32"))
|
|
if err != nil {
|
|
return err
|
|
}
|
|
frames := len(root) / 3
|
|
if frames == 0 || len(rot) != frames*22*4 {
|
|
return fmt.Errorf("invalid motion segment %d", index+1)
|
|
}
|
|
if index == 0 {
|
|
roots, rotations = root, rot
|
|
continue
|
|
}
|
|
n := overlap
|
|
if n > frames {
|
|
n = frames
|
|
}
|
|
if n > len(roots)/3 {
|
|
n = len(roots) / 3
|
|
}
|
|
for frame := 0; frame < n; frame++ {
|
|
alpha := float32(0.5)
|
|
if n > 1 {
|
|
alpha = 1 - float32(frame)/float32(n-1)
|
|
}
|
|
old := (len(roots)/3 - n + frame) * 3
|
|
newest := frame * 3
|
|
for axis := 0; axis < 3; axis++ {
|
|
roots[old+axis] = alpha*roots[old+axis] + (1-alpha)*root[newest+axis]
|
|
}
|
|
for joint := 0; joint < 22; joint++ {
|
|
oldQ := (len(rotations)/4 - n*22 + frame*22 + joint) * 4
|
|
newQ := (frame*22 + joint) * 4
|
|
blendQuaternion(rotations[oldQ:oldQ+4], append([]float32(nil), rot[newQ:newQ+4]...), alpha)
|
|
}
|
|
}
|
|
roots = append(roots, root[n*3:]...)
|
|
rotations = append(rotations, rot[n*22*4:]...)
|
|
}
|
|
if err := writeF32(filepath.Join(output, "root_positions.f32"), roots); err != nil {
|
|
return err
|
|
}
|
|
return writeF32(filepath.Join(output, "local_rotations_xyzw.f32"), rotations)
|
|
}
|
|
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)))
|
|
}
|
|
}
|
|
}
|
|
}
|
|
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")
|
|
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)
|
|
}
|
|
models := map[string]motionModel{
|
|
"smplx-rp-v1": {ID: "smplx-rp-v1", Label: "SMPL-X RP v1", Skeleton: "SMPL-X 22 joints", Upstream: "nvidia/Kimodo-SMPLX-RP-v1", Available: true, Motion: *motion},
|
|
"soma-rp-v1.1": {ID: "soma-rp-v1.1", Label: "SOMA RP v1.1", Skeleton: "SOMA 30 joints", Upstream: "nvidia/Kimodo-SOMA-RP-v1.1", Reason: "SOMA decoder and GGML conversion are being added"},
|
|
"soma-seed-v1.1": {ID: "soma-seed-v1.1", Label: "SOMA SEED v1.1", Skeleton: "SOMA 30 joints", Upstream: "nvidia/Kimodo-SOMA-SEED-v1.1", Reason: "SOMA decoder and GGML conversion are being added"},
|
|
"g1-rp-v1": {ID: "g1-rp-v1", Label: "G1 RP v1", Skeleton: "Unitree G1 34 joints", Upstream: "nvidia/Kimodo-G1-RP-v1", Reason: "G1 decoder and GGML conversion are being added"},
|
|
"g1-seed-v1": {ID: "g1-seed-v1", Label: "G1 SEED v1", Skeleton: "Unitree G1 34 joints", Upstream: "nvidia/Kimodo-G1-SEED-v1", Reason: "G1 decoder and GGML conversion are being added"},
|
|
}
|
|
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
|
|
}
|
|
}
|
|
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 > 300 {
|
|
http.Error(w, "each prompt segment must be 60..300 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
|
|
}
|
|
totalFrames -= request.TransitionFrames * (len(request.Segments) - 1)
|
|
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") {
|
|
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"
|
|
}
|
|
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))
|
|
}
|