Feature abot dit

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.

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"

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;
         }

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;

src/stable-diffusion.cpp

@@ -86,6 +86,7 @@ const char* model_version_to_str[] = {
     "Longcat-Image",
     "PiD",
     "Ideogram 4",
+    "ABot-World",
 };
 
 const char* sampling_methods_str[] = {
@@ -2125,7 +2126,7 @@ class StableDiffusionGGML {
             sd::Tensor<float> timesteps_tensor({static_cast<int64_t>(timesteps_vec.size())}, timesteps_vec);
             sd::Tensor<float> guidance_tensor({1}, std::vector<float>{guidance.distilled_guidance});
             sd::Tensor<float> 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 +2375,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 +3102,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 +3518,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 +5419,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;

src/tae.hpp

@@ -525,7 +525,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)) {

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;