monster/kimodo-mrp.ccc
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Condition multi-prompt transitions in native sampler

Browse Source
This commit is contained in:
Richard Palethorpe 2026-08-24 07:14:03 +01:00
parent ff811e49de
commit 8571c0a098
7 changed files with 186 additions and 79 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

36
demo/main.go
View File

@ -219,38 +219,24 @@ func (g *gallery) worker() {
if len(segments) == 0 { if len(segments) == 0 {
segments = []promptSegment{{Prompt: item.Prompt, Frames: item.Frames}} 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 { for index, segment := range segments {
g.mu.Lock() promptPath := filepath.Join(dir, fmt.Sprintf("segment-%02d.txt", index+1))
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")
if err = os.WriteFile(promptPath, []byte(segment.Prompt), 0600); err != nil { if err = os.WriteFile(promptPath, []byte(segment.Prompt), 0600); err != nil {
break 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") cmd.Env = append(os.Environ(), "KIMODO_BACKEND=vulkan")
output, runErr := cmd.CombinedOutput() output, runErr := cmd.CombinedOutput()
if runErr != nil { if runErr != nil {
err = fmt.Errorf("segment %d: %w: %s", index+1, runErr, strings.TrimSpace(string(output))) err = fmt.Errorf("sequence: %w: %s", 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
}
} }
} }
} }

49
demo/main_test.go
View File

@ -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])
}
}

9
include/kimodo/kimodo.hpp
View File

