diff --git a/docs/abot_world.md b/docs/abot_world.md new file mode 100644 index 000000000..a94b66218 --- /dev/null +++ b/docs/abot_world.md @@ -0,0 +1,43 @@ +# ABot-World + +[ABot-World-0-5B-LF](https://huggingface.co/acvlab/ABot-World-0-5B-LF) (AMAP CV +Lab, Apache-2.0) is a **causal, interactive** derivative of Wan2.2-TI2V-5B: it +generates an explorable world block-by-block, driven by keyboard actions, using +a rolling KV cache and a 4-step distilled sampler. + +## Status in this repo + +**Supported:** model detection, architecture instantiation, and full weight +loading (F16 / Q8_0 GGUF or safetensors), including the extra +`act_control_adapter.*` action-conditioning tensors. The model reuses the +Wan2.2-TI2V latent space (48ch, 16x) and VAE. + +**Not yet supported:** actual generation. ABot-World cannot run through the +batch paths — those are bidirectional/one-shot, whereas ABot needs a stateful +causal session (KV cache, per-block action injection, its distilled 4-step +schedule). Both batch entrypoints reject ABot models with a clear error at the +shared `GenerationRequest` stage, and the capability queries +(`sd_ctx_supports_image_generation` / `sd_ctx_supports_video_generation`) +report `false`, so front-ends pre-screen it consistently. The interactive +session API is a planned follow-up. + +## Detection + +A GGUF/safetensors checkpoint is classified `VERSION_ABOT_WORLD` when it is a +Wan model that also contains `model.diffusion_model.act_control_adapter.conv.weight`. + +## GGUF conversion + +The DiT checkpoint converts to GGUF with standard tooling. Notes: +- `patch_embedding.weight` is a Conv3d `[3072,48,1,2,2]`; stored 4D as + `[147456,1,2,2]` (ggml/ComfyUI-GGUF Conv3d convention). +- The 6 `act_control_adapter.*` tensors are kept F16 (convs); everything else + follows the usual F16/Q8_0 split. +- Q8_0 (5.86 GB) vs F16 (10.55 GB): validated near-lossless — per-block latent + cosine >= 0.9998 vs the bf16 reference over golden rollouts. + +## Validation + +`script/validate_abot_world.sh` loads an ABot GGUF, asserts detection + clean +tensor load + the guarded rejection, and (optionally, with `WAN_*` env vars) +confirms a stock Wan model is unaffected. diff --git a/examples/CMakeLists.txt b/examples/CMakeLists.txt index 2dcd1d53a..3bb2951d1 100644 --- a/examples/CMakeLists.txt +++ b/examples/CMakeLists.txt @@ -1,4 +1,12 @@ include_directories(${CMAKE_CURRENT_SOURCE_DIR}) add_subdirectory(cli) -add_subdirectory(server) \ No newline at end of file +add_subdirectory(server) +if(NOT SD_BUILD_SHARED_LIBS) + # ABot walk validation harnesses link against engine internals + # (ModelLoader, SDBackendManager, the runner classes), which are not + # exported from the shared library — static builds only. + add_subdirectory(abot-parity) + add_subdirectory(abot-session) + add_subdirectory(abot-walk) +endif() \ No newline at end of file diff --git a/examples/abot-parity/CMakeLists.txt b/examples/abot-parity/CMakeLists.txt new file mode 100644 index 000000000..6fc85fcc4 --- /dev/null +++ b/examples/abot-parity/CMakeLists.txt @@ -0,0 +1,13 @@ +set(TARGET sd-abot-parity) + +add_executable(${TARGET} main.cpp) +if(APPLE) + sd_set_macos_rpaths(${TARGET}) +endif() +target_include_directories(${TARGET} PRIVATE + "${CMAKE_CURRENT_SOURCE_DIR}/.." + "${PROJECT_SOURCE_DIR}/src" +) +install(TARGETS ${TARGET} RUNTIME) +target_link_libraries(${TARGET} PRIVATE stable-diffusion zip ${CMAKE_THREAD_LIBS_INIT}) +target_compile_features(${TARGET} PUBLIC c_std_11 cxx_std_17) diff --git a/examples/abot-parity/main.cpp b/examples/abot-parity/main.cpp new file mode 100644 index 000000000..80b641ece --- /dev/null +++ b/examples/abot-parity/main.cpp @@ -0,0 +1,280 @@ +// sd-abot-parity — per-step parity harness for the ABot-World causal walk. +// +// Loads the ABot DiT GGUF + a scene pack + one golden denoise-step input +// (xt latents, per-frame timesteps) dumped by the reference implementation +// (ABot-World/tools/dump_golden_steps.py), runs one causal forward with the +// recompute formulation, and writes the resulting x0 prediction for the block +// as raw float32 for offline comparison against the golden x0. +// +// Usage: +// sd-abot-parity --dit --scene --golden +// --block --step [--threads N] [--out out.bin] +// +// History frames are the golden final latents of blocks < b (clean, t=0), +// actions replayed from the golden walk.json. + +#include +#include +#include +#include +#include +#include + +#include "abot_world.hpp" +#include "ggml_extend_backend.h" +#include "model.h" + +// ---- minimal .npy reader (float32/float64/int64, C-order) ---- +static bool load_npy(const std::string& path, std::vector& out, std::vector& shape) { + std::ifstream f(path, std::ios::binary); + if (!f) { + fprintf(stderr, "npy: cannot open %s\n", path.c_str()); + return false; + } + char magic[6]; + f.read(magic, 6); + if (memcmp(magic, "\x93NUMPY", 6) != 0) { + return false; + } + uint8_t ver[2]; + f.read(reinterpret_cast(ver), 2); + uint32_t hlen = 0; + if (ver[0] == 1) { + uint16_t h16; + f.read(reinterpret_cast(&h16), 2); + hlen = h16; + } else { + f.read(reinterpret_cast(&hlen), 4); + } + std::string header(hlen, '\0'); + f.read(header.data(), hlen); + + bool f4 = header.find("= header.size() || header[sp] == ')') { + break; + } + shape.push_back(strtoll(header.c_str() + sp, nullptr, 10)); + while (sp < header.size() && header[sp] != ',' && header[sp] != ')') { + sp++; + } + } + int64_t n = 1; + for (auto s : shape) { + n *= s; + } + out.resize(static_cast(n)); + if (f4) { + f.read(reinterpret_cast(out.data()), n * 4); + } else if (f8) { + std::vector tmp(static_cast(n)); + f.read(reinterpret_cast(tmp.data()), n * 8); + for (int64_t i = 0; i < n; i++) { + out[static_cast(i)] = static_cast(tmp[static_cast(i)]); + } + } else { + std::vector tmp(static_cast(n)); + f.read(reinterpret_cast(tmp.data()), n * 8); + for (int64_t i = 0; i < n; i++) { + out[static_cast(i)] = static_cast(tmp[static_cast(i)]); + } + } + return f.good(); +} + +int main(int argc, char** argv) { + ggml_log_set(ggml_log_callback_default, nullptr); + std::string dit, scene_path, golden, out_path = "abot_parity_x0.bin"; + int block = 0, step = 0, n_threads = 8; + std::string backend_spec = "cpu"; + for (int i = 1; i < argc; i++) { + std::string k = argv[i]; + auto next = [&]() -> std::string { return (i + 1 < argc) ? argv[++i] : std::string(); }; + if (k == "--dit") dit = next(); + else if (k == "--scene") scene_path = next(); + else if (k == "--golden") golden = next(); + else if (k == "--block") block = std::stoi(next()); + else if (k == "--step") step = std::stoi(next()); + else if (k == "--threads") n_threads = std::stoi(next()); + else if (k == "--backend") backend_spec = next(); + else if (k == "--out") out_path = next(); + } + if (dit.empty() || scene_path.empty() || golden.empty()) { + fprintf(stderr, "usage: sd-abot-parity --dit x.gguf --scene scene.safetensors --golden dir --block b --step s\n"); + return 2; + } + + // ---- backend + weights ---- + SDBackendManager backend_manager; + std::string err; + if (!backend_manager.init(backend_spec.c_str(), nullptr, false, false, false, false, &err)) { + fprintf(stderr, "backend init failed: %s\n", err.c_str()); + return 1; + } + ggml_backend_t be = backend_manager.runtime_backend(SDBackendModule::DIFFUSION); + ggml_backend_t pbe = backend_manager.params_backend(SDBackendModule::DIFFUSION); + + ModelLoader ml; + if (!ml.init_from_file(dit, "model.diffusion_model.")) { + fprintf(stderr, "cannot open DiT: %s\n", dit.c_str()); + return 1; + } + + ABOT::AbotWorldConfig cfg; // window 8 / 3 fpb — matches the golden test config + ABOT::AbotWorldRunner runner(be, pbe, ml.get_tensor_storage_map(), "model.diffusion_model.", cfg); + if (!runner.alloc_params_buffer()) { + fprintf(stderr, "alloc params failed\n"); + return 1; + } + std::map tensors; + runner.get_param_tensors(tensors, "model.diffusion_model"); + if (!ml.load_tensors(tensors, {}, n_threads)) { + fprintf(stderr, "load DiT tensors failed\n"); + return 1; + } + printf("DiT loaded (%zu tensors)\n", tensors.size()); + + if (!runner.scene.load(scene_path)) { + return 1; + } + printf("scene pack loaded (ref slots: %d)\n", runner.scene.ref_slots); + + // ---- golden inputs ---- + std::vector xt; + std::vector xt_shape; // [1, F, C, H, W] + char name[256]; + snprintf(name, sizeof(name), "%s/b%d_s%d_xt.npy", golden.c_str(), block, step); + if (!load_npy(name, xt, xt_shape)) { + return 1; + } + std::vector tsv; + std::vector ts_shape; + snprintf(name, sizeof(name), "%s/b%d_s%d_t.npy", golden.c_str(), block, step); + if (!load_npy(name, tsv, ts_shape)) { + return 1; + } + const int Fb = static_cast(xt_shape[1]); + const int C = static_cast(xt_shape[2]); + const int H = static_cast(xt_shape[3]); + const int W = static_cast(xt_shape[4]); + runner.lat_w = W; + runner.lat_h = H; + runner.lat_c = C; + + // walk.json actions: rows of 8 ints (one row per block, W..L bit order). + // Tiny hand parser: after "actions", every inner [...] is one row until the + // outer array closes. + std::vector block_actions; + { + std::ifstream jf(golden + "/walk.json"); + std::string js((std::istreambuf_iterator(jf)), std::istreambuf_iterator()); + size_t p = js.find("\"actions\""); + size_t outer = js.find('[', p); + size_t pos = outer + 1; + while (true) { + size_t rs = js.find_first_of("[]", pos); + if (rs == std::string::npos || js[rs] == ']') { + break; // outer array closed + } + size_t re = js.find(']', rs); + if (re == std::string::npos) { + break; + } + uint8_t m = 0; + int bit = 0; + for (size_t q = rs + 1; q < re && bit < 8; q++) { + if (js[q] == '0' || js[q] == '1') { + m |= static_cast((js[q] - '0') << bit); + bit++; + } + } + block_actions.push_back(m); + pos = re + 1; + } + } + printf("parsed %zu action rows\n", block_actions.size()); + + // history latents: golden final latents of blocks < block + std::vector> hist_store; + std::vector frames; + std::vector actions; + std::vector tframes; + const size_t frame_elems = static_cast(C) * H * W; + for (int b = 0; b < block; b++) { + std::vector lat; + std::vector lshape; + snprintf(name, sizeof(name), "%s/b%d_final_latent.npy", golden.c_str(), b); + if (!load_npy(name, lat, lshape)) { + return 1; + } + for (int f = 0; f < Fb; f++) { + hist_store.emplace_back(lat.begin() + static_cast(f) * frame_elems, + lat.begin() + static_cast(f + 1) * frame_elems); + actions.push_back(b < static_cast(block_actions.size()) ? block_actions[static_cast(b)] : 0); + tframes.push_back(0.0f); // history at context_noise = 0 + } + } + for (auto& h : hist_store) { + frames.push_back(h.data()); + } + // current block xt frames + for (int f = 0; f < Fb; f++) { + frames.push_back(xt.data() + static_cast(f) * frame_elems); + actions.push_back(block < static_cast(block_actions.size()) ? block_actions[static_cast(block)] : 0); + tframes.push_back(tsv[static_cast(f)]); + } + + printf("forward: F_vis=%zu block_frames=%d t=[", frames.size(), Fb); + for (auto t : tframes) { + printf("%.1f ", t); + } + printf("]\n"); + + auto flow = runner.forward_step(frames, actions, tframes, Fb, n_threads); + if (flow.empty()) { + fprintf(stderr, "forward failed\n"); + return 1; + } + // unpatchify flow tokens {pt*ph*pw*C, tokens} -> [F, C, H, W] and x0 = xt - sigma*flow + const int h_len = H / 2, w_len = W / 2; + std::vector x0(static_cast(Fb) * frame_elems); + for (int f = 0; f < Fb; f++) { + float sigma = tframes[frames.size() - Fb + f] / 1000.0f; + for (int hh = 0; hh < h_len; hh++) { + for (int ww = 0; ww < w_len; ww++) { + const float* tok = flow.data() + ((static_cast(f) * h_len + hh) * w_len + ww) * (4 * C); + for (int ph = 0; ph < 2; ph++) { + for (int pw = 0; pw < 2; pw++) { + for (int c = 0; c < C; c++) { + // token layout: [c fastest, then pw, then ph] (pt*ph*pw*C with C innermost) + float fl = tok[(ph * 2 + pw) * C + c]; + size_t idx = ((static_cast(f) * C + c) * H + (hh * 2 + ph)) * W + (ww * 2 + pw); + x0[idx] = xt[idx] - sigma * fl; + } + } + } + } + } + } + // NOTE: x0 index math above intentionally kept explicit; xt is [F,C,H,W]. + + std::ofstream of(out_path, std::ios::binary); + of.write(reinterpret_cast(x0.data()), static_cast(x0.size() * sizeof(float))); + printf("wrote %s (%zu floats)\n", out_path.c_str(), x0.size()); + return 0; +} diff --git a/examples/abot-session/CMakeLists.txt b/examples/abot-session/CMakeLists.txt new file mode 100644 index 000000000..2739ea934 --- /dev/null +++ b/examples/abot-session/CMakeLists.txt @@ -0,0 +1,14 @@ +set(TARGET sd-abot-session) + +add_executable(${TARGET} main.cpp) +if(APPLE) + sd_set_macos_rpaths(${TARGET}) +endif() +target_include_directories(${TARGET} PRIVATE + "${CMAKE_CURRENT_SOURCE_DIR}/.." + "${PROJECT_SOURCE_DIR}/src" + "${PROJECT_SOURCE_DIR}/thirdparty" +) +install(TARGETS ${TARGET} RUNTIME) +target_link_libraries(${TARGET} PRIVATE stable-diffusion zip ${CMAKE_THREAD_LIBS_INIT}) +target_compile_features(${TARGET} PUBLIC c_std_11 cxx_std_17) diff --git a/examples/abot-session/main.cpp b/examples/abot-session/main.cpp new file mode 100644 index 000000000..9e97a51a7 --- /dev/null +++ b/examples/abot-session/main.cpp @@ -0,0 +1,249 @@ +// sd-abot-session — drive the public sd_abot_session_* C API for a fixed-scene +// interactive walk, writing decoded frames as PNGs. +// +// Modes: +// --mode walk (default) full native pipeline via the C API: DiT walk + +// taehv decode, one action per block. +// --mode decode decode-only cross-check: load latents produced by +// sd-abot-walk (raw f32 [T,C,H,W]) and decode them natively +// block-by-block with the session's overlap logic, so the +// PNGs can be PSNR-compared against the reference (Python) +// decode of the same latents. +// +// Usage: +// sd-abot-session --dit dit.gguf --taehv taew2_2.gguf --scene scene.safetensors +// [--actions idle:1,W:3] [--threads N] [--seed 42] [--outdir frames] +// sd-abot-session --mode decode --taehv taew2_2.gguf --latents walk.bin +// --lat-w 52 --lat-h 30 [--outdir frames] + +#define STB_IMAGE_WRITE_IMPLEMENTATION +#include "stb_image_write.h" + +#include +#include +#include +#include +#include +#include +#include + +#include "stable-diffusion.h" + +#include "abot_world.hpp" +#include "ggml_extend_backend.h" + +static std::vector parse_actions(const std::string& spec) { + // "idle:1,W:3" -> per-block key bitmask (W,A,S,D,I,J,K,L = bits 0..7). + // "idle"/"none" = no keys; any other character must be a valid walk key. + std::vector blocks; + const std::string keys = "WASDIJKL"; + size_t p = 0; + while (p < spec.size()) { + size_t c = spec.find(',', p); + std::string chunk = spec.substr(p, c == std::string::npos ? std::string::npos : c - p); + p = c == std::string::npos ? spec.size() : c + 1; + size_t colon = chunk.find(':'); + std::string key = chunk.substr(0, colon); + int n = colon == std::string::npos ? 1 : std::stoi(chunk.substr(colon + 1)); + uint8_t m = 0; + std::string upper; + for (char ch : key) { + upper.push_back(static_cast(toupper(ch))); + } + if (upper != "IDLE" && upper != "NONE" && !upper.empty()) { + for (char ch : upper) { + size_t bit = keys.find(ch); + if (bit == std::string::npos) { + fprintf(stderr, "unknown walk key '%c' in --actions chunk '%s' (valid: W,A,S,D,I,J,K,L or 'idle')\n", + ch, key.c_str()); + exit(2); + } + m |= static_cast(1u << bit); + } + } + for (int i = 0; i < n; i++) { + blocks.push_back(m); + } + } + return blocks; +} + +static bool write_png(const std::string& path, const uint8_t* rgb, int w, int h) { + return stbi_write_png(path.c_str(), w, h, 3, rgb, w * 3) != 0; +} + +static int run_walk(const std::string& dit, const std::string& taehv, const std::string& scene, + const std::string& actions_spec, const std::string& outdir, + int threads, int64_t seed, const std::string& backend) { + sd_abot_session_params_t params; + sd_abot_session_params_init(¶ms); + params.dit_model_path = dit.c_str(); + params.taehv_path = taehv.c_str(); + params.scene_path = scene.c_str(); + params.backend = backend.c_str(); + params.n_threads = threads; + params.seed = seed; + + sd_abot_session_t* session = sd_abot_session_new(¶ms); + if (session == nullptr) { + fprintf(stderr, "session creation failed\n"); + return 1; + } + + std::vector actions = parse_actions(actions_spec); + int frame_no = 0; + for (size_t b = 0; b < actions.size(); b++) { + int n_frames = 0; + int64_t t0 = ggml_time_ms(); + sd_image_t* fr = sd_abot_session_step(session, actions[b], &n_frames); + int64_t t1 = ggml_time_ms(); + if (fr == nullptr) { + fprintf(stderr, "step %zu failed\n", b); + sd_abot_session_free(session); + return 1; + } + printf("block %zu (mask 0x%02x): %d frames %ux%u in %.1fs\n", + b, actions[b], n_frames, fr[0].width, fr[0].height, (t1 - t0) / 1000.0f); + for (int i = 0; i < n_frames; i++) { + char name[512]; + snprintf(name, sizeof(name), "%s/frame_%04d.png", outdir.c_str(), frame_no++); + if (!write_png(name, fr[i].data, static_cast(fr[i].width), static_cast(fr[i].height))) { + fprintf(stderr, "cannot write %s\n", name); + } + } + sd_abot_session_frames_free(fr, n_frames); + } + sd_abot_session_free(session); + printf("done: %d frames -> %s\n", frame_no, outdir.c_str()); + return 0; +} + +static int run_decode(const std::string& taehv, const std::string& latents_path, + int lat_w, int lat_h, int lat_c, int fpb, + const std::string& outdir, int threads, const std::string& backend) { + std::ifstream f(latents_path, std::ios::binary); + if (!f) { + fprintf(stderr, "cannot open %s\n", latents_path.c_str()); + return 1; + } + f.seekg(0, std::ios::end); + size_t bytes = static_cast(f.tellg()); + f.seekg(0); + std::vector lat(bytes / sizeof(float)); + f.read(reinterpret_cast(lat.data()), static_cast(bytes)); + const size_t fel = static_cast(lat_c) * lat_h * lat_w; + const int T = static_cast(lat.size() / fel); + printf("latents: %d frames %dx%dx%d\n", T, lat_w, lat_h, lat_c); + + SDBackendManager bm; + std::string err; + if (!bm.init(backend.c_str(), nullptr, false, false, false, false, &err)) { + fprintf(stderr, "backend init failed: %s\n", err.c_str()); + return 1; + } + ModelLoader tl; + if (!tl.init_from_file(taehv, "tae.")) { + return 1; + } + TinyVideoAutoEncoder tae(bm.runtime_backend(SDBackendModule::VAE), + bm.params_backend(SDBackendModule::VAE), + tl.get_tensor_storage_map(), "decoder", true, VERSION_ABOT_WORLD); + if (!tae.alloc_params_buffer()) { + return 1; + } + std::map tensors; + tae.get_param_tensors(tensors, "tae"); + if (!tl.load_tensors(tensors, {}, threads)) { + return 1; + } + printf("taehv loaded (is_wide path auto-detected)\n"); + + int frame_no = 0; + for (int b0 = 0; b0 < T; b0 += fpb) { + const int lead = std::min(b0, 3); + const int Tb = std::min(fpb, T - b0) + lead; + sd::Tensor z({lat_w, lat_h, Tb, lat_c}); + for (int t = 0; t < Tb; t++) { + const float* src = lat.data() + static_cast(b0 - lead + t) * fel; + for (int c = 0; c < lat_c; c++) { + memcpy(z.data() + (static_cast(c) * Tb + t) * lat_w * lat_h, + src + static_cast(c) * lat_w * lat_h, + static_cast(lat_w) * lat_h * sizeof(float)); + } + } + sd_tiling_params_t no_tiling = {}; + sd::Tensor px = tae.decode(threads, z, no_tiling, true); + if (px.empty()) { + fprintf(stderr, "decode failed at block %d\n", b0 / fpb); + return 1; + } + const int64_t Wp = px.shape()[0], Hp = px.shape()[1], Tp = px.shape()[2]; + const int64_t keep = lead > 0 ? 4 * (Tb - lead) : Tp; + const int64_t start = Tp - keep; + std::vector rgb(static_cast(Wp) * Hp * 3); + for (int64_t t = 0; t < keep; t++) { + for (int64_t ch = 0; ch < 3; ch++) { + const float* sp = px.data() + (ch * Tp + start + t) * Wp * Hp; + for (int64_t i = 0; i < Wp * Hp; i++) { + float v = sp[i] < 0.0f ? 0.0f : (sp[i] > 1.0f ? 1.0f : sp[i]); + rgb[static_cast(i) * 3 + ch] = static_cast(v * 255.0f + 0.5f); + } + } + char name[512]; + snprintf(name, sizeof(name), "%s/frame_%04d.png", outdir.c_str(), frame_no++); + if (!write_png(name, rgb.data(), static_cast(Wp), static_cast(Hp))) { + fprintf(stderr, "cannot write %s\n", name); + } + } + printf("block %d: +%lld px frames\n", b0 / fpb, static_cast(keep)); + } + printf("done: %d frames -> %s\n", frame_no, outdir.c_str()); + return 0; +} + +static void sd_log_to_stderr(enum sd_log_level_t level, const char* text, void* data) { + (void)data; + if (level >= SD_LOG_INFO) { + fputs(text, stderr); + fflush(stderr); + } +} + +int main(int argc, char** argv) { + ggml_log_set(ggml_log_callback_default, nullptr); + sd_set_log_callback(sd_log_to_stderr, nullptr); + std::string mode = "walk", dit, taehv, scene, latents, outdir = "."; + std::string actions = "idle:1,W:3", backend = "cpu"; + int threads = 8, lat_w = 52, lat_h = 30, lat_c = 48, fpb = 3; + int64_t seed = 42; + for (int i = 1; i < argc; i++) { + std::string k = argv[i]; + auto next = [&]() -> std::string { return (i + 1 < argc) ? argv[++i] : std::string(); }; + if (k == "--mode") mode = next(); + else if (k == "--dit") dit = next(); + else if (k == "--taehv") taehv = next(); + else if (k == "--scene") scene = next(); + else if (k == "--latents") latents = next(); + else if (k == "--actions") actions = next(); + else if (k == "--outdir") outdir = next(); + else if (k == "--threads") threads = std::stoi(next()); + else if (k == "--seed") seed = std::stoll(next()); + else if (k == "--backend") backend = next(); + else if (k == "--lat-w") lat_w = std::stoi(next()); + else if (k == "--lat-h") lat_h = std::stoi(next()); + else if (k == "--lat-c") lat_c = std::stoi(next()); + else if (k == "--fpb") fpb = std::stoi(next()); + } + if (mode == "decode") { + if (taehv.empty() || latents.empty()) { + fprintf(stderr, "decode mode needs --taehv and --latents\n"); + return 2; + } + return run_decode(taehv, latents, lat_w, lat_h, lat_c, fpb, outdir, threads, backend); + } + if (dit.empty() || taehv.empty() || scene.empty()) { + fprintf(stderr, "walk mode needs --dit, --taehv and --scene\n"); + return 2; + } + return run_walk(dit, taehv, scene, actions, outdir, threads, seed, backend); +} diff --git a/examples/abot-walk/CMakeLists.txt b/examples/abot-walk/CMakeLists.txt new file mode 100644 index 000000000..35bb90eeb --- /dev/null +++ b/examples/abot-walk/CMakeLists.txt @@ -0,0 +1,13 @@ +set(TARGET sd-abot-walk) + +add_executable(${TARGET} main.cpp) +if(APPLE) + sd_set_macos_rpaths(${TARGET}) +endif() +target_include_directories(${TARGET} PRIVATE + "${CMAKE_CURRENT_SOURCE_DIR}/.." + "${PROJECT_SOURCE_DIR}/src" +) +install(TARGETS ${TARGET} RUNTIME) +target_link_libraries(${TARGET} PRIVATE stable-diffusion zip ${CMAKE_THREAD_LIBS_INIT}) +target_compile_features(${TARGET} PUBLIC c_std_11 cxx_std_17) diff --git a/examples/abot-walk/main.cpp b/examples/abot-walk/main.cpp new file mode 100644 index 000000000..05bca8efb --- /dev/null +++ b/examples/abot-walk/main.cpp @@ -0,0 +1,353 @@ +// sd-abot-walk — full causal walk through a fixed ABot-World scene. +// +// Runs the reference block loop natively: for each block of +// num_frame_per_block latent frames, 4 warped flow-matching denoise steps +// (x0 = xt - sigma*flow; renoise to the next step), history kept clean at +// t=0, actions injected per block. Two noise modes: +// --golden replay the reference's recorded noise (deterministic +// full-chain validation; final latents must match the +// golden b{b}_final_latent.npy) +// (default) xorshift/Box-Muller RNG with --seed (free walk) +// +// Output: raw f32 latents [T,C,H,W] (torch order) to --out; decode with +// taehv (native, once the taew GGUF is present) or the reference tools. +// +// Usage: +// sd-abot-walk --dit dit.gguf --scene scene.safetensors +// [--golden golden_steps_dir] [--actions idle:1,W:3] +// [--blocks N] [--threads N] [--out walk_latents.bin] + +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "abot_world.hpp" +#include "ggml_extend_backend.h" +#include "model.h" + +static bool load_npy_f32(const std::string& path, std::vector& out, std::vector& shape); + +// xorshift128+ Box-Muller normal RNG (portable, seed-stable across platforms) +struct Rng { + uint64_t s0, s1; + bool have_spare = false; + float spare = 0.0f; + explicit Rng(uint64_t seed) { + s0 = seed ^ 0x9E3779B97F4A7C15ULL; + s1 = (seed << 1) | 1; + for (int i = 0; i < 16; i++) { + next_u64(); + } + } + uint64_t next_u64() { + uint64_t x = s0, y = s1; + s0 = y; + x ^= x << 23; + s1 = x ^ y ^ (x >> 17) ^ (y >> 26); + return s1 + y; + } + float uniform() { + return static_cast((next_u64() >> 11) * (1.0 / 9007199254740992.0)); + } + float normal() { + if (have_spare) { + have_spare = false; + return spare; + } + float u1 = std::max(uniform(), 1e-12f), u2 = uniform(); + float r = std::sqrt(-2.0f * std::log(u1)), a = 6.28318530718f * u2; + spare = r * std::sin(a); + have_spare = true; + return r * std::cos(a); + } +}; + +static std::vector parse_actions(const std::string& spec) { + // "idle:1,W:3,WJ:2" -> per-block key bitmask (W,A,S,D,I,J,K,L = bits 0..7). + // "idle" / "none" mean no keys held; otherwise every character of the + // chunk must be a walk key (so a typo cannot silently become phantom keys). + std::vector blocks; + const std::string keys = "WASDIJKL"; + size_t p = 0; + while (p < spec.size()) { + size_t c = spec.find(',', p); + std::string chunk = spec.substr(p, c == std::string::npos ? std::string::npos : c - p); + p = c == std::string::npos ? spec.size() : c + 1; + size_t colon = chunk.find(':'); + std::string key = chunk.substr(0, colon); + int n = colon == std::string::npos ? 1 : std::stoi(chunk.substr(colon + 1)); + uint8_t m = 0; + std::string upper; + for (char ch : key) { + upper.push_back(static_cast(toupper(ch))); + } + if (upper != "IDLE" && upper != "NONE" && !upper.empty()) { + for (char ch : upper) { + size_t bit = keys.find(ch); + if (bit == std::string::npos) { + fprintf(stderr, "unknown walk key '%c' in --actions chunk '%s' (valid: W,A,S,D,I,J,K,L or 'idle')\n", + ch, key.c_str()); + exit(2); + } + m |= static_cast(1u << bit); + } + } + for (int i = 0; i < n; i++) { + blocks.push_back(m); + } + } + return blocks; +} + +int main(int argc, char** argv) { + ggml_log_set(ggml_log_callback_default, nullptr); + std::string dit, scene_path, golden, out_path = "walk_latents.bin", actions_spec = "idle:1,W:3"; + int n_threads = 8, blocks_n = -1; + uint64_t seed = 42; + std::string backend_spec = "cpu"; + for (int i = 1; i < argc; i++) { + std::string k = argv[i]; + auto next = [&]() -> std::string { return (i + 1 < argc) ? argv[++i] : std::string(); }; + if (k == "--dit") dit = next(); + else if (k == "--scene") scene_path = next(); + else if (k == "--golden") golden = next(); + else if (k == "--actions") actions_spec = next(); + else if (k == "--blocks") blocks_n = std::stoi(next()); + else if (k == "--threads") n_threads = std::stoi(next()); + else if (k == "--seed") seed = std::stoull(next()); + else if (k == "--backend") backend_spec = next(); + else if (k == "--out") out_path = next(); + } + if (dit.empty() || scene_path.empty()) { + fprintf(stderr, "usage: sd-abot-walk --dit x.gguf --scene scene.safetensors [--golden dir] [--actions a] [--out o.bin]\n"); + return 2; + } + + SDBackendManager backend_manager; + std::string err; + if (!backend_manager.init(backend_spec.c_str(), nullptr, false, false, false, false, &err)) { + fprintf(stderr, "backend init failed: %s\n", err.c_str()); + return 1; + } + ModelLoader ml; + if (!ml.init_from_file(dit, "model.diffusion_model.")) { + return 1; + } + ABOT::AbotWorldConfig cfg; + ABOT::AbotWorldRunner runner(backend_manager.runtime_backend(SDBackendModule::DIFFUSION), + backend_manager.params_backend(SDBackendModule::DIFFUSION), + ml.get_tensor_storage_map(), "model.diffusion_model.", cfg); + if (!runner.alloc_params_buffer()) { + return 1; + } + std::map tensors; + runner.get_param_tensors(tensors, "model.diffusion_model"); + if (!ml.load_tensors(tensors, {}, n_threads)) { + return 1; + } + if (!runner.scene.load(scene_path)) { + return 1; + } + auto ffl = runner.scene.first_frame_latents.shape(); + runner.lat_w = ffl[0]; + runner.lat_h = ffl[1]; + runner.lat_c = ffl[2]; + const int W = static_cast(runner.lat_w), H = static_cast(runner.lat_h), C = static_cast(runner.lat_c); + const size_t fel = static_cast(C) * H * W; + const int Fb = cfg.num_frame_per_block; + printf("scene: latent %dx%dx%d, ref slots %d\n", W, H, C, runner.scene.ref_slots); + + std::vector acts = parse_actions(actions_spec); + if (blocks_n > 0) { + acts.resize(static_cast(blocks_n), acts.empty() ? 0 : acts.back()); + } + Rng rng(seed); + + // walk state (torch-order [C,H,W] frames) + std::vector> hist; + std::vector hist_act; + + // first_frame latent in torch order: scene tensor is {W,H,C} ggml == torch [C,H,W] flat ✓ + std::vector ff(runner.scene.first_frame_latents.data(), + runner.scene.first_frame_latents.data() + fel); + + for (size_t b = 0; b < acts.size(); b++) { + const bool first = b == 0; + // ---- init xt for the block ---- + std::vector xt(static_cast(Fb) * fel); + if (!golden.empty()) { + std::vector g; + std::vector gs; + char nm[256]; + snprintf(nm, sizeof(nm), "%s/b%zu_s0_xt.npy", golden.c_str(), b); + if (!load_npy_f32(nm, g, gs)) { + return 1; + } + memcpy(xt.data(), g.data(), xt.size() * sizeof(float)); + } else { + for (auto& v : xt) { + v = rng.normal(); + } + if (first) { + memcpy(xt.data(), ff.data(), fel * sizeof(float)); // pin frame 0 + } + } + + // ---- 4-step denoise ---- + for (size_t s = 0; s < cfg.denoise_steps.size(); s++) { + float t_cur = cfg.denoise_steps[s]; + // assemble visible frames: history + block frames + std::vector frames; + std::vector facts; + std::vector fts; + for (size_t h = 0; h < hist.size(); h++) { + frames.push_back(hist[h].data()); + facts.push_back(hist_act[h]); + fts.push_back(cfg.context_noise_t); + } + for (int f = 0; f < Fb; f++) { + frames.push_back(xt.data() + static_cast(f) * fel); + facts.push_back(acts[b]); + fts.push_back(first && f == 0 ? 0.0f : t_cur); + } + + auto flow = runner.forward_step(frames, facts, fts, Fb, n_threads); + if (flow.empty()) { + fprintf(stderr, "forward failed (block %zu step %zu)\n", b, s); + return 1; + } + // unpatchify + x0 + renoise (host) + const int h_len = H / 2, w_len = W / 2; + std::vector x0(static_cast(Fb) * fel); + for (int f = 0; f < Fb; f++) { + float sigma = fts[hist.size() + f] / 1000.0f; + const float* xf = xt.data() + static_cast(f) * fel; + float* of = x0.data() + static_cast(f) * fel; + for (int hh = 0; hh < h_len; hh++) { + for (int ww = 0; ww < w_len; ww++) { + const float* tok = flow.data() + ((static_cast(f) * h_len + hh) * w_len + ww) * (4 * C); + for (int ph = 0; ph < 2; ph++) { + for (int pw = 0; pw < 2; pw++) { + for (int c = 0; c < C; c++) { + size_t idx = (static_cast(c) * H + (hh * 2 + ph)) * W + (ww * 2 + pw); + of[idx] = xf[idx] - sigma * tok[(ph * 2 + pw) * C + c]; + } + } + } + } + } + } + if (s + 1 < cfg.denoise_steps.size()) { + float t_next = cfg.denoise_steps[s + 1], sn = t_next / 1000.0f; + std::vector eps(static_cast(Fb) * fel); + if (!golden.empty()) { + std::vector g; + std::vector gs; + char nm[256]; + snprintf(nm, sizeof(nm), "%s/b%zu_s%zu_eps.npy", golden.c_str(), b, s); + if (!load_npy_f32(nm, g, gs)) { + return 1; + } + memcpy(eps.data(), g.data(), eps.size() * sizeof(float)); + } else { + for (auto& v : eps) { + v = rng.normal(); + } + } + for (size_t i = 0; i < xt.size(); i++) { + xt[i] = (1.0f - sn) * x0[i] + sn * eps[i]; + } + } else { + xt = x0; + } + if (first) { + memcpy(xt.data(), ff.data(), fel * sizeof(float)); // keep frame 0 pinned + } + printf("block %zu step %zu done (t=%.1f)\n", b, s, t_cur); + } + + for (int f = 0; f < Fb; f++) { + hist.emplace_back(xt.begin() + static_cast(f) * fel, + xt.begin() + static_cast(f + 1) * fel); + hist_act.push_back(acts[b]); + } + printf("block %zu finalized (history %zu frames)\n", b, hist.size()); + } + + std::ofstream of(out_path, std::ios::binary); + for (auto& fr : hist) { + of.write(reinterpret_cast(fr.data()), static_cast(fel * sizeof(float))); + } + printf("wrote %s (%zu frames x %zu floats)\n", out_path.c_str(), hist.size(), fel); + return 0; +} + +// ---- minimal .npy reader (f32/f64/i64, C-order) ---- +static bool load_npy_f32(const std::string& path, std::vector& out, std::vector& shape) { + std::ifstream f(path, std::ios::binary); + if (!f) { + fprintf(stderr, "npy: cannot open %s\n", path.c_str()); + return false; + } + char magic[6]; + f.read(magic, 6); + if (memcmp(magic, "\x93NUMPY", 6) != 0) { + return false; + } + uint8_t ver[2]; + f.read(reinterpret_cast(ver), 2); + uint32_t hlen = 0; + if (ver[0] == 1) { + uint16_t h16; + f.read(reinterpret_cast(&h16), 2); + hlen = h16; + } else { + f.read(reinterpret_cast(&hlen), 4); + } + std::string header(hlen, '\0'); + f.read(header.data(), hlen); + bool f4 = header.find("= header.size() || header[sp] == ')') { + break; + } + shape.push_back(strtoll(header.c_str() + sp, nullptr, 10)); + while (sp < header.size() && header[sp] != ',' && header[sp] != ')') { + sp++; + } + } + int64_t n = 1; + for (auto s : shape) { + n *= s; + } + out.resize(static_cast(n)); + if (f4) { + f.read(reinterpret_cast(out.data()), n * 4); + } else { + std::vector tmp(static_cast(n)); + f.read(reinterpret_cast(tmp.data()), n * 8); + for (int64_t i = 0; i < n; i++) { + out[static_cast(i)] = static_cast(tmp[static_cast(i)]); + } + } + return f.good(); +} diff --git a/include/stable-diffusion.h b/include/stable-diffusion.h index f2b6f00fc..deaf20b9a 100644 --- a/include/stable-diffusion.h +++ b/include/stable-diffusion.h @@ -500,6 +500,38 @@ SD_API sd_image_t upscale(upscaler_ctx_t* upscaler_ctx, SD_API int get_upscale_factor(upscaler_ctx_t* upscaler_ctx); +// -- ABot-World interactive walk session -------------------------------------- +// ABot-World is a causal world model: it generates video block-by-block under +// per-block keyboard actions instead of one batch call. A session owns the DiT +// + taehv decoder and a fixed scene pack; each step() generates one latent +// block and returns its decoded RGB frames. Standalone (no sd_ctx_t needed); +// the batch generate_image()/generate_video() APIs keep rejecting ABot models. + +typedef struct sd_abot_session_t sd_abot_session_t; + +typedef struct { + const char* dit_model_path; // ABot-World DiT GGUF (F16 or Q8_0) + const char* taehv_path; // taew2_2 GGUF (streaming pixel decoder) + const char* scene_path; // scene pack safetensors (prompt/first-frame/ref latents) + const char* backend; // backend spec string (NULL/"" = default, e.g. "cpu", "cuda") + int n_threads; // <= 0: physical cores + int64_t seed; // walk noise seed + int num_frame_per_block; // <= 0: model default (3) + int local_attn_size; // <= 0: scene/deploy default; latent-frame window + bool offload_params_to_cpu; +} sd_abot_session_params_t; + +SD_API void sd_abot_session_params_init(sd_abot_session_params_t* params); +SD_API sd_abot_session_t* sd_abot_session_new(const sd_abot_session_params_t* params); +// Generate the next block under `action_mask` (bit 0..7 = W,A,S,D,I,J,K,L held). +// Returns a malloc'd array of `*num_frames_out` RGB frames (free each .data and +// the array with sd_abot_session_frames_free), or NULL on failure. +SD_API sd_image_t* sd_abot_session_step(sd_abot_session_t* session, + uint32_t action_mask, + int* num_frames_out); +SD_API void sd_abot_session_frames_free(sd_image_t* frames, int num_frames); +SD_API void sd_abot_session_free(sd_abot_session_t* session); + SD_API bool convert(const char* input_path, const char* vae_path, const char* output_path, diff --git a/script/validate_abot_world.sh b/script/validate_abot_world.sh new file mode 100644 index 000000000..025fb67b0 --- /dev/null +++ b/script/validate_abot_world.sh @@ -0,0 +1,57 @@ +#!/usr/bin/env bash +# Validation / regression smoke test for ABot-World support. +# +# ABot-World-0-5B-LF is a causal, interactive derivative of Wan2.2-TI2V-5B. This +# script verifies: +# 1. an ABot GGUF is detected as "ABot-World" and all tensors load (incl. the +# act_control_adapter action-conditioning block); +# 2. the batch generate_video() path rejects it with the explanatory error +# (the causal interactive session is a separate, not-yet-landed feature); +# 3. (optional) a stock Wan model still detects + generates unchanged, proving +# no regression from the sd_version_is_wan_ti2v_family() refactor. +# +# Usage: +# ABOT_DIT=/path/abot-dit-q8_0.gguf ABOT_VAE=/path/wan2.2_vae.safetensors \ +# ABOT_T5=/path/umt5-xxl.(pth|gguf|safetensors) \ +# [WAN_DIT=... WAN_VAE=... WAN_T5=...] \ +# SD_CLI=./build/bin/sd-cli script/validate_abot_world.sh +set -euo pipefail + +SD_CLI="${SD_CLI:-./build/bin/sd-cli}" +: "${ABOT_DIT:?set ABOT_DIT}"; : "${ABOT_VAE:?set ABOT_VAE}"; : "${ABOT_T5:?set ABOT_T5}" + +log=$(mktemp) +fail() { echo "FAIL: $1"; echo "--- log tail ---"; tail -30 "$log"; exit 1; } + +echo "== ABot-World detection + load + guarded rejection ==" +set +e +"$SD_CLI" -M vid_gen --diffusion-model "$ABOT_DIT" --vae "$ABOT_VAE" \ + --t5xxl "$ABOT_T5" -p "a coastal street at dusk" -W 480 -H 832 \ + --video-frames 9 -v >"$log" 2>&1 +set -e + +grep -q "ABot-World-5B" "$log" || fail "not detected as ABot-World-5B" +grep -q "loading tensors completed" "$log" || fail "tensor load did not complete" +grep -qi "wrong shape\|not in model file" "$log" && fail "tensor shape/name mismatch on load" +grep -q "not supported by batch" "$log" || fail "missing the guarded-rejection error" +echo "PASS: ABot-World detected, all tensors loaded, batch video path correctly rejected." + +echo "== ABot-World: batch image path must be rejected too ==" +set +e +"$SD_CLI" -M img_gen --diffusion-model "$ABOT_DIT" --vae "$ABOT_VAE" \ + --t5xxl "$ABOT_T5" -p "a coastal street at dusk" -W 480 -H 832 -v >"$log" 2>&1 +set -e +grep -q "not supported by batch" "$log" || fail "generate_image() not rejected for ABot-World" +echo "PASS: batch image path correctly rejected." + +if [[ -n "${WAN_DIT:-}" ]]; then + echo "== Regression: stock Wan still generates ==" + out=$(mktemp -u).png + "$SD_CLI" -M vid_gen --diffusion-model "$WAN_DIT" --vae "${WAN_VAE:?}" \ + --t5xxl "${WAN_T5:?}" -p "a bird flying" -W 480 -H 832 --video-frames 9 \ + -o "$out" -v >"$log" 2>&1 || fail "stock Wan generation failed (regression!)" + grep -qi "not supported by batch" "$log" && fail "stock Wan wrongly rejected (regression!)" + echo "PASS: stock Wan generation unaffected." +fi + +echo "ALL CHECKS PASSED" diff --git a/src/abot_world.hpp b/src/abot_world.hpp new file mode 100644 index 000000000..69458c3ec --- /dev/null +++ b/src/abot_world.hpp @@ -0,0 +1,799 @@ +#ifndef __ABOT_WORLD_HPP__ +#define __ABOT_WORLD_HPP__ + +// ABot-World causal walk — recompute formulation. +// +// ABot-World generates video block-by-block (num_frame_per_block latent frames +// per step) with keyboard-action conditioning. The reference implementation +// keeps a rolling KV cache; this port reproduces the exact same math by +// recomputing attention over the full visible sequence each denoise step: +// +// sequence = [ ref tokens | history frames (clean, t=0) | block frames (t) ] +// +// with a per-row mask: tokens of frame f attend the ref tokens plus all frames +// of blocks <= block(f) that lie inside f's trailing local attention window +// (local_attn_size frames). Attention within a block is full (the causal +// boundary is between blocks). History frames re-derive exactly the K/V the +// reference wrote into its cache (clean latents at context_noise=0 modulation, +// each frame using its own historical action features), including under +// eviction, because each frame's row mask reproduces the window that frame saw +// when it was generated. +// +// RoPE uses window-local temporal ids (the reference's relative-RoPE scheme, +// ids clamped to local_attn_size-1 <= 20) and the reference's negative-time +// ids for ref slots. The 4-step distilled schedule (denoising_step_list warped +// by timestep_shift) runs flow-matching updates: +// x0 = xt - sigma(t) * flow, xt' = (1 - sigma(t')) * x0 + sigma(t') * eps +// where sigma(t) = t / 1000 for the warped step values. +// +// Semantics derived from the reference (github.com/amap-cvlab/ABot-World): +// pipeline/causal_inference.py generate_next_block, wan/modules/causal_model.py +// (_forward_inference, ref token prep, relative RoPE), utils/wan_wrapper.py +// (flow<->x0 conversion), utils/scheduler.py (FlowMatchScheduler). + +#include +#include +#include +#include +#include +#include + +#include +#include + +#include "ggml_extend.hpp" +#include "model.h" +#include "rope.hpp" +#include "vae.hpp" // must precede tae.hpp (TinyVideoAutoEncoder's base) +#include "tae.hpp" +#include "wan.hpp" + +namespace ABOT { + +struct AbotWorldConfig { + int num_frame_per_block = 3; + int local_attn_size = 8; // latent-frame window (test config; deployed default 21) + int rope_temporal_clamp = 20; // reference clamps window-local temporal ids to <= 20 + float context_noise_t = 0.0f; // history frames' timestep + std::vector denoise_steps = {1000.0f, 937.5f, 833.3333333f, 625.0f}; + int act_in_dim = 32; // 8 keys x 4 (repeat_interleave) + int act_downscale_factor = 16; +}; + +static inline float abot_sigma_of_t(float t) { + return t / 1000.0f; // warped steps satisfy t = sigma' * 1000 exactly +} + +// One latent frame of walk state. +struct AbotFrame { + sd::Tensor latent; // [W_l, H_l, C] (ggml order) — clean/final latent + uint8_t action_mask = 0; // 8-key bitmask active when this frame was generated +}; + +struct AbotScenePack { + sd::Tensor prompt_embeds; // [4096, 512] per token row-major -> tensor shape {4096, 512, 1} + sd::Tensor first_frame_latents; // {W_l, H_l, C, 1} + sd::Tensor ref_latents; // {32, 32, C, K} (T=1 squeezed) + sd::Tensor ref_mask; // {K} + int ref_slots = 0; + + // Minimal safetensors reader (F32 tensors only — the scene pack is fp32). + // Torch shapes map to ggml ne reversed; trailing singleton dims dropped. + bool load(const std::string& path) { + std::ifstream f(path, std::ios::binary); + if (!f) { + LOG_ERROR("scene pack: cannot open '%s'", path.c_str()); + return false; + } + uint64_t hlen = 0; + f.read(reinterpret_cast(&hlen), 8); + std::string header(hlen, '\0'); + f.read(header.data(), static_cast(hlen)); + const uint64_t data_base = 8 + hlen; + + auto fetch = [&](const std::string& name, sd::Tensor& out, + std::vector& torch_shape) -> bool { + size_t p = header.find("\"" + name + "\""); + if (p == std::string::npos) { + return false; + } + size_t sp = header.find("\"shape\"", p); + size_t sb = header.find('[', sp), se = header.find(']', sb); + torch_shape.clear(); + for (size_t q = sb + 1; q < se;) { + while (q < se && (header[q] == ' ' || header[q] == ',')) q++; + if (q >= se) break; + torch_shape.push_back(strtoll(header.c_str() + q, nullptr, 10)); + while (q < se && header[q] != ',') q++; + } + size_t op = header.find("\"data_offsets\"", p); + size_t ob = header.find('[', op); + uint64_t off0 = strtoull(header.c_str() + ob + 1, nullptr, 10); + int64_t n = 1; + for (auto s : torch_shape) n *= s; + std::vector ne; // reversed, no trailing 1s beyond 4 + for (auto it = torch_shape.rbegin(); it != torch_shape.rend(); ++it) { + if (*it != 1 || ne.size() < 1) ne.push_back(*it); + } + while (ne.size() < 1) ne.push_back(1); + out.resize(ne); + if (out.numel() != n) { + out.resize({n}); + } + f.seekg(static_cast(data_base + off0)); + f.read(reinterpret_cast(out.data()), n * static_cast(sizeof(float))); + return f.good(); + }; + + std::vector sh; + if (!fetch("prompt_embeds", prompt_embeds, sh)) { // [1,512,4096] + return false; + } + prompt_embeds.resize({sh[2], sh[1], 1, 1}); + if (!fetch("first_frame_latents", first_frame_latents, sh)) { // [1,1,C,H,W] + return false; + } + first_frame_latents.resize({sh[4], sh[3], sh[2], 1}); + if (fetch("ref_latents", ref_latents, sh)) { // [1,K,C,1,32,32] + ref_slots = static_cast(sh[1]); + // stored per-slot [C,1,32,32]; conv needs ne {32,32,K,C}: + // torch order is (K,C,32,32) contiguous -> need transpose to {W,H,K,C} + sd::Tensor re({sh[5], sh[4], sh[1], sh[2]}); + const int64_t K = sh[1], C = sh[2], HW = sh[4] * sh[5]; + for (int64_t k = 0; k < K; k++) { + for (int64_t c = 0; c < C; c++) { + memcpy(re.data() + (c * K + k) * HW, + ref_latents.data() + (k * C + c) * HW, + static_cast(HW) * sizeof(float)); + } + } + ref_latents = std::move(re); + std::vector msh; + if (!fetch("ref_mask", ref_mask, msh)) { + return false; + } + ref_mask.resize({ref_slots}); + } + return true; + } +}; + +// Wan subclass: same weights/blocks, causal recompute graph. +class AbotWan : public WAN::Wan { +public: + AbotWan() = default; + explicit AbotWan(WAN::WanParams params) + : WAN::Wan(params) {} + + // Layout conventions mirror WAN::Wan::forward_orig: + // x_all: ne {W_l, H_l, F_vis, C} latents (history clean + block xt) + // act_planes: ne {W_l/ds, H_l/ds, act_in_dim*ds*ds, F_vis}? NO — adapter input + // is the pixel-unshuffled action image at latent grid: + // ne {W_l, H_l, act_in_dim*ds*ds/ (16/ latent factor) ...} — the + // host builds it at ne {w_in, h_in, C_unsh, F_vis} such that the + // adapter conv (k2 s2) lands on the token grid (w_len, h_len). + // t_frames: ne {F_vis + 1} per-frame timesteps, extra last row = 0 (ref) + // context: ne {4096, 512, 1} prompt embeds + // ref_latents: ne {32, 32, K, C} or nullptr (T dim = slots, temporal k=1) + // ref_mask: ne {K} or nullptr + // pe: packed rope table for [ref | video] token sequence + // attn_mask: ne {L, L} additive f32 (built host-side) + // token_frame_rows: ne {L} int32 -> row in the timestep table (ref rows -> F_vis) + // returns flow tokens for the LAST block_frames frames: {pt*ph*pw*out_dim, block_tokens, 1} + ggml_tensor* forward_causal(GGMLRunnerContext* ctx, + ggml_tensor* x_all, + ggml_tensor* act_planes, + ggml_tensor* t_frames, + ggml_tensor* context, + ggml_tensor* ref_latents, + ggml_tensor* ref_mask, + ggml_tensor* pe, + ggml_tensor* attn_mask, + ggml_tensor* token_frame_rows, + int block_frames, + int64_t& out_h_len, + int64_t& out_w_len) { + auto gctx = ctx->ggml_ctx; + + auto patch_embedding = std::dynamic_pointer_cast(blocks["patch_embedding"]); + auto text_embedding_0 = std::dynamic_pointer_cast(blocks["text_embedding.0"]); + auto text_embedding_2 = std::dynamic_pointer_cast(blocks["text_embedding.2"]); + auto time_embedding_0 = std::dynamic_pointer_cast(blocks["time_embedding.0"]); + auto time_embedding_2 = std::dynamic_pointer_cast(blocks["time_embedding.2"]); + auto time_projection_1 = std::dynamic_pointer_cast(blocks["time_projection.1"]); + auto head = std::dynamic_pointer_cast(blocks["head"]); + auto adapter = std::dynamic_pointer_cast(blocks["act_control_adapter"]); + + const int64_t F_vis = x_all->ne[2]; + + // ── video tokens: patchify (+ action adapter), forward_orig layout chain ── + ggml_tensor* x = patch_embedding->forward(ctx, x_all); // {w_len, h_len, F_vis, dim} + out_w_len = x->ne[0]; + out_h_len = x->ne[1]; + const int64_t fsl = out_w_len * out_h_len; // tokens per frame + x = ggml_reshape_3d(gctx, x, x->ne[0] * x->ne[1] * x->ne[2], x->ne[3], 1); // {w*h*F, dim, 1} + x = ggml_ext_cont(gctx, ggml_ext_torch_permute(gctx, x, 1, 0, 2, 3)); // {dim, w*h*F, 1} + + if (adapter != nullptr && act_planes != nullptr) { + ggml_tensor* a = adapter->forward(ctx, act_planes); // {w_len, h_len, dim, F_vis} + a = ggml_reshape_3d(gctx, a, a->ne[0] * a->ne[1], a->ne[2], a->ne[3]); // {w*h, dim, F} + a = ggml_ext_cont(gctx, ggml_ext_torch_permute(gctx, a, 1, 0, 2, 3)); // {dim, w*h, F} + a = ggml_reshape_3d(gctx, a, a->ne[0], a->ne[1] * a->ne[2], 1); // {dim, w*h*F, 1} + x = ggml_add(gctx, x, a); // act_context_scale = 1.0 + } + + // ── ref tokens (masked after conv so conv bias is cancelled on empty slots) ── + if (ref_latents != nullptr) { + ggml_tensor* r = patch_embedding->forward(ctx, ref_latents); // {16, 16, K, dim} + ggml_tensor* m = ggml_reshape_4d(gctx, ref_mask, 1, 1, ref_mask->ne[0], 1); + r = ggml_mul(gctx, r, m); + r = ggml_reshape_3d(gctx, r, r->ne[0] * r->ne[1] * r->ne[2], r->ne[3], 1); // {16*16*K, dim, 1} + r = ggml_ext_cont(gctx, ggml_ext_torch_permute(gctx, r, 1, 0, 2, 3)); // {dim, n_ref, 1} + x = ggml_concat(gctx, r, x, 1); // ref first + } + const int64_t L = x->ne[1]; + + // ── per-frame modulation table -> per-token via get_rows ── + // rows: F_vis frame timesteps + 1 extra row (t = 0) used by ref tokens + ggml_tensor* e_tbl = ggml_ext_timestep_embedding(gctx, t_frames, params.freq_dim); // [freq_dim, F_vis+1] + e_tbl = time_embedding_0->forward(ctx, e_tbl); + e_tbl = ggml_silu_inplace(gctx, e_tbl); + e_tbl = time_embedding_2->forward(ctx, e_tbl); // [dim, F_vis+1] (this is `e`) + ggml_tensor* e0_tbl = ggml_silu(gctx, e_tbl); + e0_tbl = time_projection_1->forward(ctx, e0_tbl); // [6*dim, F_vis+1] + + ggml_tensor* e0_tok = ggml_get_rows(gctx, e0_tbl, token_frame_rows); // [6*dim, L] + e0_tok = ggml_reshape_4d(gctx, e0_tok, params.dim, 6, L, 1); // [N=1, L, 6, dim] + + // ── context ── + ggml_tensor* c = text_embedding_0->forward(ctx, context); + c = ggml_ext_gelu(gctx, c); + c = text_embedding_2->forward(ctx, c); // [dim, 512, 1] + + // ── transformer ── + for (int i = 0; i < params.num_layers; i++) { + auto block = std::dynamic_pointer_cast(blocks["blocks." + std::to_string(i)]); + x = block->forward(ctx, x, e0_tok, pe, c, 0, attn_mask); + } + + // ── head over the trailing block tokens only ── + const int64_t block_tokens = static_cast(block_frames) * fsl; + ggml_tensor* xb = ggml_ext_slice(gctx, x, 1, L - block_tokens, L); // [dim, block_tokens, 1] + // per-frame e for the block frames: rows F_vis-block_frames .. F_vis-1 + ggml_tensor* e_blk = ggml_ext_slice(gctx, e_tbl, 1, F_vis - block_frames, F_vis); // [dim, block_frames] + e_blk = ggml_reshape_3d(gctx, e_blk, params.dim, block_frames, 1); // [N=1, T, dim] + ggml_tensor* out = head->forward(ctx, xb, e_blk); // [pt*ph*pw*out_dim, block_tokens, 1] + return out; + } +}; + +// Runner: owns the AbotWan blocks + walk state, builds one graph per denoise +// forward, and runs the host-side 4-step block loop. +struct AbotWorldRunner : public GGMLRunner { + AbotWorldConfig cfg; + WAN::WanParams wan_params; + AbotWan wan; + AbotScenePack scene; + + // walk state + std::vector history; // finalized latent frames (clean) + int64_t lat_w = 0, lat_h = 0, lat_c = 0; + + std::vector pe_vec; + + AbotWorldRunner(ggml_backend_t backend, + ggml_backend_t params_backend, + const String2TensorStorage& tensor_storage_map, + const std::string prefix, + const AbotWorldConfig& config = {}) + : GGMLRunner(backend, params_backend), cfg(config) { + wan_params.model_type = "t2v"; + wan_params.dim = 3072; + wan_params.eps = 1e-06f; + wan_params.ffn_dim = 14336; + wan_params.freq_dim = 256; + wan_params.in_dim = 48; + wan_params.num_heads = 24; + wan_params.out_dim = 48; + wan_params.text_len = 512; + wan_params.num_layers = 30; + wan_params.abot_world = true; + wan = AbotWan(wan_params); + wan.init(params_ctx, tensor_storage_map, prefix); + } + + std::string get_desc() override { + return "ABot-World-5B (causal walk)"; + } + + void get_param_tensors(std::map& tensors, const std::string& prefix) { + wan.get_param_tensors(tensors, prefix); + } + + // ---- host-side helpers ---- + + // Per-token 3-axis rope ids: ref slots at negative constant times (the + // reference's _build_ref_freqs), video frames at absolute frame ids (the + // reference's REL_ROPE_CACHE fast path with base 0; valid for walks well + // below the re-base threshold of ~256 frames). + std::vector build_pe(int n_ref_slots, int ref_grid, int F_vis, int h_len, int w_len) { + std::vector> ids; + const int tokens_per_slot = ref_grid * ref_grid; + const int t_step = std::max(tokens_per_slot, 256); + for (int k = 0; k < n_ref_slots; k++) { + float t_id = -static_cast((n_ref_slots - k) * t_step); + for (int hh = 0; hh < ref_grid; hh++) { + for (int ww = 0; ww < ref_grid; ww++) { + ids.push_back({t_id, static_cast(hh), static_cast(ww)}); + } + } + } + for (int f = 0; f < F_vis; f++) { + for (int hh = 0; hh < h_len; hh++) { + for (int ww = 0; ww < w_len; ww++) { + ids.push_back({static_cast(f), static_cast(hh), static_cast(ww)}); + } + } + } + return Rope::embed_nd(ids, 1, static_cast(wan_params.theta), wan_params.axes_dim); + } + + // Additive mask {L(k), L(q)}: video rows of block b attend ref + frames of + // blocks <= b inside b's trailing window; ref rows attend ref + block 0. + std::vector build_mask(int n_ref, int F_vis, int fsl, int block_frames) { + const int64_t L = n_ref + static_cast(F_vis) * fsl; + std::vector mask(static_cast(L) * L, -INFINITY); + auto frame_of = [&](int64_t tok) -> int { // -1 for ref tokens + return tok < n_ref ? -1 : static_cast((tok - n_ref) / fsl); + }; + for (int64_t q = 0; q < L; q++) { + int fq = frame_of(q); + int bq = fq < 0 ? 0 : fq / cfg.num_frame_per_block; + // trailing window (in frames) as seen when block bq was generated + int hi = (bq + 1) * cfg.num_frame_per_block; // frames [0, hi) + int lo = std::max(0, hi - cfg.local_attn_size); + for (int64_t k = 0; k < L; k++) { + int fk = frame_of(k); + bool allowed = false; + if (fk < 0) { + allowed = true; // ref keys visible to everyone + } else if (fq < 0) { + allowed = fk < cfg.num_frame_per_block; // ref rows formed with block 0 + } else { + allowed = fk >= lo && fk < hi; + } + if (allowed) { + mask[static_cast(q) * L + k] = 0.0f; + } + } + } + return mask; + } + + // Pixel-unshuffled action planes for one frame: the reference broadcasts the + // 8-key vector to 8 channels, repeat_interleaves x4 -> 32 channels constant + // over HxW, then PixelUnshuffle(16): output channel c corresponds to input + // channel c / (16*16) -> value = key[(c / 256) / 4]. + void fill_act_plane(float* dst, uint8_t action_mask, int w_in, int h_in, int c_unsh) { + for (int c = 0; c < c_unsh; c++) { + int key = (c / (cfg.act_downscale_factor * cfg.act_downscale_factor)) / 4; + float val = (action_mask >> key) & 1 ? 1.0f : 0.0f; + float* p = dst + static_cast(c) * w_in * h_in; + std::fill(p, p + static_cast(w_in) * h_in, val); + } + } + + // One denoise forward over the full visible sequence; returns x0 for the + // current block (host applies the flow->x0 conversion). + // xt_block: {W,H,C} x block_frames latents (current noisy block) + // t_value: timestep for the block frames (frame 0 of walk pinned to 0 by caller) + sd::Tensor forward_step(const std::vector& frame_latents, // F_vis pointers {W*H*C} + const std::vector& frame_actions, + const std::vector& frame_timesteps, // F_vis values + int block_frames, + int n_threads) { + const int F_vis = static_cast(frame_latents.size()); + const int ds = cfg.act_downscale_factor; + const int w_in = static_cast(lat_w) / ds * ds; // == lat_w (multiple of 16 grid) + // adapter input grid: unshuffle by ds over the PIXEL grid == latent grid here: + // reference act image is [32, H_px, W_px] with H_px = lat_h*16? No: the + // adapter output must land on (h_len, w_len) = (lat_h/2, lat_w/2); conv is + // k2 s2, so its input grid is (lat_h, lat_w) with C = 32*ds*ds = 8192. + const int c_unsh = cfg.act_in_dim * ds * ds; + + // ---- build host buffers ---- + int h_len = static_cast(lat_h) / 2; + int w_len = static_cast(lat_w) / 2; + int fsl = h_len * w_len; + int n_ref = scene.ref_slots * 16 * 16; + + sd::Tensor x_all({lat_w, lat_h, F_vis, lat_c}); + for (int f = 0; f < F_vis; f++) { + // {W,H,T,C}: strides — for each channel c, plane at [.., f, c] + for (int64_t c = 0; c < lat_c; c++) { + float* dst = x_all.data() + (c * F_vis + f) * lat_w * lat_h; + const float* src = frame_latents[f] + c * lat_w * lat_h; + memcpy(dst, src, static_cast(lat_w) * lat_h * sizeof(float)); + } + } + + sd::Tensor act({lat_w, lat_h, c_unsh, F_vis}); + for (int f = 0; f < F_vis; f++) { + fill_act_plane(act.data() + static_cast(f) * c_unsh * lat_w * lat_h, + frame_actions[f], static_cast(lat_w), static_cast(lat_h), c_unsh); + } + + sd::Tensor tvec({F_vis + 1}); + for (int f = 0; f < F_vis; f++) { + tvec.data()[f] = frame_timesteps[f]; + } + tvec.data()[F_vis] = 0.0f; // ref/modulation row + + std::vector rows(static_cast(n_ref) + static_cast(F_vis) * fsl); + for (int i = 0; i < n_ref; i++) { + rows[i] = F_vis; + } + for (int f = 0; f < F_vis; f++) { + std::fill(rows.begin() + n_ref + static_cast(f) * fsl, + rows.begin() + n_ref + static_cast(f + 1) * fsl, f); + } + + pe_vec = build_pe(scene.ref_slots, 16, F_vis, h_len, w_len); + std::vector mask = build_mask(n_ref, F_vis, fsl, block_frames); + const int64_t L = n_ref + static_cast(F_vis) * fsl; + + // ---- graph ---- + int64_t out_h = 0, out_w = 0; + auto get_graph = [&]() -> ggml_cgraph* { + ggml_cgraph* gf = ggml_new_graph_custom(compute_ctx, WAN::WAN_GRAPH_SIZE, false); + ggml_tensor* x_in = to_backend_input(x_all); + ggml_tensor* act_in = to_backend_input(act); + ggml_tensor* t_in = to_backend_input(tvec); + ggml_tensor* ctx_in = to_backend_input(scene.prompt_embeds); + ggml_tensor* ref_in = nullptr; + ggml_tensor* rm_in = nullptr; + if (scene.ref_slots > 0) { + ref_in = to_backend_input(scene.ref_latents); + rm_in = to_backend_input(scene.ref_mask); + } + ggml_tensor* pe = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, 2, 2, + wan_params.axes_dim_sum / 2, + static_cast(pe_vec.size()) / wan_params.axes_dim_sum / 2); + set_backend_tensor_data(pe, pe_vec.data()); + ggml_tensor* mk = ggml_new_tensor_2d(compute_ctx, GGML_TYPE_F32, L, L); + set_backend_tensor_data(mk, mask.data()); + ggml_tensor* rw = ggml_new_tensor_1d(compute_ctx, GGML_TYPE_I32, static_cast(rows.size())); + set_backend_tensor_data(rw, rows.data()); + + auto runner_ctx = get_context(); + ggml_tensor* out = wan.forward_causal(&runner_ctx, x_in, act_in, t_in, ctx_in, + ref_in, rm_in, pe, mk, rw, block_frames, out_h, out_w); + ggml_build_forward_expand(gf, out); + return gf; + }; + + auto result = GGMLRunner::compute(get_graph, n_threads, false); + if (!result.has_value()) { + return {}; + } + return std::move(*result); + } + +private: + template + ggml_tensor* to_backend_input(sd::Tensor& t) { + auto shape = t.shape(); + ggml_tensor* g = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, + shape.size() > 0 ? shape[0] : 1, + shape.size() > 1 ? shape[1] : 1, + shape.size() > 2 ? shape[2] : 1, + shape.size() > 3 ? shape[3] : 1); + set_backend_tensor_data(g, t.data()); + return g; + } +}; + +// Seed-stable normal RNG (xorshift128+ / Box-Muller) so free walks reproduce +// across platforms regardless of libstdc++/libc++ std::normal_distribution. +struct AbotRng { + uint64_t s0, s1; + bool have_spare = false; + float spare = 0.0f; + explicit AbotRng(uint64_t seed = 42) { + s0 = seed ^ 0x9E3779B97F4A7C15ULL; + s1 = (seed << 1) | 1; + for (int i = 0; i < 16; i++) { + next_u64(); + } + } + uint64_t next_u64() { + uint64_t x = s0, y = s1; + s0 = y; + x ^= x << 23; + s1 = x ^ y ^ (x >> 17) ^ (y >> 26); + return s1 + y; + } + float normal() { + if (have_spare) { + have_spare = false; + return spare; + } + auto uniform = [&]() -> float { + return static_cast((next_u64() >> 11) * (1.0 / 9007199254740992.0)); + }; + float u1 = std::max(uniform(), 1e-12f), u2 = uniform(); + float r = std::sqrt(-2.0f * std::log(u1)), a = 6.28318530718f * u2; + spare = r * std::sin(a); + have_spare = true; + return r * std::cos(a); + } + void fill(float* dst, size_t n) { + for (size_t i = 0; i < n; i++) { + dst[i] = normal(); + } + } +}; + +// Interactive walk session: DiT walk core + taehv pixel decode, one keyboard +// action per generated block. Self-contained (loads its own weights); this is +// the object behind the public sd_abot_session_* C API. +class AbotWalkSession { +public: + AbotWorldConfig cfg; + AbotRng rng; + int n_threads = 8; + + // optional noise override for deterministic validation: called with the + // denoise step index (-1 = initial block noise, 0..2 = renoise eps); + // return false to fall back to the RNG. + std::function noise_override; + + std::unique_ptr runner; + std::shared_ptr tae; + + // finalized walk state (per-frame ggml {W,H,C} latents, torch [C,H,W] flat) + std::vector> history; + std::vector history_actions; + + bool load(ggml_backend_t runtime_backend, + ggml_backend_t params_backend, + ggml_backend_t vae_backend, + ggml_backend_t vae_params_backend, + const std::string& dit_path, + const std::string& taehv_path, + const std::string& scene_path, + const AbotWorldConfig& config, + uint64_t seed, + int threads) { + cfg = config; + rng = AbotRng(seed); + n_threads = threads; + + ModelLoader ml; + if (!ml.init_from_file(dit_path, "model.diffusion_model.")) { + LOG_ERROR("abot session: cannot open DiT '%s'", dit_path.c_str()); + return false; + } + runner = std::make_unique(runtime_backend, params_backend, + ml.get_tensor_storage_map(), + "model.diffusion_model.", cfg); + if (!runner->alloc_params_buffer()) { + return false; + } + std::map tensors; + runner->get_param_tensors(tensors, "model.diffusion_model"); + if (!ml.load_tensors(tensors, {}, n_threads)) { + LOG_ERROR("abot session: DiT tensor load failed"); + return false; + } + if (!runner->scene.load(scene_path)) { + return false; + } + auto ffl = runner->scene.first_frame_latents.shape(); + runner->lat_w = ffl[0]; + runner->lat_h = ffl[1]; + runner->lat_c = ffl.size() > 2 ? ffl[2] : 1; + + ModelLoader tl; + if (!tl.init_from_file(taehv_path, "tae.")) { // same prefixing as new_sd_ctx's taesd path + LOG_ERROR("abot session: cannot open taehv '%s'", taehv_path.c_str()); + return false; + } + tae = std::make_shared(vae_backend, vae_params_backend, + tl.get_tensor_storage_map(), + "decoder", true, VERSION_ABOT_WORLD); + if (!tae->alloc_params_buffer()) { + return false; + } + std::map tae_tensors; + tae->get_param_tensors(tae_tensors, "tae"); + if (!tl.load_tensors(tae_tensors, {}, n_threads)) { + LOG_ERROR("abot session: taehv tensor load failed"); + return false; + } + LOG_INFO("abot session ready: latent %dx%dx%d, ref slots %d, window %d", + (int)runner->lat_w, (int)runner->lat_h, (int)runner->lat_c, + runner->scene.ref_slots, cfg.local_attn_size); + return true; + } + + int64_t frame_elems() const { + return runner->lat_w * runner->lat_h * runner->lat_c; + } + + // Generate the next latent block under `action_mask` (bits: W,A,S,D,I,J,K,L). + // Appends num_frame_per_block clean frames to history; returns them. + bool step_latents(uint8_t action_mask, std::vector>& out_frames) { + const int Fb = cfg.num_frame_per_block; + const size_t fel = static_cast(frame_elems()); + const bool first = history.empty(); + const int block = static_cast(history.size()) / Fb; + const float* ff = runner->scene.first_frame_latents.data(); + + std::vector xt(static_cast(Fb) * fel); + if (!noise_override || !noise_override(block, -1, xt.data(), xt.size())) { + rng.fill(xt.data(), xt.size()); + if (first) { + memcpy(xt.data(), ff, fel * sizeof(float)); // pin frame 0 to the scene + } + } + + for (size_t s = 0; s < cfg.denoise_steps.size(); s++) { + const float t_cur = cfg.denoise_steps[s]; + std::vector frames; + std::vector facts; + std::vector fts; + for (size_t h = 0; h < history.size(); h++) { + frames.push_back(history[h].data()); + facts.push_back(history_actions[h]); + fts.push_back(cfg.context_noise_t); + } + for (int f = 0; f < Fb; f++) { + frames.push_back(xt.data() + static_cast(f) * fel); + facts.push_back(action_mask); + fts.push_back(first && f == 0 ? 0.0f : t_cur); + } + + auto flow = runner->forward_step(frames, facts, fts, Fb, n_threads); + if (flow.empty()) { + LOG_ERROR("abot session: forward failed (block %d step %zu)", block, s); + return false; + } + + const int H = static_cast(runner->lat_h), W = static_cast(runner->lat_w); + const int C = static_cast(runner->lat_c); + const int h_len = H / 2, w_len = W / 2; + std::vector x0(static_cast(Fb) * fel); + for (int f = 0; f < Fb; f++) { + float sigma = abot_sigma_of_t(fts[history.size() + f]); + const float* xf = xt.data() + static_cast(f) * fel; + float* of = x0.data() + static_cast(f) * fel; + for (int hh = 0; hh < h_len; hh++) { + for (int ww = 0; ww < w_len; ww++) { + const float* tok = flow.data() + + ((static_cast(f) * h_len + hh) * w_len + ww) * (4 * C); + for (int ph = 0; ph < 2; ph++) { + for (int pw = 0; pw < 2; pw++) { + for (int c = 0; c < C; c++) { + size_t idx = (static_cast(c) * H + (hh * 2 + ph)) * W + (ww * 2 + pw); + of[idx] = xf[idx] - sigma * tok[(ph * 2 + pw) * C + c]; + } + } + } + } + } + } + if (s + 1 < cfg.denoise_steps.size()) { + const float sn = abot_sigma_of_t(cfg.denoise_steps[s + 1]); + std::vector eps(static_cast(Fb) * fel); + if (!noise_override || !noise_override(block, static_cast(s), eps.data(), eps.size())) { + rng.fill(eps.data(), eps.size()); + } + for (size_t i = 0; i < xt.size(); i++) { + xt[i] = (1.0f - sn) * x0[i] + sn * eps[i]; + } + } else { + xt = std::move(x0); + } + if (first) { + memcpy(xt.data(), ff, fel * sizeof(float)); // keep frame 0 pinned + } + } + + out_frames.clear(); + for (int f = 0; f < Fb; f++) { + out_frames.emplace_back(xt.begin() + static_cast(f) * fel, + xt.begin() + static_cast(f + 1) * fel); + history.push_back(out_frames.back()); + history_actions.push_back(action_mask); + } + return true; + } + + // Decode the newest block to pixels with a decode_overlap-latent-frame + // prefix from the previous block (taehv here is stateless per call; the + // overlap replaces the reference's streaming decoder cache and its warmup + // frames are dropped). A 3-frame overlap matches the reference streaming + // decode at >33 dB PSNR on every frame (mean 46 dB) in local validation. + // Returns {W_px, H_px, T_px, 3} floats in [0, 1]. + sd::Tensor decode_last_block(int decode_overlap = 3) { + const int Fb = cfg.num_frame_per_block; + const size_t fel = static_cast(frame_elems()); + const int total = static_cast(history.size()); + if (total < Fb) { + return {}; + } + const int lead = std::min(decode_overlap, total - Fb); + const int T = Fb + lead; + const int64_t W = runner->lat_w, H = runner->lat_h, C = runner->lat_c; + + // wan layout {W, H, T, C}: per-channel frame planes + sd::Tensor z({W, H, T, C}); + for (int f = 0; f < T; f++) { + const float* src = history[static_cast(total - T + f)].data(); + for (int64_t c = 0; c < C; c++) { + memcpy(z.data() + (c * T + f) * W * H, src + c * W * H, + static_cast(W) * H * sizeof(float)); + } + } + sd_tiling_params_t no_tiling = {}; + sd::Tensor px = tae->decode(n_threads, z, no_tiling, true); + if (px.empty()) { + LOG_ERROR("abot session: taehv decode failed"); + return {}; + } + // px {W_px, H_px, T_px, 3}: the taehv decoder emits [0, 1] pixels + // directly (taesd convention). Drop warmup/overlap frames: stateless + // decode of T latents yields 4T-3 px frames; the last 4*Fb belong to + // the current block (or 4*Fb-3 when there is no lead block). + const int64_t Wp = px.shape()[0], Hp = px.shape()[1], Tp = px.shape()[2]; + const int64_t keep = lead > 0 ? 4 * Fb : Tp; + const int64_t start = Tp - keep; + sd::Tensor out({Wp, Hp, keep, 3}); + for (int64_t ch = 0; ch < 3; ch++) { + for (int64_t t = 0; t < keep; t++) { + const float* sp = px.data() + (ch * Tp + start + t) * Wp * Hp; + float* dp = out.data() + (ch * keep + t) * Wp * Hp; + for (int64_t i = 0; i < Wp * Hp; i++) { + dp[i] = std::min(std::max(sp[i], 0.0f), 1.0f); + } + } + } + return out; + } + + // Convenience: step + decode to 8-bit RGB frames (row-major, interleaved). + bool step(uint8_t action_mask, + std::vector>& rgb_frames, + int64_t& px_w, + int64_t& px_h) { + std::vector> latents; + if (!step_latents(action_mask, latents)) { + return false; + } + sd::Tensor px = decode_last_block(); + if (px.empty()) { + return false; + } + px_w = px.shape()[0]; + px_h = px.shape()[1]; + const int64_t T = px.shape()[2]; + rgb_frames.assign(static_cast(T), {}); + for (int64_t t = 0; t < T; t++) { + auto& fr = rgb_frames[static_cast(t)]; + fr.resize(static_cast(px_w) * px_h * 3); + for (int64_t ch = 0; ch < 3; ch++) { + const float* sp = px.data() + (ch * T + t) * px_w * px_h; + for (int64_t i = 0; i < px_w * px_h; i++) { + fr[static_cast(i) * 3 + ch] = + static_cast(sp[i] * 255.0f + 0.5f); + } + } + } + return true; + } +}; + +} // namespace ABOT + +#endif // __ABOT_WORLD_HPP__ diff --git a/src/model.cpp b/src/model.cpp index cac301425..07ea670da 100644 --- a/src/model.cpp +++ b/src/model.cpp @@ -420,6 +420,7 @@ SDVersion ModelLoader::get_sd_version() { bool is_flux2 = false; bool has_single_block_47 = false; bool is_wan = false; + bool has_act_control_adapter = false; int64_t patch_embedding_channels = 0; bool has_img_emb = false; bool has_middle_block_1 = false; @@ -479,6 +480,10 @@ SDVersion ModelLoader::get_sd_version() { if (tensor_storage.name.find("model.diffusion_model.blocks.0.cross_attn.norm_k.weight") != std::string::npos) { is_wan = true; } + if (tensor_storage.name.find("model.diffusion_model.act_control_adapter.conv.weight") != std::string::npos) { + // ABot-World action-conditioning adapter (keyboard control input) + has_act_control_adapter = true; + } if (tensor_storage.name.find("model.diffusion_model.patch_embedding.weight") != std::string::npos) { patch_embedding_channels = tensor_storage.ne[3]; } @@ -539,6 +544,10 @@ SDVersion ModelLoader::get_sd_version() { } if (is_wan) { LOG_DEBUG("patch_embedding_channels %d", patch_embedding_channels); + if (has_act_control_adapter) { + // Causal Wan2.2-TI2V-5B derivative with action conditioning + return VERSION_ABOT_WORLD; + } if (patch_embedding_channels == 184320 && !has_img_emb) { return VERSION_WAN2_2_I2V; } diff --git a/src/model.h b/src/model.h index 322a86146..0be71eb37 100644 --- a/src/model.h +++ b/src/model.h @@ -51,6 +51,7 @@ enum SDVersion { VERSION_LONGCAT, VERSION_PID, VERSION_IDEOGRAM4, + VERSION_ABOT_WORLD, VERSION_COUNT, }; @@ -117,13 +118,25 @@ static inline bool sd_version_is_ltxav(SDVersion version) { return false; } +static inline bool sd_version_is_abot_world(SDVersion version) { + // ABot-World-0-5B-LF: causal/interactive derivative of Wan2.2-TI2V-5B + // (extra act_control_adapter.* tensors; same latent space and VAE family). + return version == VERSION_ABOT_WORLD; +} + static inline bool sd_version_is_wan(SDVersion version) { - if (version == VERSION_WAN2 || version == VERSION_WAN2_2_I2V || version == VERSION_WAN2_2_TI2V) { + if (version == VERSION_WAN2 || version == VERSION_WAN2_2_I2V || version == VERSION_WAN2_2_TI2V || + sd_version_is_abot_world(version)) { return true; } return false; } +// Wan models that use the Wan2.2 TI2V latent space (48ch, 16x) and VAE. +static inline bool sd_version_is_wan_ti2v_family(SDVersion version) { + return version == VERSION_WAN2_2_TI2V || sd_version_is_abot_world(version); +} + static inline bool sd_version_is_qwen_image(SDVersion version) { if (version == VERSION_QWEN_IMAGE) { return true; diff --git a/src/stable-diffusion.cpp b/src/stable-diffusion.cpp index df84281b2..eff763e20 100644 --- a/src/stable-diffusion.cpp +++ b/src/stable-diffusion.cpp @@ -30,6 +30,7 @@ #include "ltx_latent_upscaler.hpp" #include "ltx_vae.hpp" #include "ltxv.hpp" +#include "abot_world.hpp" #include "mmdit.hpp" #include "pid.hpp" #include "pmid.hpp" @@ -86,6 +87,7 @@ const char* model_version_to_str[] = { "Longcat-Image", "PiD", "Ideogram 4", + "ABot-World", }; const char* sampling_methods_str[] = { @@ -2125,7 +2127,7 @@ class StableDiffusionGGML { sd::Tensor timesteps_tensor({static_cast(timesteps_vec.size())}, timesteps_vec); sd::Tensor guidance_tensor({1}, std::vector{guidance.distilled_guidance}); sd::Tensor noised_input = x * c_in; - if (!denoise_mask.empty() && (version == VERSION_WAN2_2_TI2V || sd_version_is_ltxav(version))) { + if (!denoise_mask.empty() && (sd_version_is_wan_ti2v_family(version) || sd_version_is_ltxav(version))) { noised_input = noised_input * denoise_mask + init_latent * (1.0f - denoise_mask); } @@ -2374,7 +2376,7 @@ class StableDiffusionGGML { if (sd_version_is_dit(version)) { if (sd_version_is_ltxav(version)) { latent_channel = 128; - } else if (version == VERSION_WAN2_2_TI2V) { + } else if (sd_version_is_wan_ti2v_family(version)) { latent_channel = 48; } else if (version == VERSION_HIDREAM_O1) { latent_channel = 3; @@ -3101,10 +3103,18 @@ struct sd_ctx_t { }; static bool sd_version_supports_video_generation(SDVersion version) { + // ABot-World loads like a Wan model but is causal/interactive-only: it + // supports neither batch capability until the causal session API lands. + if (sd_version_is_abot_world(version)) { + return false; + } return version == VERSION_SVD || sd_version_is_wan(version) || sd_version_is_ltxav(version); } static bool sd_version_supports_image_generation(SDVersion version) { + if (sd_version_is_abot_world(version)) { + return false; + } return !sd_version_supports_video_generation(version); } @@ -3509,7 +3519,25 @@ struct GenerationRequest { valid = false; } + // ABot-World is a causal/interactive model: it must be driven block-by-block + // with a KV cache, per-block keyboard actions, and its distilled 4-step + // schedule. Both batch entrypoints (generate_image and generate_video) build + // a GenerationRequest, so rejecting here covers every batch door with one + // check; running these weights through the batch recipes would silently + // produce corrupted output. Model loading and inspection remain supported. + void reject_abot_world_batch_generation(sd_ctx_t* sd_ctx) { + if (!sd_version_is_abot_world(sd_ctx->sd->version)) { + return; + } + LOG_ERROR( + "ABot-World models are not supported by batch generate_image()/generate_video(); " + "the causal interactive session API is not implemented yet " + "(model loading and inspection are supported)"); + valid = false; + } + void resolve(sd_ctx_t* sd_ctx) { + reject_abot_world_batch_generation(sd_ctx); align_generation_request_size(); resolve_hires(); seed = resolve_seed(seed); @@ -5392,6 +5420,8 @@ SD_API bool generate_video(sd_ctx_t* sd_ctx, } int64_t t0 = ggml_time_ms(); sd_ctx->sd->vae_tiling_params = sd_vid_gen_params->vae_tiling_params; + // ABot-World (causal/interactive-only) is rejected inside + // GenerationRequest::resolve(), which covers generate_image() too. GenerationRequest request(sd_ctx, sd_vid_gen_params); if (!request.valid) { return false; @@ -5747,3 +5777,117 @@ SD_API bool generate_video(sd_ctx_t* sd_ctx, } return true; } + +/* ABot-World interactive walk session ------------------------------------- */ + +struct sd_abot_session_t { + SDBackendManager backend_manager; + ABOT::AbotWalkSession session; +}; + +void sd_abot_session_params_init(sd_abot_session_params_t* params) { + if (params == nullptr) { + return; + } + memset(params, 0, sizeof(sd_abot_session_params_t)); + params->n_threads = -1; + params->seed = 42; + params->num_frame_per_block = 0; // model default + params->local_attn_size = 0; // config default +} + +sd_abot_session_t* sd_abot_session_new(const sd_abot_session_params_t* params) { + if (params == nullptr || + strlen(SAFE_STR(params->dit_model_path)) == 0 || + strlen(SAFE_STR(params->taehv_path)) == 0 || + strlen(SAFE_STR(params->scene_path)) == 0) { + LOG_ERROR("sd_abot_session_new: dit_model_path, taehv_path and scene_path are required"); + return nullptr; + } + auto* s = new (std::nothrow) sd_abot_session_t(); + if (s == nullptr) { + return nullptr; + } + std::string err; + if (!s->backend_manager.init(SAFE_STR(params->backend), nullptr, + params->offload_params_to_cpu, + false, false, false, &err)) { + LOG_ERROR("sd_abot_session_new: backend init failed: %s", err.c_str()); + delete s; + return nullptr; + } + ABOT::AbotWorldConfig cfg; + if (params->num_frame_per_block > 0) { + cfg.num_frame_per_block = params->num_frame_per_block; + } + if (params->local_attn_size > 0) { + cfg.local_attn_size = params->local_attn_size; + } + int n_threads = params->n_threads; + if (n_threads <= 0) { + n_threads = sd_get_num_physical_cores(); + } + if (!s->session.load(s->backend_manager.runtime_backend(SDBackendModule::DIFFUSION), + s->backend_manager.params_backend(SDBackendModule::DIFFUSION), + s->backend_manager.runtime_backend(SDBackendModule::VAE), + s->backend_manager.params_backend(SDBackendModule::VAE), + SAFE_STR(params->dit_model_path), + SAFE_STR(params->taehv_path), + SAFE_STR(params->scene_path), + cfg, + static_cast(params->seed), + n_threads)) { + delete s; + return nullptr; + } + return s; +} + +sd_image_t* sd_abot_session_step(sd_abot_session_t* session, + uint32_t action_mask, + int* num_frames_out) { + if (num_frames_out != nullptr) { + *num_frames_out = 0; + } + if (session == nullptr) { + return nullptr; + } + std::vector> rgb_frames; + int64_t px_w = 0, px_h = 0; + if (!session->session.step(static_cast(action_mask & 0xFF), rgb_frames, px_w, px_h)) { + return nullptr; + } + sd_image_t* result = (sd_image_t*)calloc(rgb_frames.size(), sizeof(sd_image_t)); + if (result == nullptr) { + return nullptr; + } + for (size_t i = 0; i < rgb_frames.size(); i++) { + result[i].width = static_cast(px_w); + result[i].height = static_cast(px_h); + result[i].channel = 3; + result[i].data = (uint8_t*)malloc(rgb_frames[i].size()); + if (result[i].data == nullptr) { + sd_abot_session_frames_free(result, static_cast(i)); + return nullptr; + } + memcpy(result[i].data, rgb_frames[i].data(), rgb_frames[i].size()); + } + if (num_frames_out != nullptr) { + *num_frames_out = static_cast(rgb_frames.size()); + } + return result; +} + +void sd_abot_session_frames_free(sd_image_t* frames, int num_frames) { + if (frames == nullptr) { + return; + } + for (int i = 0; i < num_frames; i++) { + free(frames[i].data); + } + free(frames); +} + +void sd_abot_session_free(sd_abot_session_t* session) { + delete session; +} diff --git a/src/tae.hpp b/src/tae.hpp index f43ef48fb..a945f9d86 100644 --- a/src/tae.hpp +++ b/src/tae.hpp @@ -302,7 +302,8 @@ class WideMemBlock : public GGMLBlock { } }; -ggml_tensor* +// inline: this header is included from more than one translation unit +inline ggml_tensor* patchify(ggml_context* ctx, ggml_tensor* x, int64_t patch_size, @@ -332,10 +333,10 @@ patchify(ggml_context* ctx, return x; } -ggml_tensor* unpatchify(ggml_context* ctx, - ggml_tensor* x, - int64_t patch_size, - int64_t b = 1) { +inline ggml_tensor* unpatchify(ggml_context* ctx, + ggml_tensor* x, + int64_t patch_size, + int64_t b = 1) { // x: [f, b*c*r*q, h, w] // return: [f, b*c, h*q, w*r] if (patch_size == 1) { @@ -525,7 +526,7 @@ class TAEHV : public GGMLBlock { TAEHV(bool decode_only = true, SDVersion version = VERSION_WAN2, bool is_wide = false) : decode_only(decode_only), version(version), is_wide(is_wide) { int patch = 1; - if (version == VERSION_WAN2_2_TI2V) { + if (sd_version_is_wan_ti2v_family(version)) { z_channels = 48; patch = 2; } else if (sd_version_is_ltxav(version)) { diff --git a/src/vae.hpp b/src/vae.hpp index 9f4ca1cc3..3bca0450b 100644 --- a/src/vae.hpp +++ b/src/vae.hpp @@ -69,7 +69,7 @@ struct VAE : public GGMLRunner { int scale_factor = 8; if (version == VERSION_LTXAV) { scale_factor = 32; - } else if (version == VERSION_WAN2_2_TI2V) { + } else if (sd_version_is_wan_ti2v_family(version)) { scale_factor = 16; } else if (sd_version_uses_flux2_vae(version)) { scale_factor = 16; diff --git a/src/wan.hpp b/src/wan.hpp index 68f020e25..57378e638 100644 --- a/src/wan.hpp +++ b/src/wan.hpp @@ -166,32 +166,36 @@ namespace WAN { if (feat_cache.size() > 0) { int idx = feat_idx; feat_idx += 1; - if (chunk_idx == 0) { - // feat_cache[idx] == nullptr, pass + auto time_conv_it = blocks.find("time_conv"); + GGML_ASSERT(time_conv_it != blocks.end()); + auto time_conv = std::dynamic_pointer_cast(time_conv_it->second); + GGML_ASSERT(time_conv != nullptr); + + // mirrors the reference (vae2_2.py Resample.forward): + // cache_x = x[:, :, -CACHE_T:].clone() + // if feat_cache[idx] == "Rep": x = time_conv(x) # chunk 0: causal zero pad + // else: ; x = time_conv(x, feat_cache[idx]) + // feat_cache[idx] = cache_x + auto cache_x = ggml_ext_slice(ctx->ggml_ctx, x, 2, -CACHE_T, x->ne[2]); + if (chunk_idx == 0) { // "Rep" + x = time_conv->forward(ctx, x); } else { - auto time_conv = std::dynamic_pointer_cast(blocks["time_conv"]); - - auto cache_x = ggml_ext_slice(ctx->ggml_ctx, x, 2, -CACHE_T, x->ne[2]); - if (cache_x->ne[2] < 2 && feat_cache[idx] != nullptr) { // chunk_idx >= 2 + if (cache_x->ne[2] < 2 && feat_cache[idx] != nullptr) { // cache last frame of last two chunk cache_x = ggml_concat(ctx->ggml_ctx, ggml_ext_slice(ctx->ggml_ctx, feat_cache[idx], 2, -1, feat_cache[idx]->ne[2]), cache_x, 2); } - if (chunk_idx == 1 && cache_x->ne[2] < 2) { // Rep - cache_x = ggml_pad_ext(ctx->ggml_ctx, cache_x, 0, 0, 0, 0, (int)cache_x->ne[2], 0, 0, 0); - // aka cache_x = torch.cat([torch.zeros_like(cache_x).to(cache_x.device),cache_x],dim=2) - } - if (chunk_idx == 1) { - x = time_conv->forward(ctx, x); - } else { - x = time_conv->forward(ctx, x, feat_cache[idx]); - } - feat_cache[idx] = cache_x; - x = ggml_reshape_4d(ctx->ggml_ctx, x, w * h, t, c, 2); // (2, c, t, h*w) - x = ggml_ext_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, x, 0, 3, 1, 2)); // (c, t, 2, h*w) - x = ggml_reshape_4d(ctx->ggml_ctx, x, w, h, 2 * t, c); // (c, t*2, h, w) + x = time_conv->forward(ctx, x, feat_cache[idx]); + } + feat_cache[idx] = cache_x; + x = ggml_reshape_4d(ctx->ggml_ctx, x, w * h, t, c, 2); // (2, c, t, h*w) + x = ggml_ext_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, x, 0, 3, 1, 2)); // (c, t, 2, h*w) + x = ggml_reshape_4d(ctx->ggml_ctx, x, w, h, 2 * t, c); // (c, t*2, h, w) + if (chunk_idx == 0) { + // reference: if first_chunk: x = x[:, :, 1:] + x = ggml_ext_slice(ctx->ggml_ctx, x, 2, 1, x->ne[2]); } } } @@ -1132,7 +1136,7 @@ namespace WAN { const std::string prefix = "", bool decode_only = false, SDVersion version = VERSION_WAN2) - : decode_only(decode_only), ae(decode_only, version == VERSION_WAN2_2_TI2V), VAE(version, backend, params_backend) { + : decode_only(decode_only), ae(decode_only, sd_version_is_wan_ti2v_family(version)), VAE(version, backend, params_backend) { ae.init(params_ctx, tensor_storage_map, prefix); } @@ -1607,10 +1611,12 @@ namespace WAN { ggml_tensor* e, ggml_tensor* pe, ggml_tensor* context, - int64_t context_img_len = 257) { + int64_t context_img_len = 257, + ggml_tensor* attn_mask = nullptr) { // x: [N, n_token, dim] // e: [N, 6, dim] or [N, T, 6, dim] // context: [N, context_img_len + context_txt_len, dim] + // attn_mask: optional additive self-attention mask (ABot causal walk) // return [N, n_token, dim] auto modulation = params["modulation"]; @@ -1629,7 +1635,7 @@ namespace WAN { auto y = norm1->forward(ctx, x); y = ggml_add(ctx->ggml_ctx, y, modulate_mul(ctx->ggml_ctx, y, es[1])); y = modulate_add(ctx->ggml_ctx, y, es[0]); - y = self_attn->forward(ctx, y, pe); + y = self_attn->forward(ctx, y, pe, attn_mask); x = ggml_add(ctx->ggml_ctx, x, modulate_mul(ctx->ggml_ctx, y, es[2])); @@ -1811,6 +1817,10 @@ namespace WAN { int64_t out_dim = 16; int64_t num_heads = 16; int num_layers = 32; + // ABot-World (causal Wan2.2-TI2V-5B + keyboard-action conditioning) + bool abot_world = false; + int64_t act_in_dim = 32; // 8 keys x 4 (repeat_interleave) + int act_downscale_factor = 16; // pixel-unshuffle factor before conv int vace_layers = 0; int64_t vace_in_dim = 96; std::map vace_layers_mapping = {}; @@ -1824,6 +1834,38 @@ namespace WAN { int64_t axes_dim_sum = 128; }; + // ABot-World action-conditioning adapter ("SimpleAdapter" in the reference + // implementation, github.com/amap-cvlab/ABot-World wan/modules/model.py): + // pixel_unshuffle(x16) -> Conv2d(act_in_dim*16*16 -> dim, k=2, s=2) + // -> ResidualBlock(conv1 3x3 p1 -> ReLU -> conv2 3x3 p1 -> +residual) + // The pixel-unshuffle is a host-side data rearrangement of the broadcast + // keyboard-action planes; forward() expects the already-unshuffled tensor + // [W/16, H/16, act_in_dim*256, N*T] and returns [W/32, H/32, dim, N*T], + // i.e. one feature per latent-token cell (matches patch_embedding grid). + class ActControlAdapter : public GGMLBlock { + public: + ActControlAdapter(int64_t act_in_dim, int64_t dim, int downscale_factor) { + int64_t unshuffled_dim = act_in_dim * downscale_factor * downscale_factor; + blocks["conv"] = std::shared_ptr(new Conv2d(unshuffled_dim, dim, {2, 2}, {2, 2})); + blocks["residual_blocks.0.conv1"] = std::shared_ptr(new Conv2d(dim, dim, {3, 3}, {1, 1}, {1, 1})); + blocks["residual_blocks.0.conv2"] = std::shared_ptr(new Conv2d(dim, dim, {3, 3}, {1, 1}, {1, 1})); + } + + ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) { + auto conv = std::dynamic_pointer_cast(blocks["conv"]); + auto conv1 = std::dynamic_pointer_cast(blocks["residual_blocks.0.conv1"]); + auto conv2 = std::dynamic_pointer_cast(blocks["residual_blocks.0.conv2"]); + + x = conv->forward(ctx, x); + auto residual = x; + x = conv1->forward(ctx, x); + x = ggml_relu_inplace(ctx->ggml_ctx, x); + x = conv2->forward(ctx, x); + x = ggml_add(ctx->ggml_ctx, x, residual); + return x; + } + }; + class Wan : public GGMLBlock { protected: WanParams params; @@ -1868,6 +1910,12 @@ namespace WAN { blocks["img_emb"] = std::shared_ptr(new MLPProj(1280, params.dim, params.flf_pos_embed_token_number)); } + // ABot-World keyboard-action adapter + if (params.abot_world) { + blocks["act_control_adapter"] = std::shared_ptr( + new ActControlAdapter(params.act_in_dim, params.dim, params.act_downscale_factor)); + } + // vace if (params.vace_layers > 0) { for (int i = 0; i < params.vace_layers; i++) { @@ -2141,8 +2189,15 @@ namespace WAN { } if (wan_params.num_layers == 30) { - if (version == VERSION_WAN2_2_TI2V) { - desc = "Wan2.2-TI2V-5B"; + if (version == VERSION_WAN2_2_TI2V || sd_version_is_abot_world(version)) { + // ABot-World-0-5B-LF shares the Wan2.2-TI2V-5B backbone and + // adds the act_control_adapter (keyboard-action conditioning). + if (sd_version_is_abot_world(version)) { + desc = "ABot-World-5B"; + wan_params.abot_world = true; + } else { + desc = "Wan2.2-TI2V-5B"; + } wan_params.dim = 3072; wan_params.eps = 1e-06f; wan_params.ffn_dim = 14336;