From 8571c0a09839802002141b4b96920816a25f3403 Mon Sep 17 00:00:00 2001 From: Richard Palethorpe Date: Mon, 24 Aug 2026 07:14:03 +0100 Subject: [PATCH] Condition multi-prompt transitions in native sampler --- demo/main.go | 36 +++++------------ demo/main_test.go | 49 ---------------------- include/kimodo/kimodo.hpp | 9 +++++ src/denoiser.cpp | 50 +++++++++++++++++++++-- src/denoiser.hpp | 9 +++++ src/generate.cpp | 27 ++++++++++++- src/model.cpp | 85 +++++++++++++++++++++++++++++++++++++++ 7 files changed, 186 insertions(+), 79 deletions(-) delete mode 100644 demo/main_test.go diff --git a/demo/main.go b/demo/main.go index f172f38..42a6e88 100644 --- a/demo/main.go +++ b/demo/main.go @@ -219,38 +219,24 @@ func (g *gallery) worker() { if len(segments) == 0 { segments = []promptSegment{{Prompt: item.Prompt, Frames: item.Frames}} } - segmentDirs := make([]string, 0, len(segments)) + 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 { - g.mu.Lock() - item.Progress = fmt.Sprintf("Generating segment %d of %d", index+1, len(segments)) - _ = g.save(item) - g.mu.Unlock() - segmentDir := filepath.Join(dir, fmt.Sprintf("segment-%02d", index+1)) - if err = os.MkdirAll(segmentDir, 0755); err != nil { - break - } - promptPath := filepath.Join(segmentDir, "prompt.txt") + promptPath := filepath.Join(dir, fmt.Sprintf("segment-%02d.txt", index+1)) if err = os.WriteFile(promptPath, []byte(segment.Prompt), 0600); err != nil { break } - cmd := exec.Command(g.generator, model.Motion, g.text, promptPath, fmt.Sprint(segment.Frames), fmt.Sprint(item.DiffusionSteps), fmt.Sprint(item.Seed+uint64(index)), segmentDir) + 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("segment %d: %w: %s", index+1, runErr, strings.TrimSpace(string(output))) - break - } - segmentDirs = append(segmentDirs, segmentDir) - } - if err == nil && len(segmentDirs) > 1 { - err = stitchSegments(dir, segmentDirs, item.TransitionFrames) - } - if err == nil && len(segmentDirs) == 1 { - for _, name := range []string{"root_positions.f32", "local_rotations_xyzw.f32"} { - err = copyFile(filepath.Join(dir, name), filepath.Join(segmentDirs[0], name)) - if err != nil { - break - } + err = fmt.Errorf("sequence: %w: %s", runErr, strings.TrimSpace(string(output))) } } } diff --git a/demo/main_test.go b/demo/main_test.go deleted file mode 100644 index db0ffb3..0000000 --- a/demo/main_test.go +++ /dev/null @@ -1,49 +0,0 @@ -package main - -import ( - "os" - "path/filepath" - "testing" -) - -func TestStitchSegmentsBlendsOverlap(t *testing.T) { - dir := t.TempDir() - first, second := filepath.Join(dir, "first"), filepath.Join(dir, "second") - if err := os.MkdirAll(first, 0755); err != nil { - t.Fatal(err) - } - if err := os.MkdirAll(second, 0755); err != nil { - t.Fatal(err) - } - rot := make([]float32, 2*22*4) - for index := range rot { - if index%4 == 3 { - rot[index] = 1 - } - } - if err := writeF32(filepath.Join(first, "root_positions.f32"), []float32{0, 0, 0, 1, 0, 0}); err != nil { - t.Fatal(err) - } - if err := writeF32(filepath.Join(first, "local_rotations_xyzw.f32"), rot); err != nil { - t.Fatal(err) - } - if err := writeF32(filepath.Join(second, "root_positions.f32"), []float32{3, 0, 0, 5, 0, 0}); err != nil { - t.Fatal(err) - } - if err := writeF32(filepath.Join(second, "local_rotations_xyzw.f32"), rot); err != nil { - t.Fatal(err) - } - if err := stitchSegments(dir, []string{first, second}, 1); err != nil { - t.Fatal(err) - } - root, err := readF32(filepath.Join(dir, "root_positions.f32")) - if err != nil { - t.Fatal(err) - } - if len(root) != 9 { - t.Fatalf("frames = %d, want 3", len(root)/3) - } - if root[3] != 2 { - t.Fatalf("blended boundary = %v, want 2", root[3]) - } -} diff --git a/include/kimodo/kimodo.hpp b/include/kimodo/kimodo.hpp index 37fead7..e774e5a 100644 --- a/include/kimodo/kimodo.hpp +++ b/include/kimodo/kimodo.hpp @@ -6,6 +6,7 @@ #include #include #include +#include #include #include #include @@ -21,6 +22,11 @@ struct motion_data { std::vector root_positions; }; +struct prompt_segment { + std::string prompt; + unsigned frames = 0; +}; + class KIMODO_API model { public: static std::expected, std::string> load( @@ -32,6 +38,9 @@ public: std::expected generate_text( std::string_view utf8_prompt, unsigned frames, unsigned steps, std::uint64_t seed, float text_cfg, float constraint_cfg) const; + std::expected generate_text_sequence( + std::span segments, unsigned transition_frames, + unsigned steps, std::uint64_t seed, float text_cfg, float constraint_cfg) const; ~model(); model(const model &) = delete; model &operator=(const model &) = delete; diff --git a/src/denoiser.cpp b/src/denoiser.cpp index 14548d7..4060db4 100644 --- a/src/denoiser.cpp +++ b/src/denoiser.cpp @@ -71,6 +71,9 @@ ggml_tensor *layer(ggml_context *ctx,ggml_tensor*x,const ggml_motion_weights&w,s a=linear(ctx,a,weight(w,s+"self_attn.out_proj.weight"),weight(w,s+"self_attn.out_proj.bias")); x=norm(ctx,ggml_add(ctx,x,a),weight(w,s+"norm1.weight"),weight(w,s+"norm1.bias")); auto*ff=linear(ctx,x,weight(w,s+"linear1.weight"),weight(w,s+"linear1.bias")); ff=ggml_gelu_erf(ctx,ff); ff=linear(ctx,ff,weight(w,s+"linear2.weight"),weight(w,s+"linear2.bias")); return norm(ctx,ggml_add(ctx,x,ff),weight(w,s+"norm2.weight"),weight(w,s+"norm2.bias")); } } +std::expected, std::string> run_separated_cfg_denoiser_conditioned( + const ggml_motion_weights &, std::span, std::span, + std::span, std::span, float, float, float, float, std::size_t); std::expected, std::string> run_motion_transformer(const ggml_motion_weights&w,std::string_view prefix,std::span motion,size_t motion_dim,std::span embedding,std::span timesteps,std::span headings,size_t batch,size_t frames) try { if(!batch||!frames||motion.size()!=batch*frames*motion_dim||embedding.size()!=batch*4096||timesteps.size()!=batch||headings.size()!=batch) return std::unexpected("invalid Transformer input dimensions"); const int seq=prefix_tokens+static_cast(frames); std::vector text(batch*text_tokens*4096),time(batch*width),angle(batch*2),position(size_t(seq)*width); @@ -115,12 +118,31 @@ std::expected, std::string> run_separated_cfg_denoiser( const ggml_motion_weights &weights, std::span motion, std::span embedding, float timestep, float text_weight, float constraint_weight, std::size_t frames) { + const std::vector empty(frames*273, 0.f); + return run_separated_cfg_denoiser_conditioned(weights, motion, embedding, empty, empty, + timestep, 0.f, text_weight, constraint_weight, frames); +} + +std::expected, std::string> run_separated_cfg_denoiser_conditioned( + const ggml_motion_weights &weights, std::span motion, + std::span embedding, std::span observed, + std::span observed_mask, float timestep, float heading, + float text_weight, float constraint_weight, std::size_t frames) { if (motion.size()!=frames*273 || embedding.size()!=4096 || !std::isfinite(timestep) || !std::isfinite(text_weight) || !std::isfinite(constraint_weight)) return std::unexpected("invalid separated CFG denoiser input"); - constexpr size_t cfg_batch=3; std::vector extended(cfg_batch*frames*546), text(cfg_batch*4096), times(cfg_batch,timestep), headings(cfg_batch), mask(cfg_batch*frames,1.f); - for(size_t b=0;b extended(cfg_batch*frames*546), text(cfg_batch*4096), times(cfg_batch,timestep), headings(cfg_batch,heading), mask(cfg_batch*frames,1.f); + for(size_t b=0;b, std::string> sample_motion_from_noise( } return state; } + +std::expected, std::string> sample_motion_from_noise_conditioned( + const ggml_motion_weights &weights, std::span initial, + std::span embedding, std::span observed, + std::span observed_mask, float heading, std::size_t frames, + unsigned steps, float text_weight, float constraint_weight) { + if(initial.size()!=frames*273 || observed.size()!=initial.size() || observed_mask.size()!=initial.size()) + return std::unexpected("invalid conditioned motion noise dimensions"); + auto schedule=make_cosine_schedule(1000,steps); if(!schedule)return std::unexpected(schedule.error()); + std::vector state(initial.begin(),initial.end()), next(state.size()); + for(unsigned i=steps;i-->0;) { + auto clean=run_separated_cfg_denoiser_conditioned(weights,state,embedding,observed,observed_mask, + float(schedule->use_timesteps[i]),heading,text_weight,constraint_weight,frames); + if(!clean)return std::unexpected(clean.error()); + auto stepped=ddim_step(*schedule,i,state.data(),clean->data(),next.data(),state.size()); + if(!stepped)return std::unexpected(stepped.error()); + state.swap(next); + } + return state; +} } diff --git a/src/denoiser.hpp b/src/denoiser.hpp index fd847ca..b3d3f1b 100644 --- a/src/denoiser.hpp +++ b/src/denoiser.hpp @@ -37,4 +37,13 @@ std::expected, std::string> sample_motion_from_noise( const ggml_motion_weights &weights, std::span initial_noise, std::span embedding, std::size_t frames, unsigned steps, float text_weight, float constraint_weight); + +// Multi-prompt transition sampler. `observed` and `observed_mask` are [T,273] +// normalized motion-representation values/masks. This mirrors the upstream +// concat-mask denoiser: text, constraint, and unconditional CFG branches. +std::expected, std::string> sample_motion_from_noise_conditioned( + const ggml_motion_weights &weights, std::span initial_noise, + std::span embedding, std::span observed, + std::span observed_mask, float first_heading, std::size_t frames, + unsigned steps, float text_weight, float constraint_weight); } diff --git a/src/generate.cpp b/src/generate.cpp index 8c1aabd..4f73168 100644 --- a/src/generate.cpp +++ b/src/generate.cpp @@ -8,6 +8,8 @@ #include #include #include +#include +#include namespace { void write_f32(const std::filesystem::path &path, const std::vector &values) { @@ -20,8 +22,31 @@ void write_f32(const std::filesystem::path &path, const std::vector &valu } int main(int argc, char **argv) try { + if (argc >= 10 && std::string_view(argv[3]) == "--sequence") { + if ((argc - 8) % 2 != 0) throw std::runtime_error("sequence requires FRAME PROMPT.txt pairs"); + const auto transition = static_cast(std::stoul(argv[4])); + const auto steps = static_cast(std::stoul(argv[5])); + const auto seed = static_cast(std::stoull(argv[6])); + std::vector segments; + for (int index=8; index(prompt_file), {}}; + if (!prompt_file && prompt.empty()) throw std::runtime_error("cannot read sequence prompt"); + segments.push_back({prompt, static_cast(std::stoul(argv[index]))}); + } + auto model = kimodo::model::load(argv[1], argv[2]); + if (!model) throw std::runtime_error(model.error()); + auto motion = (*model)->generate_text_sequence(segments, transition, steps, seed, 2.F, 2.F); + if (!motion) throw std::runtime_error(motion.error()); + const std::filesystem::path output(argv[7]); std::filesystem::create_directories(output); + write_f32(output / "root_positions.f32", motion->root_positions); + write_f32(output / "local_rotations_xyzw.f32", motion->local_rotations_xyzw); + std::cout << "generated " << motion->frames << " SMPL-X22 sequence frames\n"; + return 0; + } if (argc != 8) { - std::cerr << "usage: " << argv[0] << " MOTION.gguf TEXT_BUNDLE PROMPT.txt FRAMES STEPS SEED OUTPUT_DIR\n"; + std::cerr << "usage: " << argv[0] << " MOTION.gguf TEXT_BUNDLE PROMPT.txt FRAMES STEPS SEED OUTPUT_DIR\n" + << " or: " << argv[0] << " MOTION.gguf TEXT_BUNDLE --sequence TRANSITION STEPS SEED OUTPUT_DIR FRAME PROMPT.txt [FRAME PROMPT.txt ...]\n"; return 2; } std::ifstream prompt_file(argv[3]); diff --git a/src/model.cpp b/src/model.cpp index 7ddee01..98656df 100644 --- a/src/model.cpp +++ b/src/model.cpp @@ -8,6 +8,7 @@ #endif #include +#include #include namespace kimodo { @@ -94,4 +95,88 @@ std::expected model::generate_embedding( return std::unexpected("Kimodo was built without GGML support"); #endif } + +std::expected model::generate_text_sequence( + std::span segments, unsigned transition_frames, + unsigned steps, std::uint64_t seed, float text_cfg, float constraint_cfg) const { +#ifdef KIMODO_HAVE_GGML + if (!impl_->text) return std::unexpected("model was loaded without a native text bundle"); + if (segments.empty() || segments.size() > 16) return std::unexpected("sequence requires 1..16 prompt segments"); + if (steps == 0 || steps > 1000 || transition_frames == 0 || transition_frames > 60) + return std::unexpected("invalid sequence sampling parameters"); + if (!impl_->weights) { + auto loaded = detail::ggml_motion_weights::load(impl_->motion_path); + if (!loaded) return std::unexpected(loaded.error()); + impl_->weights = std::move(*loaded); + } + auto gm=impl_->weights->f32_values("stats.global_root.mean"), gs=impl_->weights->f32_values("stats.global_root.std"); + auto bm=impl_->weights->f32_values("stats.body.mean"), bs=impl_->weights->f32_values("stats.body.std"); + if (!gm || !gs || !bm || !bs) return std::unexpected("motion GGUF lacks normalization statistics"); + std::mt19937_64 rng(seed); std::normal_distribution normal(0.f, 1.f); + std::vector joined, previous; + for (size_t index=0; index 300) + return std::unexpected("each sequence segment must contain a prompt and have 2..300 frames"); + if (index && transition_frames >= segment.frames) return std::unexpected("transition must be shorter than every following segment"); + auto embedding=impl_->text->encode(segment.prompt); if (!embedding) return std::unexpected(embedding.error()); + std::vector noise(static_cast(segment.frames)*273); for (float &value : noise) value=normal(rng); + std::vector current; + if (index == 0) { + auto sampled=detail::sample_motion_from_noise(*impl_->weights,noise,*embedding,segment.frames,steps,text_cfg,constraint_cfg); + if (!sampled) return std::unexpected(sampled.error()); current=std::move(*sampled); + } else { + // Derived from NVIDIA's Apache-2.0 `_multiprompt` sampler: + // https://github.com/nv-tlabs/kimodo/blob/main/kimodo/model/kimodo_model.py + // Preserve the prior tail as observed motion for the next DDIM + // run, then use its transition frames to replace the old tail. + std::vector observed(noise.size()), observed_mask(noise.size()); + const auto overlap=static_cast(transition_frames); + const auto previous_start=previous.size()-overlap*273; + // FullBodyConstraintSet conditions smooth root, heading, local + // joint positions, and global rotations (203 values), but not + // generated velocities or contact labels. + for (size_t frame=0; frameweights,noise,*embedding,observed,observed_mask,heading,segment.frames,steps,text_cfg,constraint_cfg); + if (!sampled) return std::unexpected(sampled.error()); current=std::move(*sampled); + // `_multiprompt` samples in the translated local coordinates, then + // restores the prior segment's planar smooth-root origin. + for (size_t frame=0; frame(overlap*273),current.end()); + } + if (index == 0) joined=current; + previous=std::move(current); + } + const auto frames=static_cast(joined.size()/273); + auto decoded=detail::decode_smplx22(joined,frames,*gm,*gs,*bm,*bs); + if (!decoded) return std::unexpected(decoded.error()); + motion_data result; result.frames=frames; result.joints=22; + result.local_rotations_xyzw=std::move(decoded->local_xyzw); result.root_positions=std::move(decoded->root_positions); + return result; +#else + (void) segments; (void) transition_frames; (void) steps; (void) seed; (void) text_cfg; (void) constraint_cfg; + return std::unexpected("Kimodo was built without GGML support"); +#endif +} } // namespace kimodo