@ -6,6 +6,7 @@
#include <cstdint> #include <cstdint>
#include <expected> #include <expected>
#include <memory> #include <memory>
#include <span>
#include <string> #include <string>
#include <string_view> #include <string_view>
#include <vector> #include <vector>
@ -21,6 +22,11 @@ struct motion_data {
std::vector<float> root_positions; std::vector<float> root_positions;
}; };
struct prompt_segment {
std::string prompt;
unsigned frames = 0;
};
class KIMODO_API model { class KIMODO_API model {
public: public:
static std::expected<std::unique_ptr<model>, std::string> load( static std::expected<std::unique_ptr<model>, std::string> load(
@ -32,6 +38,9 @@ public:
std::expected<motion_data, std::string> generate_text( std::expected<motion_data, std::string> generate_text(
std::string_view utf8_prompt, unsigned frames, unsigned steps, std::uint64_t seed, std::string_view utf8_prompt, unsigned frames, unsigned steps, std::uint64_t seed,
float text_cfg, float constraint_cfg) const; float text_cfg, float constraint_cfg) const;
std::expected<motion_data, std::string> generate_text_sequence(
std::span<const prompt_segment> segments, unsigned transition_frames,
unsigned steps, std::uint64_t seed, float text_cfg, float constraint_cfg) const;
~model(); ~model();
model(const model &) = delete; model(const model &) = delete;
model &operator=(const model &) = delete; model &operator=(const model &) = delete;

50
src/denoiser.cpp
View File

@ -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")); 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::vector<float>, std::string> run_separated_cfg_denoiser_conditioned(
const ggml_motion_weights &, std::span<const float>, std::span<const float>,
std::span<const float>, std::span<const float>, float, float, float, float, std::size_t);
std::expected<std::vector<float>, std::string> run_motion_transformer(const ggml_motion_weights&w,std::string_view prefix,std::span<const float> motion,size_t motion_dim,std::span<const float> embedding,std::span<const float> timesteps,std::span<const float> headings,size_t batch,size_t frames) try { std::expected<std::vector<float>, std::string> run_motion_transformer(const ggml_motion_weights&w,std::string_view prefix,std::span<const float> motion,size_t motion_dim,std::span<const float> embedding,std::span<const float> timesteps,std::span<const float> 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"); 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<int>(frames); std::vector<float> text(batch*text_tokens*4096),time(batch*width),angle(batch*2),position(size_t(seq)*width); const int seq=prefix_tokens+static_cast<int>(frames); std::vector<float> text(batch*text_tokens*4096),time(batch*width),angle(batch*2),position(size_t(seq)*width);
@ -115,12 +118,31 @@ std::expected<std::vector<float>, std::string> run_separated_cfg_denoiser(
const ggml_motion_weights &weights, std::span<const float> motion, const ggml_motion_weights &weights, std::span<const float> motion,
std::span<const float> embedding, float timestep, float text_weight, std::span<const float> embedding, float timestep, float text_weight,
float constraint_weight, std::size_t frames) { float constraint_weight, std::size_t frames) {
const std::vector<float> 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::vector<float>, std::string> run_separated_cfg_denoiser_conditioned(
const ggml_motion_weights &weights, std::span<const float> motion,
std::span<const float> embedding, std::span<const float> observed,
std::span<const float> 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)) 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"); return std::unexpected("invalid separated CFG denoiser input");
constexpr size_t cfg_batch=3; std::vector<float> extended(cfg_batch*frames*546), text(cfg_batch*4096), times(cfg_batch,timestep), headings(cfg_batch), mask(cfg_batch*frames,1.f); if (observed.size()!=frames*273 || observed_mask.size()!=frames*273 || !std::isfinite(heading))
for(size_t b=0;b<cfg_batch;++b) for(size_t t=0;t<frames;++t) std::memcpy(extended.data()+(b*frames+t)*546,motion.data()+t*273,273*sizeof(float)); return std::unexpected("invalid separated CFG condition dimensions");
// Upstream order is text, constraint, unconditional. No constraints constexpr size_t cfg_batch=3; std::vector<float> extended(cfg_batch*frames*546), text(cfg_batch*4096), times(cfg_batch,timestep), headings(cfg_batch,heading), mask(cfg_batch*frames,1.f);
// means all motion-mask channels are zero; only batch zero has text. for(size_t b=0;b<cfg_batch;++b) for(size_t t=0;t<frames;++t) {
auto *dst=extended.data()+(b*frames+t)*546;
std::memcpy(dst,motion.data()+t*273,273*sizeof(float));
// Upstream separated CFG is [text, constraint, unconditional]. Only
// the constraint branch receives observed motion and its feature mask.
if (b==1) for (size_t d=0;d<273;++d) dst[d]=motion[t*273+d]*(1.f-observed_mask[t*273+d])+observed[t*273+d]*observed_mask[t*273+d];
if (b==1) std::memcpy(dst+273,observed_mask.data()+t*273,273*sizeof(float));
}
// Only branch zero has text. Branch one is constraint-only; branch two is
// unconditional. This is the upstream separated-CFG batch order.
std::memcpy(text.data(),embedding.data(),4096*sizeof(float)); std::memcpy(text.data(),embedding.data(),4096*sizeof(float));
auto all=run_two_stage_denoiser(weights,extended,text,times,headings,mask,cfg_batch,frames); auto all=run_two_stage_denoiser(weights,extended,text,times,headings,mask,cfg_batch,frames);
if(!all)return std::unexpected(all.error()); if(!all)return std::unexpected(all.error());
@ -145,4 +167,24 @@ std::expected<std::vector<float>, std::string> sample_motion_from_noise(
} }
return state; return state;
} }
std::expected<std::vector<float>, std::string> sample_motion_from_noise_conditioned(
const ggml_motion_weights &weights, std::span<const float> initial,
std::span<const float> embedding, std::span<const float> observed,
std::span<const float> 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<float> 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;
}
} }

9
src/denoiser.hpp
View File

@ -37,4 +37,13 @@ std::expected<std::vector<float>, std::string> sample_motion_from_noise(
const ggml_motion_weights &weights, std::span<const float> initial_noise, const ggml_motion_weights &weights, std::span<const float> initial_noise,
std::span<const float> embedding, std::size_t frames, unsigned steps, std::span<const float> embedding, std::size_t frames, unsigned steps,
float text_weight, float constraint_weight); 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::vector<float>, std::string> sample_motion_from_noise_conditioned(
const ggml_motion_weights &weights, std::span<const float> initial_noise,
std::span<const float> embedding, std::span<const float> observed,
std::span<const float> observed_mask, float first_heading, std::size_t frames,
unsigned steps, float text_weight, float constraint_weight);
} }

27
src/generate.cpp
View File

@ -8,6 +8,8 @@
#include <iostream> #include <iostream>
#include <iterator> #include <iterator>
#include <string> #include <string>
#include <string_view>
#include <vector>
namespace { namespace {
void write_f32(const std::filesystem::path &path, const std::vector<float> &values) { void write_f32(const std::filesystem::path &path, const std::vector<float> &values) {
@ -20,8 +22,31 @@ void write_f32(const std::filesystem::path &path, const std::vector<float> &valu
} }
int main(int argc, char **argv) try { 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<unsigned>(std::stoul(argv[4]));
const auto steps = static_cast<unsigned>(std::stoul(argv[5]));
const auto seed = static_cast<std::uint64_t>(std::stoull(argv[6]));
std::vector<kimodo::prompt_segment> segments;
for (int index=8; index<argc; index+=2) {
std::ifstream prompt_file(argv[index+1]);
const std::string prompt{std::istreambuf_iterator<char>(prompt_file), {}};
if (!prompt_file && prompt.empty()) throw std::runtime_error("cannot read sequence prompt");
segments.push_back({prompt, static_cast<unsigned>(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) { 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; return 2;
} }
std::ifstream prompt_file(argv[3]); std::ifstream prompt_file(argv[3]);

85
src/model.cpp
View File

@ -8,6 +8,7 @@
#endif #endif
#include <cmath> #include <cmath>
#include <algorithm>
#include <random> #include <random>
namespace kimodo { namespace kimodo {
@ -94,4 +95,88 @@ std::expected<motion_data, std::string> model::generate_embedding(
return std::unexpected("Kimodo was built without GGML support"); return std::unexpected("Kimodo was built without GGML support");
#endif #endif
} }
std::expected<motion_data, std::string> model::generate_text_sequence(
std::span<const prompt_segment> 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<float> normal(0.f, 1.f);
std::vector<float> joined, previous;
for (size_t index=0; index<segments.size(); ++index) {
const auto &segment=segments[index];
if (segment.prompt.empty() || segment.frames < 2 || segment.frames > 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<float> noise(static_cast<size_t>(segment.frames)*273); for (float &value : noise) value=normal(rng);
std::vector<float> 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<float> observed(noise.size()), observed_mask(noise.size());
const auto overlap=static_cast<size_t>(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; frame<overlap; ++frame) {
std::copy_n(previous.data()+previous_start+frame*273,203,observed.data()+frame*273);
std::fill_n(observed_mask.data()+frame*273,203,1.f);
}
const float origin_x=observed[0]*(*gs)[0]+(*gm)[0];
const float origin_z=observed[2]*(*gs)[2]+(*gm)[2];
for (size_t frame=0; frame<overlap; ++frame) {
auto *row=observed.data()+frame*273;
row[0]=((row[0]*(*gs)[0]+(*gm)[0])-origin_x)/(*gs)[0];
row[2]=((row[2]*(*gs)[2]+(*gm)[2])-origin_z)/(*gs)[2];
}
const auto p=(previous.size()/273-overlap)*273;
const float heading=std::atan2(previous[p+4]*(*gs)[4]+(*gm)[4],previous[p+3]*(*gs)[3]+(*gm)[3]);
auto sampled=detail::sample_motion_from_noise_conditioned(*impl_->weights,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<segment.frames; ++frame) {
auto *row=current.data()+frame*273;
row[0]=((row[0]*(*gs)[0]+(*gm)[0])+origin_x)/(*gs)[0];
row[2]=((row[2]*(*gs)[2]+(*gm)[2])+origin_z)/(*gs)[2];
}
const auto start=joined.size()-overlap*273;
for (size_t frame=0; frame<overlap; ++frame) {
const float alpha=overlap==1?.5f:1.f-float(frame)/float(overlap-1);
for (size_t d=0; d<273; ++d) joined[start+frame*273+d]=alpha*joined[start+frame*273+d]+(1.f-alpha)*current[frame*273+d];
}
joined.insert(joined.end(),current.begin()+static_cast<std::ptrdiff_t>(overlap*273),current.end());
}
if (index == 0) joined=current;
previous=std::move(current);
}
const auto frames=static_cast<unsigned>(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 } // namespace kimodo