feat: add LTX rational latent upscaler (#1549)

* Temporal tile size + overlap

* add --extra-tiling-args support

---------

Co-authored-by: leejet <leejet714@gmail.com>

examples/cli/README.md

@@ -107,6 +107,8 @@ Generation Options:
   --extra-sample-args <string>             extra sampler/scheduler args, key=value list. lcm supports noise_clip_std,
                                            noise_scale_start, noise_scale_end; ltx2 supports max_shift, base_shift,
                                            stretch, terminal; euler_ge supports gamma
+  --extra-tiling-args <string>             extra VAE tiling args, key=value list. LTX video VAE supports
+                                           temporal_tile_frames (default: 4), temporal_tile_overlap (default: 1)
   -H, --height <int>                       image height, in pixel space (default: 512)
   -W, --width <int>                        image width, in pixel space (default: 512)
   --steps <int>                            number of sample steps (default: 20)

examples/common/common.cpp

@@ -835,6 +835,10 @@ ArgOptions SDGenerationParams::get_options() {
          "--extra-sample-args",
          "extra sampler/scheduler args, key=value list. lcm supports noise_clip_std, noise_scale_start, noise_scale_end; ltx2 supports max_shift, base_shift, stretch, terminal; euler_ge supports gamma",
          &extra_sample_args},
+        {"",
+         "--extra-tiling-args",
+         "extra VAE tiling args, key=value list. LTX video VAE supports temporal_tile_frames (default: 4), temporal_tile_overlap (default: 1)",
+         &extra_tiling_args},
     };
 
     options.int_options = {
@@ -1780,6 +1784,9 @@ bool SDGenerationParams::from_json_str(
         if (tiling_json.contains("rel_size_y") && tiling_json["rel_size_y"].is_number()) {
             vae_tiling_params.rel_size_y = tiling_json["rel_size_y"];
         }
+        if (tiling_json.contains("extra_tiling_args") && tiling_json["extra_tiling_args"].is_string()) {
+            extra_tiling_args = tiling_json["extra_tiling_args"].get<std::string>();
+        }
     }
 
     if (!parse_lora_json_field(j, lora_path_resolver, lora_map, high_noise_lora_map)) {
@@ -2002,6 +2009,8 @@ bool SDGenerationParams::initialize_cache_params() {
 }
 
 bool SDGenerationParams::resolve(const std::string& lora_model_dir, const std::string& hires_upscalers_dir, bool strict) {
+    vae_tiling_params.extra_tiling_args = extra_tiling_args.empty() ? nullptr : extra_tiling_args.c_str();
+
     if (high_noise_sample_params.sample_steps <= 0) {
         high_noise_sample_params.sample_steps = -1;
     }
@@ -2188,6 +2197,7 @@ sd_img_gen_params_t SDGenerationParams::to_sd_img_gen_params_t() {
     sample_params.custom_sigmas_count                 = static_cast<int>(custom_sigmas.size());
     sample_params.extra_sample_args                   = extra_sample_args.empty() ? nullptr : extra_sample_args.c_str();
     high_noise_sample_params.extra_sample_args        = high_noise_extra_sample_args.empty() ? nullptr : high_noise_extra_sample_args.c_str();
+    vae_tiling_params.extra_tiling_args               = extra_tiling_args.empty() ? nullptr : extra_tiling_args.c_str();
     cache_params.scm_mask                             = scm_mask.empty() ? nullptr : scm_mask.c_str();
 
     sd_pm_params_t pm_params = {
@@ -2261,6 +2271,7 @@ sd_vid_gen_params_t SDGenerationParams::to_sd_vid_gen_params_t() {
     sample_params.custom_sigmas_count                 = static_cast<int>(custom_sigmas.size());
     sample_params.extra_sample_args                   = extra_sample_args.empty() ? nullptr : extra_sample_args.c_str();
     high_noise_sample_params.extra_sample_args        = high_noise_extra_sample_args.empty() ? nullptr : high_noise_extra_sample_args.c_str();
+    vae_tiling_params.extra_tiling_args               = extra_tiling_args.empty() ? nullptr : extra_tiling_args.c_str();
     cache_params.scm_mask                             = scm_mask.empty() ? nullptr : scm_mask.c_str();
 
     params.loras                     = lora_vec.empty() ? nullptr : lora_vec.data();
@@ -2386,7 +2397,8 @@ std::string SDGenerationParams::to_string() const {
         << vae_tiling_params.tile_size_y << ", "
         << vae_tiling_params.target_overlap << ", "
         << vae_tiling_params.rel_size_x << ", "
-        << vae_tiling_params.rel_size_y << " },\n"
+        << vae_tiling_params.rel_size_y << ", "
+        << "\"" << extra_tiling_args << "\" },\n"
         << "}";
     return oss.str();
 }
@@ -2565,14 +2577,18 @@ std::string build_sdcpp_image_metadata_json(const SDContextParams& ctx_params,
         };
     }
 
-    if (gen_params.vae_tiling_params.enabled) {
+    if (gen_params.vae_tiling_params.enabled ||
+        gen_params.vae_tiling_params.temporal_tiling ||
+        !gen_params.extra_tiling_args.empty()) {
         root["vae_tiling"] = {
             {"enabled", gen_params.vae_tiling_params.enabled},
+            {"temporal_tiling", gen_params.vae_tiling_params.temporal_tiling},
             {"tile_size_x", gen_params.vae_tiling_params.tile_size_x},
             {"tile_size_y", gen_params.vae_tiling_params.tile_size_y},
             {"target_overlap", gen_params.vae_tiling_params.target_overlap},
             {"rel_size_x", gen_params.vae_tiling_params.rel_size_x},
             {"rel_size_y", gen_params.vae_tiling_params.rel_size_y},
+            {"extra_tiling_args", gen_params.extra_tiling_args},
         };
     }
 

examples/common/common.h

@@ -189,7 +189,8 @@ struct SDGenerationParams {
     int video_frames                     = 1;
     int fps                              = 16;
     float vace_strength                  = 1.f;
-    sd_tiling_params_t vae_tiling_params = {false, false, 0, 0, 0.5f, 0.0f, 0.0f};
+    sd_tiling_params_t vae_tiling_params = {false, false, 0, 0, 0.5f, 0.0f, 0.0f, nullptr};
+    std::string extra_tiling_args;
 
     std::string pm_id_images_dir;
     std::string pm_id_embed_path;

examples/server/README.md

@@ -209,6 +209,8 @@ Default Generation Options:
   --extra-sample-args <string>             extra sampler/scheduler args, key=value list. lcm supports noise_clip_std,
                                            noise_scale_start, noise_scale_end; ltx2 supports max_shift, base_shift,
                                            stretch, terminal; euler_ge supports gamma
+  --extra-tiling-args <string>             extra VAE tiling args, key=value list. LTX video VAE supports
+                                           temporal_tile_frames (default: 4), temporal_tile_overlap (default: 1)
   -H, --height <int>                       image height, in pixel space (default: 512)
   -W, --width <int>                        image width, in pixel space (default: 512)
   --steps <int>                            number of sample steps (default: 20)
@@ -264,6 +266,7 @@ Default Generation Options:
   --disable-auto-resize-ref-image          disable auto resize of ref images
   --disable-image-metadata                 do not embed generation metadata on image files
   --vae-tiling                             process vae in tiles to reduce memory usage
+  --temporal-tiling                        enable temporal tiling for LTX video VAE decode
   --hires                                  enable highres fix
   -s, --seed                               RNG seed (default: 42, use random seed for < 0)
   --sampling-method                        sampling method, one of [euler, euler_a, heun, dpm2, dpm++2s_a, dpm++2m,

examples/server/api.md

@@ -504,11 +504,13 @@ Shared default fields used by both `img_gen` and `vid_gen`:
 | `sample_params.guidance.slg.scale` | `number` |
 | `vae_tiling_params` | `object` |
 | `vae_tiling_params.enabled` | `boolean` |
+| `vae_tiling_params.temporal_tiling` | `boolean` |
 | `vae_tiling_params.tile_size_x` | `integer` |
 | `vae_tiling_params.tile_size_y` | `integer` |
 | `vae_tiling_params.target_overlap` | `number` |
 | `vae_tiling_params.rel_size_x` | `number` |
 | `vae_tiling_params.rel_size_y` | `number` |
+| `vae_tiling_params.extra_tiling_args` | `string` |
 | `cache_mode` | `string` |
 | `cache_option` | `string` |
 | `scm_mask` | `string` |
@@ -516,6 +518,8 @@ Shared default fields used by both `img_gen` and `vid_gen`:
 | `output_format` | `string` |
 | `output_compression` | `integer` |
 
+`vae_tiling_params.extra_tiling_args` accepts a key=value list. For LTX video VAE temporal tiling, `temporal_tile_frames` defaults to `4` and `temporal_tile_overlap` defaults to `1`.
+
 `img_gen`-specific default fields:
 
 | Field | Type |
@@ -692,11 +696,13 @@ Example:
 
   "vae_tiling_params": {
     "enabled": false,
+    "temporal_tiling": false,
     "tile_size_x": 0,
     "tile_size_y": 0,
     "target_overlap": 0.5,
     "rel_size_x": 0.0,
-    "rel_size_y": 0.0
+    "rel_size_y": 0.0,
+    "extra_tiling_args": ""
   },
 
   "cache_mode": "disabled",
@@ -804,6 +810,14 @@ Other native fields:
 | `hires.custom_sigmas` | `array<number>` |
 | `hires.upscale_tile_size` | `integer` |
 | `vae_tiling_params` | `object` |
+| `vae_tiling_params.enabled` | `boolean` |
+| `vae_tiling_params.temporal_tiling` | `boolean` |
+| `vae_tiling_params.tile_size_x` | `integer` |
+| `vae_tiling_params.tile_size_y` | `integer` |
+| `vae_tiling_params.target_overlap` | `number` |
+| `vae_tiling_params.rel_size_x` | `number` |
+| `vae_tiling_params.rel_size_y` | `number` |
+| `vae_tiling_params.extra_tiling_args` | `string` |
 | `cache_mode` | `string` |
 | `cache_option` | `string` |
 | `scm_mask` | `string` |
@@ -1012,11 +1026,13 @@ Example:
 
   "vae_tiling_params": {
     "enabled": false,
+    "temporal_tiling": false,
     "tile_size_x": 0,
     "tile_size_y": 0,
     "target_overlap": 0.5,
     "rel_size_x": 0.0,
-    "rel_size_y": 0.0
+    "rel_size_y": 0.0,
+    "extra_tiling_args": ""
   },
 
   "cache_mode": "disabled",
@@ -1134,6 +1150,14 @@ Other native fields:
 | Field | Type |
 | --- | --- |
 | `vae_tiling_params` | `object` |
+| `vae_tiling_params.enabled` | `boolean` |
+| `vae_tiling_params.temporal_tiling` | `boolean` |
+| `vae_tiling_params.tile_size_x` | `integer` |
+| `vae_tiling_params.tile_size_y` | `integer` |
+| `vae_tiling_params.target_overlap` | `number` |
+| `vae_tiling_params.rel_size_x` | `number` |
+| `vae_tiling_params.rel_size_y` | `number` |
+| `vae_tiling_params.extra_tiling_args` | `string` |
 | `cache_mode` | `string` |
 | `cache_option` | `string` |
 | `scm_mask` | `string` |

examples/server/routes_sdcpp.cpp

@@ -56,11 +56,13 @@ static const char* capability_sample_method_name(enum sample_method_t sample_met
 static json make_vae_tiling_json(const sd_tiling_params_t& params) {
     return {
         {"enabled", params.enabled},
+        {"temporal_tiling", params.temporal_tiling},
         {"tile_size_x", params.tile_size_x},
         {"tile_size_y", params.tile_size_y},
         {"target_overlap", params.target_overlap},
         {"rel_size_x", params.rel_size_x},
         {"rel_size_y", params.rel_size_y},
+        {"extra_tiling_args", params.extra_tiling_args ? params.extra_tiling_args : ""},
     };
 }
 

include/stable-diffusion.h

@@ -160,6 +160,7 @@ typedef struct {
     float target_overlap;
     float rel_size_x;
     float rel_size_y;
+    const char* extra_tiling_args;
 } sd_tiling_params_t;
 
 typedef struct {

src/denoiser.hpp

@@ -496,84 +496,26 @@ struct LTX2Scheduler : SigmaScheduler {
         parse_extra_sample_args(extra_sample_args);
     }
 
-    static std::string trim(std::string value) {
-        const char* whitespace = " \t\r\n";
-        size_t begin           = value.find_first_not_of(whitespace);
-        if (begin == std::string::npos) {
-            return "";
-        }
-        size_t end = value.find_last_not_of(whitespace);
-        return value.substr(begin, end - begin + 1);
-    }
-
     void parse_extra_sample_args(const char* extra_sample_args) {
-        if (extra_sample_args == nullptr || extra_sample_args[0] == '\0') {
-            return;
-        }
-
-        std::string raw(extra_sample_args);
-        size_t start   = 0;
-        auto parse_arg = [&](const std::string& item) {
-            std::string token = trim(item);
-            if (token.empty()) {
-                return;
-            }
-            size_t eq = token.find('=');
-            if (eq == std::string::npos) {
-                LOG_WARN("ignoring invalid ltx2 scheduler arg '%s'", token.c_str());
-                return;
-            }
-
-            std::string key   = trim(token.substr(0, eq));
-            std::string value = trim(token.substr(eq + 1));
-            auto parse_float  = [&](float* out) -> bool {
-                try {
-                    size_t consumed = 0;
-                    float parsed    = std::stof(value, &consumed);
-                    if (!trim(value.substr(consumed)).empty()) {
-                        return false;
-                    }
-                    *out = parsed;
-                    return true;
-                } catch (const std::exception&) {
-                    return false;
+        for (const auto& [key, value] : parse_key_value_args(extra_sample_args, "ltx2 scheduler arg")) {
+            if (key == "max_shift") {
+                if (!parse_strict_float(value, max_shift)) {
+                    LOG_WARN("ignoring invalid ltx2 scheduler arg '%s=%s'", key.c_str(), value.c_str());
                 }
-            };
-            try {
-                if (key == "max_shift") {
-                    if (!parse_float(&max_shift)) {
-                        LOG_WARN("ignoring invalid ltx2 scheduler arg '%s'", token.c_str());
-                    }
-                } else if (key == "base_shift") {
-                    if (!parse_float(&base_shift)) {
-                        LOG_WARN("ignoring invalid ltx2 scheduler arg '%s'", token.c_str());
-                    }
-                } else if (key == "terminal") {
-                    if (!parse_float(&terminal)) {
-                        LOG_WARN("ignoring invalid ltx2 scheduler arg '%s'", token.c_str());
-                    }
-                } else if (key == "stretch") {
-                    std::string v = value;
-                    std::transform(v.begin(), v.end(), v.begin(), [](unsigned char c) { return static_cast<char>(std::tolower(c)); });
-                    if (v == "1" || v == "true" || v == "yes" || v == "on") {
-                        stretch = true;
-                    } else if (v == "0" || v == "false" || v == "no" || v == "off") {
-                        stretch = false;
-                    } else {
-                        LOG_WARN("ignoring invalid ltx2 scheduler arg '%s'", token.c_str());
-                    }
-                } else {
-                    LOG_WARN("ignoring unknown ltx2 scheduler arg '%s'", key.c_str());
+            } else if (key == "base_shift") {
+                if (!parse_strict_float(value, base_shift)) {
+                    LOG_WARN("ignoring invalid ltx2 scheduler arg '%s=%s'", key.c_str(), value.c_str());
                 }
-            } catch (const std::exception&) {
-                LOG_WARN("ignoring invalid ltx2 scheduler arg '%s'", token.c_str());
-            }
-        };
-
-        for (size_t pos = 0; pos <= raw.size(); ++pos) {
-            if (pos == raw.size() || raw[pos] == ',' || raw[pos] == ';') {
-                parse_arg(raw.substr(start, pos - start));
-                start = pos + 1;
+            } else if (key == "terminal") {
+                if (!parse_strict_float(value, terminal)) {
+                    LOG_WARN("ignoring invalid ltx2 scheduler arg '%s=%s'", key.c_str(), value.c_str());
+                }
+            } else if (key == "stretch") {
+                if (!parse_strict_bool(value, stretch)) {
+                    LOG_WARN("ignoring invalid ltx2 scheduler arg '%s=%s'", key.c_str(), value.c_str());
+                }
+            } else {
+                LOG_WARN("ignoring unknown ltx2 scheduler arg '%s'", key.c_str());
             }
         }
     }
@@ -1276,7 +1218,7 @@ static sd::Tensor<float> sample_dpmpp_2m_v2(denoise_cb_t model,
     return x;
 }
 
-using SamplerExtraArgs = std::vector<std::pair<std::string, std::string>>;
+using SamplerExtraArgs = KeyValueArgs;
 
 static sd::Tensor<float> sample_lcm(denoise_cb_t model,
                                     sd::Tensor<float> x,
@@ -1296,15 +1238,8 @@ static sd::Tensor<float> sample_lcm(denoise_cb_t model,
 
     for (const auto& [key, value] : extra_sample_args) {
         float parsed = 0.0f;
-        try {
-            size_t consumed = 0;
-            parsed          = std::stof(value, &consumed);
-            if (trim(value.substr(consumed)).size() != 0) {
-                LOG_WARN("ignoring invalid lcm extra sample arg '%s'", key.c_str());
-                continue;
-            }
-        } catch (const std::exception&) {
-            LOG_WARN("ignoring invalid lcm extra sample arg '%s=%s'", key.c_str());
+        if (!parse_strict_float(value, parsed)) {
+            LOG_WARN("ignoring invalid lcm extra sample arg '%s=%s'", key.c_str(), value.c_str());
             continue;
         }
         if (key == "noise_clip_std") {
@@ -1861,15 +1796,8 @@ static sd::Tensor<float> sample_gradient_estimation(denoise_cb_t model,
 
     for (const auto& [key, value] : extra_sample_args) {
         float parsed = 0.0f;
-        try {
-            size_t consumed = 0;
-            parsed          = std::stof(value, &consumed);
-            if (trim(value.substr(consumed)).size() != 0) {
-                LOG_WARN("ignoring invalid euler_ge extra sample arg '%s'", key.c_str());
-                continue;
-            }
-        } catch (const std::exception&) {
-            LOG_WARN("ignoring invalid euler_ge extra sample arg '%s'", key.c_str());
+        if (!parse_strict_float(value, parsed)) {
+            LOG_WARN("ignoring invalid euler_ge extra sample arg '%s=%s'", key.c_str(), value.c_str());
             continue;
         }
         if (key == "gamma") {
@@ -1916,46 +1844,6 @@ static sd::Tensor<float> sample_gradient_estimation(denoise_cb_t model,
     return x;
 }
 
-static SamplerExtraArgs parse_sampler_args(const char* extra_sample_args) {
-    SamplerExtraArgs pairs;
-
-    if (extra_sample_args == nullptr || extra_sample_args[0] == '\0') {
-        return pairs;
-    }
-
-    auto trim = [](std::string value) -> std::string {
-        const char* whitespace = " \t\r\n";
-        size_t begin           = value.find_first_not_of(whitespace);
-        if (begin == std::string::npos) {
-            return "";
-        }
-        size_t end = value.find_last_not_of(whitespace);
-        return value.substr(begin, end - begin + 1);
-    };
-
-    std::string raw(extra_sample_args);
-    size_t start = 0;
-
-    for (size_t pos = 0; pos <= raw.size(); ++pos) {
-        if (pos == raw.size() || raw[pos] == ',' || raw[pos] == ';') {
-            std::string item  = raw.substr(start, pos - start);
-            std::string token = trim(item);
-
-            if (!token.empty()) {
-                size_t eq = token.find('=');
-                if (eq != std::string::npos) {
-                    std::string key   = trim(token.substr(0, eq));
-                    std::string value = trim(token.substr(eq + 1));
-                    pairs.emplace_back(std::move(key), std::move(value));
-                }
-            }
-            start = pos + 1;
-        }
-    }
-
-    return pairs;
-}
-
 // k diffusion reverse ODE: dx = (x - D(x;\sigma)) / \sigma dt; \sigma(t) = t
 static sd::Tensor<float> sample_k_diffusion(sample_method_t method,
                                             denoise_cb_t model,
@@ -1965,7 +1853,7 @@ static sd::Tensor<float> sample_k_diffusion(sample_method_t method,
                                             float eta,
                                             bool is_flow_denoiser,
                                             const char* extra_sample_args) {
-    SamplerExtraArgs extra_args = parse_sampler_args(extra_sample_args);
+    SamplerExtraArgs extra_args = parse_key_value_args(extra_sample_args, "extra sample arg");
     switch (method) {
         case EULER_A_SAMPLE_METHOD:
             return sample_euler_ancestral(model, std::move(x), sigmas, rng, is_flow_denoiser, eta);

src/ggml_extend.hpp

@@ -1602,6 +1602,23 @@ __STATIC_INLINE__ size_t ggml_tensor_num(ggml_context* ctx) {
     return num;
 }
 
+__STATIC_INLINE__ ggml_tensor* ggml_ext_vec_concat(ggml_context* ctx,
+                                                   std::vector<ggml_tensor*>& tensors,
+                                                   int dim) {
+    while (tensors.size() > 1) {
+        std::vector<ggml_tensor*> next_level;
+        for (size_t i = 0; i < tensors.size(); i += 2) {
+            if (i + 1 < tensors.size()) {
+                next_level.push_back(ggml_concat(ctx, tensors[i], tensors[i + 1], dim));
+            } else {
+                next_level.push_back(tensors[i]);
+            }
+        }
+        tensors = std::move(next_level);
+    }
+    return tensors[0];
+}
+
 /* SDXL with LoRA requires more space */
 #define MAX_PARAMS_TENSOR_NUM 32768
 #define MAX_GRAPH_SIZE 327680
@@ -3139,6 +3156,163 @@ public:
     }
 };
 
+class Conv2d_grouped : public UnaryBlock {
+protected:
+    int64_t in_channels;
+    int64_t out_channels;
+    int groups;
+    std::pair<int, int> kernel_size;
+    std::pair<int, int> stride;
+    std::pair<int, int> padding;
+    std::pair<int, int> dilation;
+    bool bias;
+    float scale = 1.f;
+    std::string prefix;
+
+    void init_params(ggml_context* ctx, const String2TensorStorage& tensor_storage_map, const std::string prefix = "") override {
+        this->prefix         = prefix;
+        enum ggml_type wtype = GGML_TYPE_F16;
+        params["weight"]     = ggml_new_tensor_4d(ctx, wtype, kernel_size.second, kernel_size.first, in_channels / groups, out_channels);
+        if (bias) {
+            enum ggml_type wtype = GGML_TYPE_F32;
+            params["bias"]       = ggml_new_tensor_1d(ctx, wtype, out_channels);
+        }
+    }
+
+public:
+    Conv2d_grouped(int64_t in_channels,
+                   int64_t out_channels,
+                   int groups,
+                   std::pair<int, int> kernel_size,
+                   std::pair<int, int> stride   = {1, 1},
+                   std::pair<int, int> padding  = {0, 0},
+                   std::pair<int, int> dilation = {1, 1},
+                   bool bias                    = true)
+        : in_channels(in_channels),
+          out_channels(out_channels),
+          groups(groups),
+          kernel_size(kernel_size),
+          stride(stride),
+          padding(padding),
+          dilation(dilation),
+          bias(bias) {}
+
+    void set_scale(float scale_value) {
+        scale = scale_value;
+    }
+
+    std::string get_desc() {
+        return "Conv2d_grouped";
+    }
+
+    ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) {
+        ggml_tensor* w = params["weight"];
+        ggml_tensor* b = nullptr;
+        if (bias) {
+            b = params["bias"];
+        }
+
+        if (groups == 1) {
+            if (ctx->weight_adapter) {
+                WeightAdapter::ForwardParams forward_params;
+                forward_params.op_type           = WeightAdapter::ForwardParams::op_type_t::OP_CONV2D;
+                forward_params.conv2d.s0         = stride.second;
+                forward_params.conv2d.s1         = stride.first;
+                forward_params.conv2d.p0         = padding.second;
+                forward_params.conv2d.p1         = padding.first;
+                forward_params.conv2d.d0         = dilation.second;
+                forward_params.conv2d.d1         = dilation.first;
+                forward_params.conv2d.direct     = ctx->conv2d_direct_enabled;
+                forward_params.conv2d.circular_x = ctx->circular_x_enabled;
+                forward_params.conv2d.circular_y = ctx->circular_y_enabled;
+                forward_params.conv2d.scale      = scale;
+                return ctx->weight_adapter->forward_with_lora(ctx->ggml_ctx, ctx->backend, x, w, b, prefix, forward_params);
+            }
+            return ggml_ext_conv_2d(ctx->ggml_ctx, x, w, b,
+                                    stride.second, stride.first,
+                                    padding.second, padding.first,
+                                    dilation.second, dilation.first,
+                                    ctx->conv2d_direct_enabled,
+                                    ctx->circular_x_enabled,
+                                    ctx->circular_y_enabled,
+                                    scale);
+        }
+
+        if (groups == in_channels && groups == out_channels) {
+            ggml_tensor* res;
+            if (ctx->conv2d_direct_enabled) {
+                res = ggml_conv_2d_dw_direct(ctx->ggml_ctx, x, w,
+                                             stride.second, stride.first,
+                                             padding.second, padding.first,
+                                             dilation.second, dilation.first);
+            } else {
+                res = ggml_conv_2d_dw(ctx->ggml_ctx, x, w,
+                                      stride.second, stride.first,
+                                      padding.second, padding.first,
+                                      dilation.second, dilation.first);
+            }
+            if (b) {
+                res = ggml_add(ctx->ggml_ctx, res, b);
+            }
+            return res;
+        }
+
+        int64_t ic_g = in_channels / groups;
+        int64_t oc_g = out_channels / groups;
+
+        std::vector<ggml_tensor*> out_slices(groups);
+
+        for (int i = 0; i < groups; ++i) {
+            size_t x_offset  = i * ic_g * x->nb[2];
+            ggml_tensor* x_i = ggml_view_4d(ctx->ggml_ctx, x,
+                                            x->ne[0], x->ne[1], ic_g, x->ne[3],
+                                            x->nb[1], x->nb[2], x->nb[3],
+                                            x_offset);
+
+            size_t w_offset  = i * oc_g * w->nb[3];
+            ggml_tensor* w_i = ggml_view_4d(ctx->ggml_ctx, w,
+                                            w->ne[0], w->ne[1], w->ne[2], oc_g,
+                                            w->nb[1], w->nb[2], w->nb[3],
+                                            w_offset);
+
+            ggml_tensor* b_i = nullptr;
+            if (b) {
+                size_t b_offset = i * oc_g * b->nb[0];
+                b_i             = ggml_view_1d(ctx->ggml_ctx, b, oc_g, b_offset);
+            }
+
+            if (ctx->weight_adapter) {
+                WeightAdapter::ForwardParams forward_params;
+                forward_params.op_type           = WeightAdapter::ForwardParams::op_type_t::OP_CONV2D;
+                forward_params.conv2d.s0         = stride.second;
+                forward_params.conv2d.s1         = stride.first;
+                forward_params.conv2d.p0         = padding.second;
+                forward_params.conv2d.p1         = padding.first;
+                forward_params.conv2d.d0         = dilation.second;
+                forward_params.conv2d.d1         = dilation.first;
+                forward_params.conv2d.direct     = ctx->conv2d_direct_enabled;
+                forward_params.conv2d.circular_x = ctx->circular_x_enabled;
+                forward_params.conv2d.circular_y = ctx->circular_y_enabled;
+                forward_params.conv2d.scale      = scale;
+                out_slices[i]                    = ctx->weight_adapter->forward_with_lora(ctx->ggml_ctx, ctx->backend, x_i, w_i, b_i, prefix, forward_params);
+            } else {
+                out_slices[i] = ggml_ext_conv_2d(ctx->ggml_ctx, x_i, w_i, b_i,
+                                                 stride.second, stride.first,
+                                                 padding.second, padding.first,
+                                                 dilation.second, dilation.first,
+                                                 ctx->conv2d_direct_enabled,
+                                                 ctx->circular_x_enabled,
+                                                 ctx->circular_y_enabled,
+                                                 scale);
+            }
+        }
+
+        ggml_tensor* out = ggml_ext_vec_concat(ctx->ggml_ctx, out_slices, 2);
+
+        return out;
+    }
+};
+
 class Conv3d : public UnaryBlock {
 protected:
     int64_t in_channels;

src/ltx_audio_vae.h

@@ -2,6 +2,7 @@
 #define __SD_LTX_AUDIO_VAE_H__
 
 #include <cmath>
+#include <limits>
 #include <numeric>
 #include <string>
 #include <vector>
@@ -171,90 +172,59 @@ namespace LTXV {
         }
     };
 
-    static sd::Tensor<float> squeeze_trailing_singleton_dims(sd::Tensor<float> tensor) {
-        while (tensor.dim() > 0 && tensor.shape().back() == 1) {
-            tensor = tensor.squeeze(static_cast<size_t>(tensor.dim() - 1));
-        }
-        return tensor;
-    }
+    static ggml_tensor* compute_log_mel_spectrogram(GGMLRunnerContext* runner_ctx,
+                                                    ggml_tensor* waveform,
+                                                    ggml_tensor* forward_basis,
+                                                    ggml_tensor* mel_basis,
+                                                    int hop_length) {
+        auto ctx = runner_ctx->ggml_ctx;
+        GGML_ASSERT(ctx != nullptr);
+        GGML_ASSERT(waveform != nullptr);
+        GGML_ASSERT(forward_basis != nullptr);
+        GGML_ASSERT(mel_basis != nullptr);
+        GGML_ASSERT(waveform->type == GGML_TYPE_F32);
+        GGML_ASSERT(forward_basis->type == GGML_TYPE_F32);
+        GGML_ASSERT(mel_basis->type == GGML_TYPE_F32);
+        GGML_ASSERT(forward_basis->ne[1] == 1);
 
-    static sd::Tensor<float> normalize_waveform_for_host(sd::Tensor<float> waveform) {
-        waveform = squeeze_trailing_singleton_dims(std::move(waveform));
-        if (waveform.empty()) {
-            return waveform;
-        }
-        if (waveform.dim() == 1) {
-            return waveform.reshape({waveform.shape()[0], 1, 1});
-        }
-        if (waveform.dim() == 2) {
-            return waveform.reshape({waveform.shape()[0], waveform.shape()[1], 1});
-        }
-        if (waveform.dim() == 3) {
-            return waveform;
-        }
-        throw std::runtime_error("Unsupported waveform rank for host processing: rank=" + std::to_string(waveform.dim()));
-    }
+        const int64_t time          = waveform->ne[0];
+        const int64_t channels      = waveform->ne[1];
+        const int64_t batch         = waveform->ne[2];
+        const int64_t filter_len    = forward_basis->ne[0];
+        const int64_t stft_channels = forward_basis->ne[2];
+        const int64_t n_freqs       = stft_channels / 2;
+        const int64_t n_mels        = mel_basis->ne[1];
+        const int64_t left_pad      = std::max<int64_t>(0, filter_len - hop_length);
+        const int64_t padded_time   = time + left_pad;
+        const int64_t frame_count   = padded_time < filter_len ? 0 : 1 + (padded_time - filter_len) / hop_length;
 
-    static sd::Tensor<float> load_param_tensor_f32(ggml_tensor* tensor) {
-        GGML_ASSERT(tensor != nullptr);
-        return squeeze_trailing_singleton_dims(sd::make_sd_tensor_from_ggml<float>(tensor));
-    }
+        GGML_ASSERT(stft_channels % 2 == 0);
+        GGML_ASSERT(mel_basis->ne[0] == n_freqs);
+        GGML_ASSERT(waveform->ne[3] == 1);
+        GGML_ASSERT(frame_count > 0);
 
-    static sd::Tensor<float> compute_log_mel_spectrogram(const sd::Tensor<float>& waveform_in,
-                                                         const sd::Tensor<float>& forward_basis,
-                                                         const sd::Tensor<float>& mel_basis,
-                                                         int hop_length) {
-        auto waveform = normalize_waveform_for_host(waveform_in);
-        GGML_ASSERT(forward_basis.dim() >= 3);
-        GGML_ASSERT(mel_basis.dim() >= 2);
-
-        const int64_t time        = waveform.shape()[0];
-        const int64_t channels    = waveform.shape()[1];
-        const int64_t batch       = waveform.shape()[2];
-        const int64_t filter_len  = forward_basis.shape()[0];
-        const int64_t basis_freq2 = forward_basis.shape().back();
-        const int64_t n_freqs     = basis_freq2 / 2;
-        const int64_t n_mels      = mel_basis.shape()[1];
-        const int64_t left_pad    = std::max<int64_t>(0, filter_len - hop_length);
-        const int64_t padded_time = time + left_pad;
-        const int64_t frame_count = padded_time < filter_len ? 0 : 1 + (padded_time - filter_len) / hop_length;
-
-        sd::Tensor<float> log_mel({n_mels, frame_count, channels, batch});
-        std::vector<float> padded(static_cast<size_t>(padded_time), 0.0f);
-        std::vector<float> magnitude(static_cast<size_t>(n_freqs), 0.0f);
-
-        for (int64_t b = 0; b < batch; ++b) {
-            for (int64_t c = 0; c < channels; ++c) {
-                std::fill(padded.begin(), padded.end(), 0.0f);
-                for (int64_t t = 0; t < time; ++t) {
-                    padded[static_cast<size_t>(t + left_pad)] = waveform.index(t, c, b);
-                }
-
-                for (int64_t frame = 0; frame < frame_count; ++frame) {
-                    const int64_t frame_offset = frame * hop_length;
-                    for (int64_t f = 0; f < n_freqs; ++f) {
-                        double real = 0.0;
-                        double imag = 0.0;
-                        for (int64_t k = 0; k < filter_len; ++k) {
-                            const float sample = padded[static_cast<size_t>(frame_offset + k)];
-                            real += static_cast<double>(sample) * static_cast<double>(forward_basis.index(k, 0, f));
-                            imag += static_cast<double>(sample) * static_cast<double>(forward_basis.index(k, 0, f + n_freqs));
-                        }
-                        magnitude[static_cast<size_t>(f)] = static_cast<float>(std::sqrt(real * real + imag * imag));
-                    }
-
-                    for (int64_t m = 0; m < n_mels; ++m) {
-                        double mel_value = 0.0;
-                        for (int64_t f = 0; f < n_freqs; ++f) {
-                            mel_value += static_cast<double>(mel_basis.index(f, m)) * static_cast<double>(magnitude[static_cast<size_t>(f)]);
-                        }
-                        log_mel.index(m, frame, c, b) = static_cast<float>(std::log(std::max(mel_value, 1e-5)));
-                    }
-                }
-            }
+        auto x = ggml_reshape_3d(ctx, waveform, time, 1, channels * batch);
+        if (left_pad > 0) {
+            x = ggml_pad_ext(ctx, x, static_cast<int>(left_pad), 0, 0, 0, 0, 0, 0, 0);
         }
 
-        return log_mel;
+        auto frames = ggml_conv_1d(ctx, forward_basis, x, hop_length, 0, 1);
+        GGML_ASSERT(frames->ne[0] == frame_count);
+        GGML_ASSERT(frames->ne[1] == stft_channels);
+        GGML_ASSERT(frames->ne[2] == channels * batch);
+
+        auto real      = ggml_ext_slice(ctx, frames, 1, 0, n_freqs);
+        auto imag      = ggml_ext_slice(ctx, frames, 1, n_freqs, stft_channels);
+        auto magnitude = ggml_sqrt(ctx,
+                                   ggml_add(ctx,
+                                            ggml_sqr(ctx, real),
+                                            ggml_sqr(ctx, imag)));
+
+        magnitude = ggml_cont(ctx, ggml_permute(ctx, magnitude, 1, 0, 2, 3));
+        auto mel  = ggml_mul_mat(ctx, mel_basis, magnitude);
+        mel       = ggml_log(ctx, ggml_clamp(ctx, mel, 1e-5f, std::numeric_limits<float>::max()));
+
+        return ggml_reshape_4d(ctx, mel, n_mels, frame_count, channels, batch);
     }
 
     static std::vector<float> build_hann_resample_filter(int ratio) {
@@ -276,75 +246,6 @@ namespace LTXV {
         return filter;
     }
 
-    static sd::Tensor<float> upsample_waveform_hann(const sd::Tensor<float>& waveform_in, int ratio) {
-        auto waveform = normalize_waveform_for_host(waveform_in);
-        if (ratio <= 1) {
-            return waveform;
-        }
-
-        const int lowpass_filter_width = 6;
-        const double rolloff           = 0.99;
-        const int width                = static_cast<int>(std::ceil(static_cast<double>(lowpass_filter_width) / rolloff));
-        const int kernel_size          = 2 * width * ratio + 1;
-        const int pad                  = width;
-        const int pad_left             = 2 * width * ratio;
-        const int pad_right            = kernel_size - ratio;
-        const int64_t time             = waveform.shape()[0];
-        const int64_t channels         = waveform.shape()[1];
-        const int64_t batch            = waveform.shape()[2];
-        const int64_t padded_time      = time + 2 * pad;
-        const int64_t conv_out_time    = (padded_time - 1) * ratio + kernel_size;
-        const int64_t cropped_time     = conv_out_time - pad_left - pad_right;
-        auto filter                    = build_hann_resample_filter(ratio);
-
-        sd::Tensor<float> output({cropped_time, channels, batch});
-        std::vector<float> padded(static_cast<size_t>(padded_time), 0.0f);
-        std::vector<float> conv_out(static_cast<size_t>(conv_out_time), 0.0f);
-
-        for (int64_t b = 0; b < batch; ++b) {
-            for (int64_t c = 0; c < channels; ++c) {
-                std::fill(padded.begin(), padded.end(), 0.0f);
-                const float first = waveform.index(0, c, b);
-                const float last  = waveform.index(time - 1, c, b);
-                for (int i = 0; i < pad; ++i) {
-                    padded[static_cast<size_t>(i)]              = first;
-                    padded[static_cast<size_t>(pad + time + i)] = last;
-                }
-                for (int64_t t = 0; t < time; ++t) {
-                    padded[static_cast<size_t>(pad + t)] = waveform.index(t, c, b);
-                }
-
-                std::fill(conv_out.begin(), conv_out.end(), 0.0f);
-                for (int64_t t = 0; t < padded_time; ++t) {
-                    const double sample    = static_cast<double>(padded[static_cast<size_t>(t)]) * ratio;
-                    const int64_t out_base = t * ratio;
-                    for (int k = 0; k < kernel_size; ++k) {
-                        conv_out[static_cast<size_t>(out_base + k)] += static_cast<float>(sample * filter[static_cast<size_t>(k)]);
-                    }
-                }
-
-                for (int64_t t = 0; t < cropped_time; ++t) {
-                    output.index(t, c, b) = conv_out[static_cast<size_t>(t + pad_left)];
-                }
-            }
-        }
-
-        return output;
-    }
-
-    static sd::Tensor<float> crop_waveform_samples(const sd::Tensor<float>& waveform_in, int64_t target_samples) {
-        auto waveform = normalize_waveform_for_host(waveform_in);
-        if (waveform.shape()[0] == target_samples) {
-            return waveform;
-        }
-        if (waveform.shape()[0] > target_samples) {
-            return sd::ops::slice(waveform, 0, 0, target_samples);
-        }
-        sd::Tensor<float> output({target_samples, waveform.shape()[1], waveform.shape()[2]});
-        sd::ops::slice_assign(&output, 0, 0, waveform.shape()[0], waveform);
-        return output;
-    }
-
     static ggml_type audio_conv_weight_type(ggml_type type) {
         return type == GGML_TYPE_BF16 ? GGML_TYPE_F16 : type;
     }
@@ -413,22 +314,101 @@ namespace LTXV {
         return ggml_reshape_4d(ctx, out, out->ne[0], out->ne[1], 1, 1);
     }
 
+    static ggml_tensor* reverse_1d_filter(ggml_context* ctx, ggml_tensor* filter) {
+        GGML_ASSERT(ctx != nullptr);
+        GGML_ASSERT(filter != nullptr);
+        GGML_ASSERT(filter->ne[1] == 1);
+        GGML_ASSERT(filter->ne[2] == 1);
+        GGML_ASSERT(filter->ne[3] == 1);
+
+        ggml_tensor* reversed = nullptr;
+        for (int64_t k = filter->ne[0] - 1; k >= 0; --k) {
+            auto slice = ggml_ext_slice(ctx, filter, 0, k, k + 1);
+            reversed   = reversed == nullptr ? slice : ggml_concat(ctx, reversed, slice, 0);
+        }
+        return reversed;
+    }
+
     static ggml_tensor* depthwise_conv_transpose1d(ggml_context* ctx,
                                                    ggml_tensor* x,
                                                    ggml_tensor* filter,
                                                    int stride) {
         GGML_ASSERT(x->ne[2] == 1 && x->ne[3] == 1);
         GGML_ASSERT(filter->ne[1] == 1);
+        GGML_ASSERT(filter->ne[2] == 1 && filter->ne[3] == 1);
 
-        ggml_tensor* out = nullptr;
-        for (int64_t c = 0; c < x->ne[1]; ++c) {
-            auto xi = ggml_ext_slice(ctx, x, 1, c, c + 1);
-            auto yi = ggml_conv_transpose_1d(ctx, filter, xi, stride, 0, 1);
-            yi      = ggml_ext_scale(ctx, yi, static_cast<float>(stride));
-            yi      = ggml_reshape_4d(ctx, yi, yi->ne[0], 1, 1, 1);
-            out     = out == nullptr ? yi : ggml_concat(ctx, out, yi, 1);
+        const int64_t time        = x->ne[0];
+        const int64_t channels    = x->ne[1];
+        const int64_t kernel_size = filter->ne[0];
+        const int64_t out_time    = (time - 1) * stride + kernel_size;
+
+        auto x_flat = ggml_reshape_3d(ctx, x, 1, time, channels);
+        if (stride > 1) {
+            auto zero_unit = ggml_ext_scale(ctx, x_flat, 0.0f);
+            auto zero_tail = zero_unit;
+            for (int i = 1; i < stride - 1; ++i) {
+                zero_tail = ggml_concat(ctx, zero_tail, zero_unit, 0);
+            }
+            x_flat = ggml_concat(ctx, x_flat, zero_tail, 0);
         }
-        return out;
+        x_flat = ggml_reshape_3d(ctx, x_flat, time * stride, 1, channels);
+
+        auto reversed_filter = reverse_1d_filter(ctx, filter);
+        auto out             = ggml_conv_1d(ctx, reversed_filter, x_flat, 1, static_cast<int>(kernel_size - 1), 1);
+        if (out->ne[0] > out_time) {
+            out = ggml_ext_slice(ctx, out, 0, 0, out_time);
+        }
+        GGML_ASSERT(out->ne[0] == out_time);
+        GGML_ASSERT(out->ne[1] == 1);
+        GGML_ASSERT(out->ne[2] == channels);
+
+        out = ggml_ext_scale(ctx, out, static_cast<float>(stride));
+        return ggml_reshape_4d(ctx, out, out_time, channels, 1, 1);
+    }
+
+    static ggml_tensor* upsample_waveform_hann(GGMLRunnerContext* runner_ctx,
+                                               ggml_tensor* waveform,
+                                               ggml_tensor* filter,
+                                               int ratio) {
+        auto ctx = runner_ctx->ggml_ctx;
+        GGML_ASSERT(ctx != nullptr);
+        GGML_ASSERT(waveform != nullptr);
+        GGML_ASSERT(filter != nullptr);
+        GGML_ASSERT(waveform->ne[3] == 1);
+        if (ratio <= 1) {
+            return waveform;
+        }
+
+        const int lowpass_filter_width = 6;
+        const double rolloff           = 0.99;
+        const int width                = static_cast<int>(std::ceil(static_cast<double>(lowpass_filter_width) / rolloff));
+        const int kernel_size          = 2 * width * ratio + 1;
+        const int pad                  = width;
+        const int pad_left             = 2 * width * ratio;
+        const int pad_right            = kernel_size - ratio;
+        const int64_t time             = waveform->ne[0];
+        const int64_t channels         = waveform->ne[1];
+        const int64_t batch            = waveform->ne[2];
+
+        GGML_ASSERT(filter->ne[0] == kernel_size);
+
+        auto x = ggml_reshape_3d(ctx, waveform, time, channels * batch, 1);
+        x      = replicate_pad_1d(runner_ctx, x, pad, pad);
+        x      = depthwise_conv_transpose1d(ctx, x, filter, ratio);
+        x      = ggml_ext_slice(ctx, x, 0, pad_left, x->ne[0] - pad_right);
+        return ggml_reshape_3d(ctx, x, x->ne[0], channels, batch);
+    }
+
+    static ggml_tensor* crop_waveform_samples(ggml_context* ctx,
+                                              ggml_tensor* waveform,
+                                              int64_t target_samples) {
+        GGML_ASSERT(ctx != nullptr);
+        GGML_ASSERT(waveform != nullptr);
+        if (waveform->ne[0] == target_samples) {
+            return waveform;
+        }
+        GGML_ASSERT(waveform->ne[0] > target_samples);
+        return ggml_ext_slice(ctx, waveform, 0, 0, target_samples);
     }
 
     struct PixelNorm2D : public UnaryBlock {
@@ -950,41 +930,66 @@ namespace LTXV {
             }
         }
 
-        ggml_tensor* decode_to_mel(GGMLRunnerContext* ctx,
-                                   ggml_tensor* latent,
-                                   int target_time,
-                                   int target_freq) {
-            auto mean    = params["audio_vae.per_channel_statistics.mean-of-means"];
-            auto stddev  = params["audio_vae.per_channel_statistics.std-of-means"];
-            auto decoder = std::dynamic_pointer_cast<AudioDecoder>(blocks["audio_vae.decoder"]);
-            return decoder->forward(ctx, latent, mean, stddev, target_time, target_freq);
-        }
+        ggml_tensor* decode(GGMLRunnerContext* ctx,
+                            ggml_tensor* latent,
+                            ggml_tensor* bwe_skip_filter) {
+            int target_time = static_cast<int>(latent->ne[1]) * config.latent_downsample_factor() -
+                              (config.latent_downsample_factor() - 1);
+            int target_freq = config.mel_bins;
 
-        ggml_tensor* run_vocoder(GGMLRunnerContext* ctx, ggml_tensor* mel) {
-            auto vocoder = std::dynamic_pointer_cast<Vocoder>(blocks["vocoder.vocoder"]);
-            return vocoder->forward(ctx, mel);
-        }
+            auto decoder  = std::dynamic_pointer_cast<AudioDecoder>(blocks["audio_vae.decoder"]);
+            auto mean     = params["audio_vae.per_channel_statistics.mean-of-means"];
+            auto stddev   = params["audio_vae.per_channel_statistics.std-of-means"];
+            auto mel      = decoder->forward(ctx, latent, mean, stddev, target_time, target_freq);
+            auto vocoder  = std::dynamic_pointer_cast<Vocoder>(blocks["vocoder.vocoder"]);
+            auto waveform = vocoder->forward(ctx, mel);
 
-        ggml_tensor* run_bwe_generator(GGMLRunnerContext* ctx, ggml_tensor* mel) {
-            GGML_ASSERT(config.has_bwe);
-            auto bwe_generator = std::dynamic_pointer_cast<Vocoder>(blocks["vocoder.bwe_generator"]);
-            return bwe_generator->forward(ctx, mel);
-        }
+            if (config.has_bwe) {
+                GGML_ASSERT(bwe_skip_filter != nullptr);
+                const int bwe_ratio    = config.bwe_output_sample_rate / config.bwe_input_sample_rate;
+                const int64_t low_time = waveform->ne[0];
+                const int64_t out_time = low_time * bwe_ratio;
+                int64_t remainder      = low_time % config.bwe_hop_length;
+                auto bwe_waveform      = waveform;
+                if (remainder != 0) {
+                    bwe_waveform = ggml_pad_ext(ctx->ggml_ctx,
+                                                bwe_waveform,
+                                                0,
+                                                static_cast<int>(config.bwe_hop_length - remainder),
+                                                0,
+                                                0,
+                                                0,
+                                                0,
+                                                0,
+                                                0);
+                }
 
-        ggml_tensor* mel_basis_tensor() const {
-            auto iter = params.find("vocoder.mel_stft.mel_basis");
-            return iter == params.end() ? nullptr : iter->second;
-        }
+                auto mel_basis  = params["vocoder.mel_stft.mel_basis"];
+                auto stft_basis = params["vocoder.mel_stft.stft_fn.forward_basis"];
+                GGML_ASSERT(mel_basis != nullptr && stft_basis != nullptr);
+                auto bwe_mel       = compute_log_mel_spectrogram(ctx, bwe_waveform, stft_basis, mel_basis, config.bwe_hop_length);
+                auto bwe_generator = std::dynamic_pointer_cast<Vocoder>(blocks["vocoder.bwe_generator"]);
+                auto residual      = bwe_generator->forward(ctx, bwe_mel);
 
-        ggml_tensor* stft_forward_basis_tensor() const {
-            auto iter = params.find("vocoder.mel_stft.stft_fn.forward_basis");
-            return iter == params.end() ? nullptr : iter->second;
+                auto skip = upsample_waveform_hann(ctx,
+                                                   bwe_waveform,
+                                                   bwe_skip_filter,
+                                                   bwe_ratio);
+                waveform  = ggml_clamp(ctx->ggml_ctx,
+                                       ggml_add(ctx->ggml_ctx, residual, skip),
+                                       -1.0f,
+                                       1.0f);
+                waveform  = crop_waveform_samples(ctx->ggml_ctx, waveform, out_time);
+            }
+
+            return waveform;
         }
     };
 
     struct LTXAudioVAERunner : public GGMLRunner {
         LTXAudioVAEConfig config;
         LTXAudioVAE model;
+        sd::Tensor<float> bwe_skip_filter_tensor;
 
         LTXAudioVAERunner(ggml_backend_t backend,
                           ggml_backend_t params_backend,
@@ -994,6 +999,10 @@ namespace LTXV {
               config(LTXAudioVAEConfig::detect_from_weights(tensor_storage_map)),
               model(config) {
             model.init(params_ctx, tensor_storage_map, prefix);
+            if (config.has_bwe) {
+                const int bwe_ratio    = config.bwe_output_sample_rate / config.bwe_input_sample_rate;
+                bwe_skip_filter_tensor = sd::Tensor<float>::from_vector(build_hann_resample_filter(bwe_ratio));
+            }
         }
 
         void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors, const std::string prefix) {
@@ -1008,77 +1017,22 @@ namespace LTXV {
             return "ltx_audio_vae";
         }
 
-        ggml_cgraph* build_base_graph(const sd::Tensor<float>& latent_tensor) {
-            auto latent     = make_input(latent_tensor);
-            int target_time = static_cast<int>(latent_tensor.shape()[1]) * config.latent_downsample_factor() -
-                              (config.latent_downsample_factor() - 1);
-            int target_freq = config.mel_bins;
-
-            ggml_cgraph* gf = new_graph_custom(655360);
-            auto runner_ctx = GGMLRunner::get_context();
-            auto mel        = model.decode_to_mel(&runner_ctx, latent, target_time, target_freq);
-            auto waveform   = model.run_vocoder(&runner_ctx, mel);
-            ggml_build_forward_expand(gf, waveform);
-            return gf;
-        }
-
-        ggml_cgraph* build_bwe_graph(const sd::Tensor<float>& mel_tensor) {
-            auto mel        = make_input(mel_tensor);
-            ggml_cgraph* gf = new_graph_custom(655360);
-            auto runner_ctx = GGMLRunner::get_context();
-            auto residual   = model.run_bwe_generator(&runner_ctx, mel);
-            ggml_build_forward_expand(gf, residual);
-            return gf;
-        }
-
-        sd::Tensor<float> compute_base_waveform(int n_threads,
-                                                const sd::Tensor<float>& latent_tensor) {
-            auto get_graph = [&]() -> ggml_cgraph* {
-                return build_base_graph(latent_tensor);
-            };
-            return restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false), 4);
-        }
-
-        sd::Tensor<float> compute_bwe_residual(int n_threads,
-                                               const sd::Tensor<float>& mel_tensor) {
-            auto get_graph = [&]() -> ggml_cgraph* {
-                return build_bwe_graph(mel_tensor);
-            };
-            return restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false), 4);
-        }
-
         sd::Tensor<float> decode(int n_threads,
                                  const sd::Tensor<float>& latent_tensor) {
-            auto waveform = compute_base_waveform(n_threads, latent_tensor);
-            if (!config.has_bwe || waveform.empty()) {
-                return waveform;
-            }
-
-            auto waveform_host     = normalize_waveform_for_host(waveform);
-            const int64_t low_time = waveform_host.shape()[0];
-            const int64_t out_time = low_time * config.bwe_output_sample_rate / config.bwe_input_sample_rate;
-            int64_t remainder      = low_time % config.bwe_hop_length;
-            if (remainder != 0) {
-                sd::Tensor<float> padded({low_time + (config.bwe_hop_length - remainder), waveform_host.shape()[1], waveform_host.shape()[2]});
-                sd::ops::slice_assign(&padded, 0, 0, low_time, waveform_host);
-                waveform_host = std::move(padded);
-            }
-
-            auto mel_basis_tensor  = model.mel_basis_tensor();
-            auto stft_basis_tensor = model.stft_forward_basis_tensor();
-            GGML_ASSERT(mel_basis_tensor != nullptr && stft_basis_tensor != nullptr);
-            auto mel_basis     = load_param_tensor_f32(mel_basis_tensor);
-            auto forward_basis = load_param_tensor_f32(stft_basis_tensor);
-            auto bwe_mel       = compute_log_mel_spectrogram(waveform_host, forward_basis, mel_basis, config.bwe_hop_length);
-            auto residual_raw  = compute_bwe_residual(n_threads, bwe_mel);
-            if (residual_raw.empty()) {
-                return waveform;
-            }
-            auto residual = normalize_waveform_for_host(residual_raw);
-            auto skip     = upsample_waveform_hann(waveform_host, config.bwe_output_sample_rate / config.bwe_input_sample_rate);
-            auto combined = sd::ops::clamp(residual + skip, -1.0f, 1.0f);
-            auto cropped  = crop_waveform_samples(combined, out_time);
-            return restore_trailing_singleton_dims(cropped, 4);
+            int64_t t0     = ggml_time_ms();
+            auto get_graph = [&]() -> ggml_cgraph* {
+                auto latent                  = make_input(latent_tensor);
+                ggml_tensor* bwe_skip_filter = config.has_bwe ? make_input(bwe_skip_filter_tensor) : nullptr;
+                ggml_cgraph* gf              = new_graph_custom(655360);
+                auto runner_ctx              = GGMLRunner::get_context();
+                auto waveform                = model.decode(&runner_ctx, latent, bwe_skip_filter);
+                ggml_build_forward_expand(gf, waveform);
+                return gf;
+            };
+            auto result = restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, false), 4);
+            int64_t t1  = ggml_time_ms();
+            LOG_INFO("ltx audio vae decode completed, taking %.2fs", (t1 - t0) * 1.0f / 1000);
+            return result;
         }
 
         void test(const std::string& input_path) {

src/ltx_latent_upscaler.hpp

@@ -6,8 +6,10 @@
 #include <cstdlib>
 #include <map>
 #include <memory>
+#include <set>
 #include <string>
 #include <utility>
+#include <vector>
 
 #include "common_dit.hpp"
 #include "ggml_extend.hpp"
@@ -26,6 +28,9 @@ namespace LTXVUpsampler {
         bool spatial_upsample    = true;
         bool temporal_upsample   = false;
         bool rational_resampler  = false;
+        float spatial_scale      = 2.f;
+        int spatial_up_num       = 2;
+        int spatial_down_den     = 1;
     };
 
     static inline bool has_tensor(const String2TensorStorage& tensor_storage_map,
@@ -33,14 +38,21 @@ namespace LTXVUpsampler {
         return tensor_storage_map.find(name) != tensor_storage_map.end();
     }
 
+    static inline int64_t get_tensor_ne(const String2TensorStorage& tensor_storage_map,
+                                        const std::string& name,
+                                        int axis,
+                                        int64_t fallback) {
+        auto it = tensor_storage_map.find(name);
+        if (it == tensor_storage_map.end() || axis < 0 || axis >= GGML_MAX_DIMS) {
+            return fallback;
+        }
+        return it->second.ne[axis];
+    }
+
     static inline int64_t get_tensor_ne0(const String2TensorStorage& tensor_storage_map,
                                          const std::string& name,
                                          int64_t fallback) {
-        auto it = tensor_storage_map.find(name);
-        if (it == tensor_storage_map.end()) {
-            return fallback;
-        }
-        return it->second.ne[0];
+        return get_tensor_ne(tensor_storage_map, name, 0, fallback);
     }
 
     static inline int count_module_blocks(const String2TensorStorage& tensor_storage_map,
@@ -71,8 +83,32 @@ namespace LTXVUpsampler {
         if (detected_blocks > 0) {
             config.num_blocks_per_stage = detected_blocks;
         }
-        config.spatial_upsample  = has_tensor(tensor_storage_map, "upsampler.0.weight");
-        config.temporal_upsample = has_tensor(tensor_storage_map, "temporal_upsampler.0.weight");
+        config.rational_resampler = has_tensor(tensor_storage_map, "upsampler.conv.weight");
+        config.spatial_upsample   = config.rational_resampler || has_tensor(tensor_storage_map, "upsampler.0.weight");
+        config.temporal_upsample  = has_tensor(tensor_storage_map, "temporal_upsampler.0.weight");
+        if (config.rational_resampler) {
+            int64_t out_channels = get_tensor_ne(tensor_storage_map,
+                                                 "upsampler.conv.weight",
+                                                 3,
+                                                 config.mid_channels * 9);
+            if (config.mid_channels > 0 && out_channels % config.mid_channels == 0) {
+                int64_t ratio = out_channels / config.mid_channels;
+                int num       = static_cast<int>(std::round(std::sqrt(static_cast<double>(ratio))));
+                if (num > 0 && static_cast<int64_t>(num) * num == ratio) {
+                    config.spatial_up_num = num;
+                }
+            }
+            if (config.spatial_up_num == 3) {
+                config.spatial_down_den = 2;
+                config.spatial_scale    = 1.5f;
+            } else if (config.spatial_up_num == 4) {
+                config.spatial_down_den = 1;
+                config.spatial_scale    = 4.f;
+            } else {
+                config.spatial_down_den = 1;
+                config.spatial_scale    = static_cast<float>(config.spatial_up_num);
+            }
+        }
         return config;
     }
 
@@ -160,16 +196,111 @@ namespace LTXVUpsampler {
             : upscale_factor(upscale_factor) {}
 
         ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) override {
-            GGML_ASSERT(upscale_factor == 2);
+            GGML_ASSERT(upscale_factor > 0);
             int64_t h = x->ne[1];
             int64_t w = x->ne[0];
-            // x: [b*f, c*4, h, w] -> [b*f, c, h*2, w*2]
-            x = ggml_ext_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, x, 2, 0, 1, 3));  // [b*f, h, w, c*4]
-            x = ggml_reshape_3d(ctx->ggml_ctx, x, x->ne[0], x->ne[1] * x->ne[2], x->ne[3]);          // [b*f, h*w, c*4]
+            GGML_ASSERT(x->ne[2] % (upscale_factor * upscale_factor) == 0);
+            // x: [b*f, c*p1*p2, h, w] -> [b*f, c, h*p1, w*p2]
+            x = ggml_ext_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, x, 2, 0, 1, 3));  // [b*f, h, w, c*p1*p2]
+            x = ggml_reshape_3d(ctx->ggml_ctx, x, x->ne[0], x->ne[1] * x->ne[2], x->ne[3]);          // [b*f, h*w, c*p1*p2]
             return DiT::unpatchify(ctx->ggml_ctx, x, h, w, upscale_factor, upscale_factor, true);
         }
     };
 
+    class BlurDownsample : public GGMLBlock {
+    protected:
+        int64_t channels;
+        int stride;
+        ggml_tensor* kernel = nullptr;
+        std::vector<float> kernel_data;
+
+        void init_params(ggml_context* ctx,
+                         const String2TensorStorage& tensor_storage_map = {},
+                         const std::string prefix                       = "") override {
+            SD_UNUSED(tensor_storage_map);
+            if (stride == 1) {
+                return;
+            }
+            kernel           = ggml_new_tensor_4d(ctx, GGML_TYPE_F32, 5, 5, 1, channels);
+            std::string name = prefix + "kernel";
+            ggml_set_name(kernel, name.c_str());
+
+            static const float binomial[5] = {1.f, 4.f, 6.f, 4.f, 1.f};
+            kernel_data.resize(static_cast<size_t>(5 * 5 * channels));
+            for (int64_t c = 0; c < channels; ++c) {
+                for (int y = 0; y < 5; ++y) {
+                    for (int x = 0; x < 5; ++x) {
+                        kernel_data[static_cast<size_t>(x + 5 * (y + 5 * c))] =
+                            binomial[y] * binomial[x] / 256.f;
+                    }
+                }
+            }
+        }
+
+    public:
+        BlurDownsample(int64_t channels, int stride)
+            : channels(channels),
+              stride(stride) {
+            GGML_ASSERT(stride >= 1);
+        }
+
+        void load_fixed_tensors() {
+            if (kernel == nullptr || kernel_data.empty()) {
+                return;
+            }
+            ggml_backend_tensor_set(kernel, kernel_data.data(), 0, kernel_data.size() * sizeof(float));
+        }
+
+        ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) {
+            if (stride == 1) {
+                return x;
+            }
+            GGML_ASSERT(kernel != nullptr);
+            GGML_ASSERT(x->ne[2] == channels);
+            if (ctx->conv2d_direct_enabled) {
+                return ggml_conv_2d_dw_direct(ctx->ggml_ctx, kernel, x, stride, stride, 2, 2, 1, 1);
+            }
+            return ggml_conv_2d_dw(ctx->ggml_ctx, kernel, x, stride, stride, 2, 2, 1, 1);
+        }
+    };
+
+    class SpatialRationalResampler : public GGMLBlock {
+    protected:
+        int64_t mid_channels;
+        int num;
+        int den;
+
+    public:
+        SpatialRationalResampler(int64_t mid_channels, int num, int den)
+            : mid_channels(mid_channels),
+              num(num),
+              den(den) {
+            GGML_ASSERT(num >= 1);
+            GGML_ASSERT(den >= 1);
+            blocks["conv"]          = std::shared_ptr<GGMLBlock>(new Conv2d(mid_channels, num * num * mid_channels, {3, 3}, {1, 1}, {1, 1}));
+            blocks["pixel_shuffle"] = std::shared_ptr<GGMLBlock>(new PixelShuffleND(num));
+            blocks["blur_down"]     = std::shared_ptr<GGMLBlock>(new BlurDownsample(mid_channels, den));
+        }
+
+        void load_fixed_tensors() {
+            auto blur_down = std::dynamic_pointer_cast<BlurDownsample>(blocks["blur_down"]);
+            blur_down->load_fixed_tensors();
+        }
+
+        ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) {
+            auto conv          = std::dynamic_pointer_cast<Conv2d>(blocks["conv"]);
+            auto pixel_shuffle = std::dynamic_pointer_cast<PixelShuffleND>(blocks["pixel_shuffle"]);
+            auto blur_down     = std::dynamic_pointer_cast<BlurDownsample>(blocks["blur_down"]);
+
+            // rearrange(x, "b c f h w -> (b f) c h w")
+            x = ggml_ext_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, x, 0, 1, 3, 2));
+            x = conv->forward(ctx, x);
+            x = pixel_shuffle->forward(ctx, x);
+            x = blur_down->forward(ctx, x);
+            return ggml_ext_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, x, 0, 1, 3, 2));
+        }
+    };
+
     class LatentUpsampler : public GGMLBlock {
     public:
         LatentUpsamplerConfig config;
@@ -179,7 +310,6 @@ namespace LTXVUpsampler {
             GGML_ASSERT(this->config.dims == 3);
             GGML_ASSERT(this->config.spatial_upsample);
             GGML_ASSERT(!this->config.temporal_upsample);
-            GGML_ASSERT(!this->config.rational_resampler);
 
             blocks["initial_conv"] = std::shared_ptr<GGMLBlock>(new Conv3d(this->config.in_channels,
                                                                            this->config.mid_channels,
@@ -190,12 +320,18 @@ namespace LTXVUpsampler {
             for (int i = 0; i < this->config.num_blocks_per_stage; ++i) {
                 blocks["res_blocks." + std::to_string(i)] = std::shared_ptr<GGMLBlock>(new ResBlock(this->config.mid_channels, this->config.dims));
             }
-            blocks["upsampler.0"] = std::shared_ptr<GGMLBlock>(new Conv2d(this->config.mid_channels,
-                                                                          4 * this->config.mid_channels,
-                                                                          {3, 3},
-                                                                          {1, 1},
-                                                                          {1, 1}));
-            blocks["upsampler.1"] = std::shared_ptr<GGMLBlock>(new PixelShuffleND(2));
+            if (this->config.rational_resampler) {
+                blocks["upsampler"] = std::shared_ptr<GGMLBlock>(new SpatialRationalResampler(this->config.mid_channels,
+                                                                                              this->config.spatial_up_num,
+                                                                                              this->config.spatial_down_den));
+            } else {
+                blocks["upsampler.0"] = std::shared_ptr<GGMLBlock>(new Conv2d(this->config.mid_channels,
+                                                                              4 * this->config.mid_channels,
+                                                                              {3, 3},
+                                                                              {1, 1},
+                                                                              {1, 1}));
+                blocks["upsampler.1"] = std::shared_ptr<GGMLBlock>(new PixelShuffleND(2));
+            }
             for (int i = 0; i < this->config.num_blocks_per_stage; ++i) {
                 blocks["post_upsample_res_blocks." + std::to_string(i)] = std::shared_ptr<GGMLBlock>(new ResBlock(this->config.mid_channels, this->config.dims));
             }
@@ -207,13 +343,11 @@ namespace LTXVUpsampler {
         }
 
         ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) {
-            // x: [b*c, f, h, w]
-            // return: [b*c, f, h*2, w*2]
-            auto initial_conv  = std::dynamic_pointer_cast<Conv3d>(blocks["initial_conv"]);
-            auto initial_norm  = std::dynamic_pointer_cast<VideoGroupNorm>(blocks["initial_norm"]);
-            auto upsample_conv = std::dynamic_pointer_cast<Conv2d>(blocks["upsampler.0"]);
-            auto pixel_shuffle = std::dynamic_pointer_cast<PixelShuffleND>(blocks["upsampler.1"]);
-            auto final_conv    = std::dynamic_pointer_cast<Conv3d>(blocks["final_conv"]);
+            // x: [b, c, f, h, w]
+            // return: [b, c, f, scaled_h, scaled_w]
+            auto initial_conv = std::dynamic_pointer_cast<Conv3d>(blocks["initial_conv"]);
+            auto initial_norm = std::dynamic_pointer_cast<VideoGroupNorm>(blocks["initial_norm"]);
+            auto final_conv   = std::dynamic_pointer_cast<Conv3d>(blocks["final_conv"]);
 
             x = initial_conv->forward(ctx, x);
             x = initial_norm->forward(ctx, x);
@@ -226,11 +360,19 @@ namespace LTXVUpsampler {
                 sd::ggml_graph_cut::mark_graph_cut(x, "ltx_latent_upsampler.res_blocks." + std::to_string(i), "x");
             }
 
-            // rearrange(x, "b c f h w -> (b f) c h w"),
-            x = ggml_ext_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, x, 0, 1, 3, 2));  // [b*f, c, h, w]
-            x = upsample_conv->forward(ctx, x);                                                      // [b*f, c*4, h, w]
-            x = pixel_shuffle->forward(ctx, x);                                                      // [b*f, c, h*2, w*2]
-            x = ggml_ext_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, x, 0, 1, 3, 2));  // [b*c, f, h, w]
+            if (config.rational_resampler) {
+                auto upsampler = std::dynamic_pointer_cast<SpatialRationalResampler>(blocks["upsampler"]);
+                x              = upsampler->forward(ctx, x);
+            } else {
+                auto upsample_conv = std::dynamic_pointer_cast<Conv2d>(blocks["upsampler.0"]);
+                auto pixel_shuffle = std::dynamic_pointer_cast<PixelShuffleND>(blocks["upsampler.1"]);
+
+                // rearrange(x, "b c f h w -> (b f) c h w"),
+                x = ggml_ext_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, x, 0, 1, 3, 2));  // [b*f, c, h, w]
+                x = upsample_conv->forward(ctx, x);                                                      // [b*f, c*4, h, w]
+                x = pixel_shuffle->forward(ctx, x);                                                      // [b*f, c, h*2, w*2]
+                x = ggml_ext_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, x, 0, 1, 3, 2));  // [b*c, f, h, w]
+            }
             sd::ggml_graph_cut::mark_graph_cut(x, "ltx_latent_upsampler.spatial_up", "x");
 
             for (int i = 0; i < config.num_blocks_per_stage; ++i) {
@@ -243,6 +385,14 @@ namespace LTXVUpsampler {
             sd::ggml_graph_cut::mark_graph_cut(x, "ltx_latent_upsampler.final", "x");
             return x;
         }
+
+        void load_fixed_tensors() {
+            if (!config.rational_resampler) {
+                return;
+            }
+            auto upsampler = std::dynamic_pointer_cast<SpatialRationalResampler>(blocks["upsampler"]);
+            upsampler->load_fixed_tensors();
+        }
     };
 
     struct LatentUpsamplerRunner : public GGMLRunner {
@@ -265,20 +415,23 @@ namespace LTXVUpsampler {
             }
 
             const auto& tensor_storage_map = model_loader.get_tensor_storage_map();
+            bool has_regular_spatial       = has_tensor(tensor_storage_map, "upsampler.0.weight");
+            bool has_rational_spatial      = has_tensor(tensor_storage_map, "upsampler.conv.weight");
             if (!has_tensor(tensor_storage_map, "post_upsample_res_blocks.0.conv2.bias") ||
-                !has_tensor(tensor_storage_map, "upsampler.0.weight")) {
+                (!has_regular_spatial && !has_rational_spatial)) {
                 LOG_ERROR("unsupported LTX latent upsampler weights: expected spatial upsampler tensors");
                 return false;
             }
 
             LatentUpsamplerConfig config = detect_config_from_weights(tensor_storage_map);
             if (config.dims != 3 || !config.spatial_upsample || config.temporal_upsample ||
-                config.rational_resampler) {
-                LOG_ERROR("unsupported LTX latent upsampler config: dims=%d spatial=%d temporal=%d rational=%d",
+                config.spatial_up_num < 1 || config.spatial_down_den < 1) {
+                LOG_ERROR("unsupported LTX latent upsampler config: dims=%d spatial=%d temporal=%d rational=%d scale=%.3f",
                           config.dims,
                           config.spatial_upsample,
                           config.temporal_upsample,
-                          config.rational_resampler);
+                          config.rational_resampler,
+                          config.spatial_scale);
                 return false;
             }
 
@@ -291,15 +444,22 @@ namespace LTXVUpsampler {
 
             std::map<std::string, ggml_tensor*> tensors;
             model->get_param_tensors(tensors);
-            if (!model_loader.load_tensors(tensors, {}, n_threads)) {
+            std::set<std::string> ignore_tensors;
+            if (config.rational_resampler) {
+                ignore_tensors.insert("upsampler.blur_down.kernel");
+            }
+            if (!model_loader.load_tensors(tensors, ignore_tensors, n_threads)) {
                 LOG_ERROR("load LTX latent upsampler tensors failed");
                 return false;
             }
+            model->load_fixed_tensors();
 
-            LOG_INFO("LTX latent upsampler loaded: in_channels=%" PRId64 ", mid_channels=%" PRId64 ", blocks=%d",
+            LOG_INFO("LTX latent upsampler loaded: in_channels=%" PRId64 ", mid_channels=%" PRId64 ", blocks=%d, scale=%.3f, rational=%d",
                      config.in_channels,
                      config.mid_channels,
-                     config.num_blocks_per_stage);
+                     config.num_blocks_per_stage,
+                     config.spatial_scale,
+                     config.rational_resampler);
             return true;
         }
 

src/ltx_vae.hpp

@@ -1,6 +1,7 @@
 #ifndef __SD_LTX_VAE_HPP__
 #define __SD_LTX_VAE_HPP__
 
+#include <algorithm>
 #include <fstream>
 #include <memory>
 #include <string>
@@ -143,16 +144,25 @@ namespace LTXVAE {
                              std::vector<ggml_tensor*>& feat_map,
                              int& feat_idx,
                              int chunk_idx,
-                             bool causal = true) {
+                             bool causal      = true,
+                             int temporal_pad = 0) {
             auto conv     = std::dynamic_pointer_cast<Conv3d>(blocks["conv"]);
             const int pad = causal ? (time_kernel_size - 1) : (time_kernel_size - 1) / 2;
 
             ggml_tensor* prev = (feat_idx < (int)feat_map.size()) ? feat_map[feat_idx] : nullptr;
 
+            GGML_ASSERT(x->ne[2] >= temporal_pad);
+
+            int end_idx   = x->ne[2] - temporal_pad;
+            int start_idx = std::max(end_idx - pad, 0);
+
             // Save a contiguous copy of the last `pad` frames so the large `x`
             // tensor is not kept alive across iterations by a dangling view.
-            if (feat_idx < (int)feat_map.size() && pad > 0 && x->ne[2] >= pad) {
-                auto slice         = ggml_ext_slice(ctx->ggml_ctx, x, 2, x->ne[2] - pad, x->ne[2]);
+            if (feat_idx < (int)feat_map.size() && end_idx - start_idx > 0) {
+                GGML_ASSERT(start_idx >= 0);
+                GGML_ASSERT(end_idx > 0);
+
+                auto slice         = ggml_ext_slice(ctx->ggml_ctx, x, 2, start_idx, end_idx);
                 feat_map[feat_idx] = ggml_cont(ctx->ggml_ctx, slice);
             }
             feat_idx++;
@@ -284,7 +294,8 @@ namespace LTXVAE {
                              bool causal,
                              std::vector<ggml_tensor*>& feat_map,
                              int& feat_idx,
-                             int chunk_idx) {
+                             int chunk_idx,
+                             int temporal_pad = 0) {
             auto norm1 = std::dynamic_pointer_cast<PixelNorm3D>(blocks["norm1"]);
             auto conv1 = std::dynamic_pointer_cast<CausalConv3d>(blocks["conv1"]);
             auto norm2 = std::dynamic_pointer_cast<PixelNorm3D>(blocks["norm2"]);
@@ -311,14 +322,14 @@ namespace LTXVAE {
                 h = apply_scale_shift(ctx->ggml_ctx, h, scale1, shift1);
             }
             h = ggml_silu_inplace(ctx->ggml_ctx, h);
-            h = conv1->forward(ctx, h, feat_map, feat_idx, chunk_idx, causal);
+            h = conv1->forward(ctx, h, feat_map, feat_idx, chunk_idx, causal, temporal_pad);
 
             h = norm2->forward(ctx, h);
             if (timestep_conditioning) {
                 h = apply_scale_shift(ctx->ggml_ctx, h, scale2, shift2);
             }
             h = ggml_silu_inplace(ctx->ggml_ctx, h);
-            h = conv2->forward(ctx, h, feat_map, feat_idx, chunk_idx, causal);
+            h = conv2->forward(ctx, h, feat_map, feat_idx, chunk_idx, causal, temporal_pad);
 
             return ggml_add(ctx->ggml_ctx, h, x);
         }
@@ -367,7 +378,8 @@ namespace LTXVAE {
                              bool causal,
                              std::vector<ggml_tensor*>& feat_map,
                              int& feat_idx,
-                             int chunk_idx) {
+                             int chunk_idx,
+                             int temporal_pad = 0) {
             ggml_tensor* timestep_embed = nullptr;
             if (timestep_conditioning) {
                 GGML_ASSERT(timestep != nullptr);
@@ -376,7 +388,7 @@ namespace LTXVAE {
             }
             for (int i = 0; i < num_layers; i++) {
                 auto resnet = std::dynamic_pointer_cast<ResnetBlock3D>(blocks["res_blocks." + std::to_string(i)]);
-                x           = resnet->forward(ctx, x, timestep_embed, causal, feat_map, feat_idx, chunk_idx);
+                x           = resnet->forward(ctx, x, timestep_embed, causal, feat_map, feat_idx, chunk_idx, temporal_pad);
             }
             return x;
         }
@@ -437,7 +449,8 @@ namespace LTXVAE {
                              bool causal,
                              std::vector<ggml_tensor*>& feat_map,
                              int& feat_idx,
-                             int chunk_idx) {
+                             int chunk_idx,
+                             int temporal_pad = 0) {
             auto conv = std::dynamic_pointer_cast<CausalConv3d>(blocks["conv"]);
 
             bool drop_first = (chunk_idx == 0) && (factor_t > 1);
@@ -453,7 +466,7 @@ namespace LTXVAE {
                 x_in = res;
             }
 
-            x = conv->forward(ctx, x, feat_map, feat_idx, chunk_idx, causal);
+            x = conv->forward(ctx, x, feat_map, feat_idx, chunk_idx, causal, temporal_pad);
             x = depth_to_space_3d(ctx->ggml_ctx, x, get_output_channels(), factor_t, factor_s, drop_first);
             if (residual) {
                 x = ggml_add(ctx->ggml_ctx, x, x_in);
@@ -986,7 +999,8 @@ namespace LTXVAE {
                                          ggml_tensor* timestep,
                                          std::vector<ggml_tensor*>& feat_map,
                                          int& feat_idx,
-                                         int chunk_idx) {
+                                         int chunk_idx,
+                                         int& temporal_pad) {
             auto conv_in       = std::dynamic_pointer_cast<CausalConv3d>(blocks["conv_in"]);
             auto conv_norm_out = std::dynamic_pointer_cast<PixelNorm3D>(blocks["conv_norm_out"]);
             auto conv_out      = std::dynamic_pointer_cast<CausalConv3d>(blocks["conv_out"]);
@@ -998,7 +1012,7 @@ namespace LTXVAE {
             }
 
             // conv_in with feat_map for left temporal context
-            x = conv_in->forward(ctx, x, feat_map, feat_idx, chunk_idx, causal_decoder);
+            x = conv_in->forward(ctx, x, feat_map, feat_idx, chunk_idx, causal_decoder, temporal_pad);
 
             // up_blocks
             int block_idx = 0;
@@ -1006,12 +1020,13 @@ namespace LTXVAE {
                 auto mid_block = std::dynamic_pointer_cast<UNetMidBlock3D>(blocks["up_blocks." + std::to_string(block_idx)]);
                 if (mid_block) {
                     x = mid_block->forward(ctx, x, scaled_timestep, causal_decoder,
-                                           feat_map, feat_idx, chunk_idx);
+                                           feat_map, feat_idx, chunk_idx, temporal_pad);
                 } else {
                     auto upsample = std::dynamic_pointer_cast<DepthToSpaceUpsample>(
                         blocks["up_blocks." + std::to_string(block_idx)]);
                     x = upsample->forward(ctx, x, causal_decoder,
-                                          feat_map, feat_idx, chunk_idx);
+                                          feat_map, feat_idx, chunk_idx, temporal_pad);
+                    temporal_pad *= upsample->factor_t;
                 }
                 block_idx++;
             }
@@ -1028,7 +1043,7 @@ namespace LTXVAE {
                 x                       = apply_scale_shift(ctx->ggml_ctx, x, scale, shift);
             }
             x = ggml_silu_inplace(ctx->ggml_ctx, x);
-            x = conv_out->forward(ctx, x, feat_map, feat_idx, chunk_idx, causal_decoder);
+            x = conv_out->forward(ctx, x, feat_map, feat_idx, chunk_idx, causal_decoder, temporal_pad);
             return x;
         }
     };
@@ -1084,7 +1099,9 @@ namespace LTXVAE {
         // tensors can be freed by GGML before the next iteration starts.
         ggml_tensor* decode_tiled(GGMLRunnerContext* ctx,
                                   ggml_tensor* z,
-                                  ggml_tensor* timestep) {
+                                  ggml_tensor* timestep,
+                                  int temporal_window_size  = 1,
+                                  int temporal_tile_overlap = 0) {
             auto decoder   = std::dynamic_pointer_cast<Decoder>(blocks["decoder"]);
             auto processor = std::dynamic_pointer_cast<PerChannelStatistics>(blocks["per_channel_statistics"]);
             auto latents   = processor->un_normalize(ctx, z);
@@ -1099,18 +1116,69 @@ namespace LTXVAE {
             // 128 slots is generous enough for any supported decoder configuration.
             std::vector<ggml_tensor*> feat_map(128, nullptr);
 
+            // Ensure window size is at least 1
+            int window  = std::max(1, temporal_window_size);
+            int overlap = std::max(0, temporal_tile_overlap);
+
+            if (overlap >= window) {
+                LOG_WARN("temporal_tile_overlap (%d) is greater than or equal to temporal_tile_frames (%d), adjusting values to avoid empty decode windows",
+                         overlap, window);
+                overlap = window - 1;
+            }
+            LOG_DEBUG("Using temporal tiling: temporal_tile_frames = %d, temporal_tile_overlap = %d, total frames = %d, resulting in %d tiles",
+                      window,
+                      overlap,
+                      (int)T,
+                      (T + window - overlap - 1) / (window - overlap));
             ggml_tensor* out = nullptr;
-            for (int i = 0; i < (int)T; i++) {
+            for (int i = 0; i < (int)T - overlap; i += (window - overlap)) {
                 int feat_idx = 0;
-                auto z_i     = ggml_ext_slice(ctx->ggml_ctx, latents, 2, i, i + 1);
-                auto out_i   = decoder->forward_tiled_frame(ctx, z_i, timestep,
-                                                            feat_map, feat_idx, i);
-                out          = (out == nullptr) ? out_i : ggml_concat(ctx->ggml_ctx, out, out_i, 2);
+
+                // Calculate the end index for the current temporal chunk
+                int end_i = std::min((int)T, i + window);
+                if (end_i >= (int)T) {
+                    overlap = 0;  // avoid overlap issues in the last chunk
+                }
+
+                int chunk_overlap = overlap;  // modified by forward_tiled_frame temporal inflation
+
+                auto z_chunk = ggml_ext_slice(ctx->ggml_ctx, latents, 2, i, end_i);
+
+                auto out_chunk = decoder->forward_tiled_frame(ctx, z_chunk, timestep,
+                                                              feat_map, feat_idx, i, chunk_overlap);
+
+                // discard overlap frames if it's not the final chunk
+                if (overlap > 0 && end_i < (int)T) {
+                    out_chunk = ggml_ext_slice(ctx->ggml_ctx, out_chunk, 2, 0, out_chunk->ne[2] - chunk_overlap);
+                }
+
+                out = (out == nullptr) ? out_chunk : ggml_concat(ctx->ggml_ctx, out, out_chunk, 2);
             }
 
             return WAN::WanVAE::unpatchify(ctx->ggml_ctx, out, patch_size, 1);
         }
 
+        ggml_tensor* decode_tiled_chunk(GGMLRunnerContext* ctx,
+                                        ggml_tensor* z,
+                                        ggml_tensor* timestep,
+                                        std::vector<ggml_tensor*>& feat_map,
+                                        int chunk_idx,
+                                        int temporal_tile_overlap,
+                                        int& feat_idx) {
+            auto decoder   = std::dynamic_pointer_cast<Decoder>(blocks["decoder"]);
+            auto processor = std::dynamic_pointer_cast<PerChannelStatistics>(blocks["per_channel_statistics"]);
+            auto latents   = processor->un_normalize(ctx, z);
+
+            feat_idx          = 0;
+            int chunk_overlap = temporal_tile_overlap;  // modified by forward_tiled_frame temporal inflation
+            auto out_chunk    = decoder->forward_tiled_frame(ctx, latents, timestep,
+                                                             feat_map, feat_idx, chunk_idx, chunk_overlap);
+            if (chunk_overlap > 0) {
+                out_chunk = ggml_ext_slice(ctx->ggml_ctx, out_chunk, 2, 0, out_chunk->ne[2] - chunk_overlap);
+            }
+            return WAN::WanVAE::unpatchify(ctx->ggml_ctx, out_chunk, patch_size, 1);
+        }
+
         ggml_tensor* encode(GGMLRunnerContext* ctx,
                             ggml_tensor* x) {
             GGML_ASSERT(!decode_only);
@@ -1140,8 +1208,13 @@ namespace LTXVAE {
 }  // namespace LTXVAE
 
 struct LTXVideoVAE : public VAE {
+    static constexpr int DEFAULT_TEMPORAL_TILE_FRAMES  = 4;
+    static constexpr int DEFAULT_TEMPORAL_TILE_OVERLAP = 1;
+
     bool decode_only;
     bool temporal_tiling_enabled = false;
+    int temporal_tile_frames     = DEFAULT_TEMPORAL_TILE_FRAMES;
+    int temporal_tile_overlap    = DEFAULT_TEMPORAL_TILE_OVERLAP;
     int ltx_vae_version;
     bool timestep_conditioning;
     int patch_size;
@@ -1178,10 +1251,64 @@ struct LTXVideoVAE : public VAE {
         temporal_tiling_enabled = enabled;
     }
 
+    void set_tiling_params(const sd_tiling_params_t& params) override {
+        temporal_tiling_enabled = params.temporal_tiling;
+        temporal_tile_frames    = DEFAULT_TEMPORAL_TILE_FRAMES;
+        temporal_tile_overlap   = DEFAULT_TEMPORAL_TILE_OVERLAP;
+
+        for (const auto& [key, value] : parse_key_value_args(params.extra_tiling_args, "LTX VAE extra tiling arg")) {
+            int parsed = 0;
+            if (!parse_strict_int(value, parsed)) {
+                LOG_WARN("ignoring invalid LTX VAE extra tiling arg '%s=%s'", key.c_str(), value.c_str());
+            } else if (key == "temporal_tile_frames") {
+                temporal_tile_frames = std::max(1, parsed);
+            } else if (key == "temporal_tile_overlap") {
+                temporal_tile_overlap = std::max(0, parsed);
+            } else {
+                LOG_WARN("ignoring unknown LTX VAE extra tiling arg '%s'", key.c_str());
+            }
+        }
+    }
+
     void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors, const std::string prefix) override {
         vae.get_param_tensors(tensors, prefix);
     }
 
+    struct TemporalTilePlan {
+        int frames    = 1;
+        int overlap   = 0;
+        int stride    = 1;
+        int num_tiles = 1;
+    };
+
+    TemporalTilePlan resolve_temporal_tile_plan(int64_t total_frames) const {
+        TemporalTilePlan plan;
+        plan.frames  = std::max(1, temporal_tile_frames);
+        plan.overlap = std::max(0, temporal_tile_overlap);
+
+        if (plan.overlap >= plan.frames) {
+            LOG_WARN("temporal_tile_overlap (%d) is greater than or equal to temporal_tile_frames (%d), adjusting values to avoid empty decode windows",
+                     plan.overlap,
+                     plan.frames);
+            plan.overlap = plan.frames - 1;
+        }
+        if (total_frames > 1 && plan.overlap >= total_frames) {
+            LOG_WARN("temporal_tile_overlap (%d) is greater than or equal to total latent frames (%lld), adjusting values to decode at least one tile",
+                     plan.overlap,
+                     (long long)total_frames);
+            plan.overlap = static_cast<int>(total_frames - 1);
+        }
+
+        plan.stride          = std::max(1, plan.frames - plan.overlap);
+        int64_t tiled_frames = std::max<int64_t>(1, total_frames - plan.overlap);
+        plan.num_tiles       = total_frames > 0 ? static_cast<int>((tiled_frames + plan.stride - 1) / plan.stride) : 0;
+        return plan;
+    }
+
+    std::string temporal_feat_cache_name(size_t feat_idx) const {
+        return "ltx_vae_temporal_feat:" + std::to_string(feat_idx);
+    }
+
     ggml_cgraph* build_graph(const sd::Tensor<float>& z_tensor, bool decode_graph) {
         ggml_cgraph* gf       = new_graph_custom(20480);
         ggml_tensor* z        = make_input(z_tensor);
@@ -1192,18 +1319,97 @@ struct LTXVideoVAE : public VAE {
 
         auto runner_ctx = get_context();
         ggml_tensor* out;
-        bool use_tiled = decode_graph && temporal_tiling_enabled &&
-                         z_tensor.dim() == 5 && z_tensor.shape()[2] > 1;
-        if (use_tiled) {
-            out = vae.decode_tiled(&runner_ctx, z, timestep);
-        } else {
-            out = decode_graph ? vae.decode(&runner_ctx, z, timestep) : vae.encode(&runner_ctx, z);
-        }
+        out = decode_graph ? vae.decode(&runner_ctx, z, timestep) : vae.encode(&runner_ctx, z);
         ggml_build_forward_expand(gf, out);
 
         return gf;
     }
 
+    ggml_cgraph* build_temporal_tile_graph(const sd::Tensor<float>& z_chunk_tensor,
+                                           int chunk_idx,
+                                           int chunk_overlap) {
+        ggml_cgraph* gf       = new_graph_custom(20480);
+        ggml_tensor* z        = make_input(z_chunk_tensor);
+        ggml_tensor* timestep = nullptr;
+        if (timestep_conditioning) {
+            timestep = make_input(decode_timestep_tensor);
+        }
+
+        std::vector<ggml_tensor*> feat_map(128, nullptr);
+        for (size_t feat_idx = 0; feat_idx < feat_map.size(); ++feat_idx) {
+            feat_map[feat_idx] = get_cache_tensor_by_name(temporal_feat_cache_name(feat_idx));
+        }
+
+        auto runner_ctx  = get_context();
+        int feat_count   = 0;
+        ggml_tensor* out = vae.decode_tiled_chunk(&runner_ctx,
+                                                  z,
+                                                  timestep,
+                                                  feat_map,
+                                                  chunk_idx,
+                                                  chunk_overlap,
+                                                  feat_count);
+
+        for (int feat_idx = 0; feat_idx < feat_count && feat_idx < static_cast<int>(feat_map.size()); ++feat_idx) {
+            ggml_tensor* feat_cache = feat_map[static_cast<size_t>(feat_idx)];
+            if (feat_cache != nullptr) {
+                cache(temporal_feat_cache_name(static_cast<size_t>(feat_idx)), feat_cache);
+                ggml_build_forward_expand(gf, feat_cache);
+            }
+        }
+
+        ggml_build_forward_expand(gf, out);
+        return gf;
+    }
+
+    sd::Tensor<float> decode_temporal_tiled_streaming(const int n_threads,
+                                                      const sd::Tensor<float>& input,
+                                                      size_t expected_dim) {
+        const int64_t total_frames = input.shape()[2];
+        TemporalTilePlan plan      = resolve_temporal_tile_plan(total_frames);
+
+        LOG_DEBUG("Using streaming temporal tiling: temporal_tile_frames=%d, temporal_tile_overlap=%d, total latent frames=%lld, resulting in %d tiles",
+                  plan.frames,
+                  plan.overlap,
+                  (long long)total_frames,
+                  plan.num_tiles);
+
+        free_cache_ctx_and_buffer();
+        cache_tensor_map.clear();
+
+        sd::Tensor<float> output;
+        for (int64_t start = 0; start < total_frames - plan.overlap; start += plan.stride) {
+            const int64_t end       = std::min<int64_t>(total_frames, start + plan.frames);
+            const int chunk_overlap = end < total_frames ? plan.overlap : 0;
+            auto z_chunk            = sd::ops::slice(input, 2, start, end);
+
+            LOG_DEBUG("LTX VAE temporal tile %lld/%d: latent frames [%lld, %lld), overlap=%d",
+                      (long long)(start / plan.stride + 1),
+                      plan.num_tiles,
+                      (long long)start,
+                      (long long)end,
+                      chunk_overlap);
+
+            auto get_graph = [&]() -> ggml_cgraph* {
+                return build_temporal_tile_graph(z_chunk,
+                                                 static_cast<int>(start),
+                                                 chunk_overlap);
+            };
+            auto chunk = restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, true),
+                                                         expected_dim);
+            if (chunk.empty()) {
+                free_cache_ctx_and_buffer();
+                cache_tensor_map.clear();
+                return {};
+            }
+            output = output.empty() ? std::move(chunk) : sd::ops::concat(output, chunk, 2);
+        }
+
+        free_cache_ctx_and_buffer();
+        cache_tensor_map.clear();
+        return output;
+    }
+
     ggml_cgraph* build_latent_statistics_graph(const sd::Tensor<float>& z_tensor, bool normalize) {
         ggml_cgraph* gf = new_graph_custom(1024);
         ggml_tensor* z  = make_input(z_tensor);
@@ -1239,6 +1445,9 @@ struct LTXVideoVAE : public VAE {
                 input = sd::ops::slice(input, 2, 0, cropped_t);
             }
         }
+        if (decode_graph && temporal_tiling_enabled && input.dim() == 5 && input.shape()[2] > 1) {
+            return decode_temporal_tiled_streaming(n_threads, input, expected_dim);
+        }
         auto get_graph = [&]() -> ggml_cgraph* {
             return build_graph(input, decode_graph);
         };

src/stable-diffusion.cpp

@@ -151,7 +151,7 @@ public:
     bool apply_lora_immediately = false;
 
     std::string taesd_path;
-    sd_tiling_params_t vae_tiling_params = {false, false, 0, 0, 0.5f, 0, 0};
+    sd_tiling_params_t vae_tiling_params = {false, false, 0, 0, 0.5f, 0, 0, nullptr};
     bool offload_params_to_cpu           = false;
     float max_vram                       = 0.f;
     bool use_pmid                        = false;
@@ -759,7 +759,7 @@ public:
                     LOG_INFO("using TAE for preview");
                     preview_vae = create_tae();
                     preview_vae->set_max_graph_vram_bytes(max_graph_vram_bytes);
-                    get_param_tensors_p(first_stage_model, vae_mmap, "vae");
+                    get_param_tensors_p(preview_vae, vae_mmap, "tae");
                 }
             }
 
@@ -1610,23 +1610,18 @@ public:
                        std::function<void(int, int, sd_image_t*, bool, void*)> step_callback,
                        void* step_callback_data,
                        bool is_noisy) {
+        bool is_video = preview_latent_tensor_is_video(latents);
+        uint32_t dim  = is_video ? static_cast<uint32_t>(latents.shape()[3]) : static_cast<uint32_t>(latents.shape()[2]);
+        int channels  = get_latent_channel();
+        auto _latents = channels != dim ? is_video ? sd::ops::slice(latents, 3, 0, channels)
+                                                   : sd::ops::slice(latents, 2, 0, channels)
+                                        : latents;
         if (preview_mode == PREVIEW_PROJ) {
-            sd::Tensor<float> _latents       = latents;
             int patch_sz                     = 1;
             const float(*latent_rgb_proj)[3] = nullptr;
             float* latent_rgb_bias           = nullptr;
-            bool is_video                    = preview_latent_tensor_is_video(latents);
-            uint32_t dim                     = is_video ? static_cast<uint32_t>(latents.shape()[3]) : static_cast<uint32_t>(latents.shape()[2]);
-            if (version == VERSION_LTXAV) {
-                if (is_video) {
-                    _latents = sd::ops::slice(_latents, 3, 0, 128);
-                } else {
-                    _latents = sd::ops::slice(_latents, 2, 0, 128);
-                }
-                dim = 128;
-            }
 
-            if (dim == 128) {
+            if (channels == 128) {
                 if (sd_version_uses_flux2_vae(version)) {
                     latent_rgb_proj = flux2_latent_rgb_proj;
                     latent_rgb_bias = flux2_latent_rgb_bias;
@@ -1638,7 +1633,7 @@ public:
                     LOG_WARN("No latent to RGB projection known for this model");
                     return;
                 }
-            } else if (dim == 48) {
+            } else if (channels == 48) {
                 if (sd_version_is_wan(version)) {
                     latent_rgb_proj = wan_22_latent_rgb_proj;
                     latent_rgb_bias = wan_22_latent_rgb_bias;
@@ -1646,7 +1641,7 @@ public:
                     LOG_WARN("No latent to RGB projection known for this model");
                     return;
                 }
-            } else if (dim == 16) {
+            } else if (channels == 16) {
                 if (sd_version_is_sd3(version)) {
                     latent_rgb_proj = sd3_latent_rgb_proj;
                     latent_rgb_bias = sd3_latent_rgb_bias;
@@ -1660,7 +1655,7 @@ public:
                     LOG_WARN("No latent to RGB projection known for this model");
                     return;
                 }
-            } else if (dim == 4) {
+            } else if (channels == 4) {
                 if (sd_version_is_sdxl(version)) {
                     latent_rgb_proj = sdxl_latent_rgb_proj;
                     latent_rgb_bias = sdxl_latent_rgb_bias;
@@ -1671,8 +1666,8 @@ public:
                     LOG_WARN("No latent to RGB projection known for this model");
                     return;
                 }
-            } else if (dim != 3) {
-                LOG_WARN("No latent to RGB projection known for this model");
+            } else if (channels != 3) {
+                LOG_WARN("No latent to RGB projection known for this model (dim = %d)", dim);
                 return;
             }
 
@@ -1697,14 +1692,13 @@ public:
         if (preview_mode == PREVIEW_VAE || preview_mode == PREVIEW_TAE) {
             sd::Tensor<float> vae_latents;
             sd::Tensor<float> decoded;
-            bool is_video = preview_latent_tensor_is_video(latents);
             if (preview_vae) {
                 preview_vae->set_temporal_tiling_enabled(vae_tiling_params.temporal_tiling);
-                vae_latents = preview_vae->diffusion_to_vae_latents(latents);
+                vae_latents = preview_vae->diffusion_to_vae_latents(_latents);
                 decoded     = preview_vae->decode(n_threads, vae_latents, vae_tiling_params, is_video, circular_x, circular_y, true);
             } else {
                 first_stage_model->set_temporal_tiling_enabled(vae_tiling_params.temporal_tiling);
-                vae_latents = first_stage_model->diffusion_to_vae_latents(latents);
+                vae_latents = first_stage_model->diffusion_to_vae_latents(_latents);
                 decoded     = first_stage_model->decode(n_threads, vae_latents, vae_tiling_params, is_video, circular_x, circular_y, true);
             }
             if (decoded.empty()) {
@@ -2679,7 +2673,7 @@ void sd_img_gen_params_init(sd_img_gen_params_t* sd_img_gen_params) {
     sd_img_gen_params->batch_count       = 1;
     sd_img_gen_params->control_strength  = 0.9f;
     sd_img_gen_params->pm_params         = {nullptr, 0, nullptr, 20.f};
-    sd_img_gen_params->vae_tiling_params = {false, false, 0, 0, 0.5f, 0.0f, 0.0f};
+    sd_img_gen_params->vae_tiling_params = {false, false, 0, 0, 0.5f, 0.0f, 0.0f, nullptr};
     sd_cache_params_init(&sd_img_gen_params->cache);
     sd_hires_params_init(&sd_img_gen_params->hires);
 }
@@ -2708,7 +2702,7 @@ char* sd_img_gen_params_to_str(const sd_img_gen_params_t* sd_img_gen_params) {
              "increase_ref_index: %s\n"
              "control_strength: %.2f\n"
              "photo maker: {style_strength = %.2f, id_images_count = %d, id_embed_path = %s}\n"
-             "VAE tiling: %s (temporal=%s)\n"
+             "VAE tiling: %s (temporal=%s, extra_tiling_args=%s)\n"
              "hires: {enabled=%s, upscaler=%s, model_path=%s, scale=%.2f, target=%dx%d, steps=%d, denoising_strength=%.2f}\n",
              SAFE_STR(sd_img_gen_params->prompt),
              SAFE_STR(sd_img_gen_params->negative_prompt),
@@ -2728,6 +2722,7 @@ char* sd_img_gen_params_to_str(const sd_img_gen_params_t* sd_img_gen_params) {
              SAFE_STR(sd_img_gen_params->pm_params.id_embed_path),
              BOOL_STR(sd_img_gen_params->vae_tiling_params.enabled),
              BOOL_STR(sd_img_gen_params->vae_tiling_params.temporal_tiling),
+             SAFE_STR(sd_img_gen_params->vae_tiling_params.extra_tiling_args),
              BOOL_STR(sd_img_gen_params->hires.enabled),
              sd_hires_upscaler_name(sd_img_gen_params->hires.upscaler),
              SAFE_STR(sd_img_gen_params->hires.model_path),
@@ -2765,7 +2760,7 @@ void sd_vid_gen_params_init(sd_vid_gen_params_t* sd_vid_gen_params) {
     sd_vid_gen_params->fps                                   = 16;
     sd_vid_gen_params->moe_boundary                          = 0.875f;
     sd_vid_gen_params->vace_strength                         = 1.f;
-    sd_vid_gen_params->vae_tiling_params                     = {false, false, 0, 0, 0.5f, 0.0f, 0.0f};
+    sd_vid_gen_params->vae_tiling_params                     = {false, false, 0, 0, 0.5f, 0.0f, 0.0f, nullptr};
     sd_vid_gen_params->hires.enabled                         = false;
     sd_vid_gen_params->hires.upscaler                        = SD_HIRES_UPSCALER_LATENT;
     sd_vid_gen_params->hires.scale                           = 2.f;
@@ -4797,7 +4792,7 @@ static sd::Tensor<float> upscale_ltx_spatial_video_latent(sd_ctx_t* sd_ctx,
         audio_latent = unpack_ltxav_audio_latent(packed_latent, audio_length, latent_channels);
     }
 
-    LOG_INFO("LTX latent spatial upscale: latent %dx%dx%dx%d -> x2",
+    LOG_INFO("LTX latent spatial upscale: latent %dx%dx%dx%d -> model output",
              (int)video_latent.shape()[0],
              (int)video_latent.shape()[1],
              (int)video_latent.shape()[2],
@@ -5218,14 +5213,24 @@ SD_API bool generate_video(sd_ctx_t* sd_ctx,
         sd_ctx->sd->diffusion_model->free_params_buffer();
     }
 
+    int64_t latent_end = ggml_time_ms();
+    LOG_INFO("generating latent video completed, taking %.2fs", (latent_end - latent_start) * 1.0f / 1000);
+
     sd_audio_t* generated_audio = nullptr;
     if (sd_version_is_ltxav(sd_ctx->sd->version) &&
         latents.audio_length > 0 &&
         sd_ctx->sd->audio_vae_model != nullptr) {
+        int64_t audio_latent_decode_start = ggml_time_ms();
+
         auto audio_latent = unpack_ltxav_audio_latent(final_latent,
                                                       latents.audio_length,
                                                       sd_ctx->sd->get_latent_channel());
         if (!audio_latent.empty()) {
+            LOG_DEBUG("decode audio latent %dx%dx%dx%d",
+                      (int)audio_latent.shape()[0],
+                      (int)audio_latent.shape()[1],
+                      (int)audio_latent.shape()[2],
+                      (int)audio_latent.shape()[3]);
             auto waveform = sd_ctx->sd->decode_ltx_audio_latent(audio_latent);
             if (!waveform.empty()) {
                 generated_audio = waveform_to_sd_audio(sd_ctx->sd, waveform);
@@ -5233,6 +5238,8 @@ SD_API bool generate_video(sd_ctx_t* sd_ctx,
                 LOG_WARN("LTX audio latent decode failed; continuing with silent video output");
             }
         }
+        int64_t audio_latent_decode_end = ggml_time_ms();
+        LOG_INFO("decoding audio latent completed, taking %.2fs", (audio_latent_decode_end - audio_latent_decode_start) * 1.0f / 1000);
     }
 
     if (latents.video_conditioning_frame_count > 0) {
@@ -5245,9 +5252,6 @@ SD_API bool generate_video(sd_ctx_t* sd_ctx,
         final_latent = sd::ops::slice(final_latent, 2, latents.ref_image_num, final_latent.shape()[2]);
     }
 
-    int64_t latent_end = ggml_time_ms();
-    LOG_INFO("generating latent video completed, taking %.2fs", (latent_end - latent_start) * 1.0f / 1000);
-
     auto result = decode_video_outputs(sd_ctx, latent_upscale_enabled ? hires_request : request, final_latent, num_frames_out);
     if (result == nullptr) {
         free_sd_audio(generated_audio);

src/tae.hpp

@@ -259,10 +259,54 @@ public:
     }
 };
 
-ggml_tensor* patchify(ggml_context* ctx,
-                      ggml_tensor* x,
-                      int64_t patch_size,
-                      int64_t b = 1) {
+class WideMemBlock : public GGMLBlock {
+    bool has_skip_conv = false;
+
+public:
+    WideMemBlock(int channels, int out_channels)
+        : has_skip_conv(channels != out_channels) {
+        int groups       = std::max(1, out_channels / 64);
+        blocks["conv.0"] = std::shared_ptr<GGMLBlock>(new Conv2d(channels * 2, out_channels, {1, 1}, {1, 1}));
+        blocks["conv.2"] = std::shared_ptr<GGMLBlock>(new Conv2d_grouped(out_channels, out_channels, groups, {3, 3}, {1, 1}, {1, 1}));
+        blocks["conv.4"] = std::shared_ptr<GGMLBlock>(new Conv2d(out_channels, out_channels, {1, 1}, {1, 1}));
+        blocks["conv.6"] = std::shared_ptr<GGMLBlock>(new Conv2d_grouped(out_channels, out_channels, groups, {3, 3}, {1, 1}, {1, 1}));
+        if (has_skip_conv) {
+            blocks["skip"] = std::shared_ptr<GGMLBlock>(new Conv2d(channels, out_channels, {1, 1}, {1, 1}, {0, 0}, {1, 1}, false));
+        }
+    }
+
+    ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x, ggml_tensor* past) {
+        // x: [n, channels, h, w]
+        auto conv0 = std::dynamic_pointer_cast<Conv2d>(blocks["conv.0"]);
+        auto conv1 = std::dynamic_pointer_cast<Conv2d_grouped>(blocks["conv.2"]);
+        auto conv2 = std::dynamic_pointer_cast<Conv2d>(blocks["conv.4"]);
+        auto conv3 = std::dynamic_pointer_cast<Conv2d_grouped>(blocks["conv.6"]);
+
+        auto h = ggml_concat(ctx->ggml_ctx, x, past, 2);
+        h      = conv0->forward(ctx, h);
+        h      = ggml_relu_inplace(ctx->ggml_ctx, h);
+        h      = conv1->forward(ctx, h);
+        h      = ggml_relu_inplace(ctx->ggml_ctx, h);
+        h      = conv2->forward(ctx, h);
+        h      = ggml_relu_inplace(ctx->ggml_ctx, h);
+        h      = conv3->forward(ctx, h);
+
+        auto skip = x;
+        if (has_skip_conv) {
+            auto skip_conv = std::dynamic_pointer_cast<Conv2d>(blocks["skip"]);
+            skip           = skip_conv->forward(ctx, x);
+        }
+        h = ggml_add_inplace(ctx->ggml_ctx, h, skip);
+        h = ggml_relu_inplace(ctx->ggml_ctx, h);
+        return h;
+    }
+};
+
+ggml_tensor*
+patchify(ggml_context* ctx,
+         ggml_tensor* x,
+         int64_t patch_size,
+         int64_t b = 1) {
     // x: [f, b*c, h*q, w*r]
     // return: [f, b*c*r*q, h, w]
     if (patch_size == 1) {
@@ -325,7 +369,6 @@ public:
     int t_downscale = 1;
     TinyVideoEncoder(int z_channels = 4, int patch_size = 1, std::vector<bool> time_downscale = {true, true, false})
         : z_channels(z_channels), patch_size(patch_size) {
-        //         self.t_downscale = 2**sum(t.stride == 2 for t in self.encoder if isinstance(t, TPool))
         t_downscale = 1;
         for (bool downscale : time_downscale) {
             if (downscale) {
@@ -384,11 +427,18 @@ class TinyVideoDecoder : public UnaryBlock {
     int channels[num_layers + 1] = {256, 128, 64, 64};
     int patch_size               = 1;
     int t_upscale                = 1;
+    bool is_wide                 = false;
 
 public:
-    TinyVideoDecoder(int z_channels = 4, int patch_size = 1, std::vector<bool> time_upscale = {false, true, true})
-        : z_channels(z_channels), patch_size(patch_size) {
+    TinyVideoDecoder(int z_channels = 4, int patch_size = 1, std::vector<bool> time_upscale = {false, true, true}, bool is_wide = false)
+        : z_channels(z_channels), patch_size(patch_size), is_wide(is_wide) {
         t_upscale = 1;
+        if (is_wide) {
+            channels[0] = 1024;
+            channels[1] = 512;
+            channels[2] = 256;
+        }
+
         for (bool upscale : time_upscale) {
             if (upscale) {
                 t_upscale *= 2;
@@ -400,7 +450,11 @@ public:
         for (int i = 0; i < num_layers; i++) {
             int stride = time_upscale[i] ? 2 : 1;
             for (int j = 0; j < num_blocks; j++) {
-                blocks[std::to_string(index++)] = std::shared_ptr<GGMLBlock>(new MemBlock(channels[i], channels[i]));
+                if (is_wide) {
+                    blocks[std::to_string(index++)] = std::shared_ptr<GGMLBlock>(new WideMemBlock(channels[i], channels[i]));
+                } else {
+                    blocks[std::to_string(index++)] = std::shared_ptr<GGMLBlock>(new MemBlock(channels[i], channels[i]));
+                }
             }
             index++;  // nn.Upsample()
             blocks[std::to_string(index++)] = std::shared_ptr<GGMLBlock>(new TGrow(channels[i], stride));
@@ -425,10 +479,15 @@ public:
         int index = 3;
         for (int i = 0; i < num_layers; i++) {
             for (int j = 0; j < num_blocks; j++) {
-                auto block = std::dynamic_pointer_cast<MemBlock>(blocks[std::to_string(index++)]);
-                auto mem   = ggml_pad_ext(ctx->ggml_ctx, h, 0, 0, 0, 0, 0, 0, 1, 0);
-                mem        = ggml_view_4d(ctx->ggml_ctx, mem, h->ne[0], h->ne[1], h->ne[2], h->ne[3], h->nb[1], h->nb[2], h->nb[3], 0);
-                h          = block->forward(ctx, h, mem);
+                auto mem = ggml_pad_ext(ctx->ggml_ctx, h, 0, 0, 0, 0, 0, 0, 1, 0);
+                mem      = ggml_view_4d(ctx->ggml_ctx, mem, h->ne[0], h->ne[1], h->ne[2], h->ne[3], h->nb[1], h->nb[2], h->nb[3], 0);
+                if (is_wide) {
+                    auto block = std::dynamic_pointer_cast<WideMemBlock>(blocks[std::to_string(index++)]);
+                    h          = block->forward(ctx, h, mem);
+                } else {
+                    auto block = std::dynamic_pointer_cast<MemBlock>(blocks[std::to_string(index++)]);
+                    h          = block->forward(ctx, h, mem);
+                }
             }
             // upsample
             index++;
@@ -455,6 +514,7 @@ class TAEHV : public GGMLBlock {
 protected:
     bool decode_only;
     SDVersion version;
+    bool is_wide;
 
 public:
     int z_channels                   = 16;
@@ -462,8 +522,8 @@ public:
     std::vector<bool> time_upscale   = {false, true, true};
 
 public:
-    TAEHV(bool decode_only = true, SDVersion version = VERSION_WAN2)
-        : decode_only(decode_only), version(version) {
+    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) {
             z_channels = 48;
@@ -474,7 +534,7 @@ public:
             time_downscale = {true, true, true};
             time_upscale   = {true, true, true};
         }
-        blocks["decoder"] = std::shared_ptr<GGMLBlock>(new TinyVideoDecoder(z_channels, patch, time_upscale));
+        blocks["decoder"] = std::shared_ptr<GGMLBlock>(new TinyVideoDecoder(z_channels, patch, time_upscale, is_wide));
         if (!decode_only) {
             blocks["encoder"] = std::shared_ptr<GGMLBlock>(new TinyVideoEncoder(z_channels, patch, time_downscale));
         }
@@ -623,6 +683,7 @@ struct TinyImageAutoEncoder : public VAE {
 struct TinyVideoAutoEncoder : public VAE {
     TAEHV taehv;
     bool decode_only = false;
+    bool is_wide     = false;
 
     TinyVideoAutoEncoder(ggml_backend_t backend,
                          ggml_backend_t params_backend,
@@ -631,8 +692,14 @@ struct TinyVideoAutoEncoder : public VAE {
                          bool decoder_only = true,
                          SDVersion version = VERSION_WAN2)
         : decode_only(decoder_only),
-          taehv(decoder_only, version),
           VAE(version, backend, params_backend) {
+        for (auto tensor_storage : tensor_storage_map) {
+            if (tensor_storage.first.find(prefix + ".3.conv.6.weight") != std::string::npos) {
+                is_wide = true;
+                break;
+            }
+        }
+        taehv       = TAEHV(decoder_only, version, is_wide);
         scale_input = false;
         taehv.init(params_ctx, tensor_storage_map, prefix);
     }
@@ -663,7 +730,8 @@ struct TinyVideoAutoEncoder : public VAE {
     }
 
     ggml_cgraph* build_graph(const sd::Tensor<float>& z_tensor, bool decode_graph) {
-        ggml_cgraph* gf  = ggml_new_graph(compute_ctx);
+        ggml_cgraph* gf  = decode_graph && is_wide ? ggml_new_graph_custom(compute_ctx, 4096, false)
+                                                   : ggml_new_graph(compute_ctx);
         ggml_tensor* z   = make_input(z_tensor);
         auto runner_ctx  = get_context();
         ggml_tensor* out = decode_graph ? taehv.decode(&runner_ctx, z) : taehv.encode(&runner_ctx, z);

src/util.cpp

@@ -1,8 +1,10 @@
 #include "util.h"
 #include <algorithm>
+#include <cctype>
 #include <cmath>
 #include <codecvt>
 #include <cstdarg>
+#include <exception>
 #include <fstream>
 #include <locale>
 #include <regex>
@@ -406,6 +408,88 @@ std::vector<std::string> split_string(const std::string& str, char delimiter) {
     return result;
 }
 
+KeyValueArgs parse_key_value_args(const char* args, const char* context) {
+    KeyValueArgs pairs;
+
+    if (args == nullptr || args[0] == '\0') {
+        return pairs;
+    }
+
+    std::string raw(args);
+    size_t start = 0;
+    for (size_t pos = 0; pos <= raw.size(); ++pos) {
+        if (pos != raw.size() && raw[pos] != ',' && raw[pos] != ';') {
+            continue;
+        }
+
+        std::string token = trim(raw.substr(start, pos - start));
+        if (!token.empty()) {
+            size_t eq = token.find('=');
+            if (eq == std::string::npos) {
+                const char* log_context = context ? context : "key=value arg";
+                LOG_WARN("ignoring malformed %s '%s'", log_context, token.c_str());
+            } else {
+                std::string key   = trim(token.substr(0, eq));
+                std::string value = trim(token.substr(eq + 1));
+                pairs.emplace_back(std::move(key), std::move(value));
+            }
+        }
+
+        start = pos + 1;
+    }
+
+    return pairs;
+}
+
+KeyValueArgs parse_key_value_args(const std::string& args, const char* context) {
+    return parse_key_value_args(args.c_str(), context);
+}
+
+bool parse_strict_float(const std::string& text, float& value) {
+    try {
+        size_t consumed = 0;
+        float parsed    = std::stof(text, &consumed);
+        if (!trim(text.substr(consumed)).empty()) {
+            return false;
+        }
+        value = parsed;
+        return true;
+    } catch (const std::exception&) {
+        return false;
+    }
+}
+
+bool parse_strict_int(const std::string& text, int& value) {
+    try {
+        size_t consumed = 0;
+        int parsed      = std::stoi(text, &consumed);
+        if (!trim(text.substr(consumed)).empty()) {
+            return false;
+        }
+        value = parsed;
+        return true;
+    } catch (const std::exception&) {
+        return false;
+    }
+}
+
+bool parse_strict_bool(const std::string& text, bool& value) {
+    std::string lowered = trim(text);
+    std::transform(lowered.begin(), lowered.end(), lowered.begin(), [](unsigned char c) {
+        return static_cast<char>(std::tolower(c));
+    });
+
+    if (lowered == "1" || lowered == "true" || lowered == "yes" || lowered == "on") {
+        value = true;
+        return true;
+    }
+    if (lowered == "0" || lowered == "false" || lowered == "no" || lowered == "off") {
+        value = false;
+        return true;
+    }
+    return false;
+}
+
 static std::string build_progress_bar(int step, int steps) {
     std::string progress = "  |";
     int max_progress     = 50;

src/util.h

@@ -4,6 +4,7 @@
 #include <cstdint>
 #include <memory>
 #include <string>
+#include <utility>
 #include <vector>
 
 #include "ggml-backend.h"
@@ -65,6 +66,15 @@ protected:
 
 std::string path_join(const std::string& p1, const std::string& p2);
 std::vector<std::string> split_string(const std::string& str, char delimiter);
+
+using KeyValueArgs = std::vector<std::pair<std::string, std::string>>;
+
+KeyValueArgs parse_key_value_args(const char* args, const char* context = "key=value arg");
+KeyValueArgs parse_key_value_args(const std::string& args, const char* context = "key=value arg");
+bool parse_strict_float(const std::string& text, float& value);
+bool parse_strict_int(const std::string& text, int& value);
+bool parse_strict_bool(const std::string& text, bool& value);
+
 void pretty_progress(int step, int steps, float time);
 void pretty_bytes_progress(int step, int steps, uint64_t bytes_processed, float elapsed_seconds);
 

src/vae.hpp

@@ -167,6 +167,7 @@ public:
         int64_t t0              = ggml_time_ms();
         sd::Tensor<float> input = x;
         sd::Tensor<float> output;
+        set_tiling_params(tiling_params);
 
         if (tiling_params.enabled) {
             const int scale_factor = get_scale_factor();
@@ -216,6 +217,9 @@ public:
     virtual void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors, const std::string prefix)         = 0;
     virtual void set_conv2d_scale(float scale) { SD_UNUSED(scale); };
     virtual void set_temporal_tiling_enabled(bool enabled) { SD_UNUSED(enabled); };
+    virtual void set_tiling_params(const sd_tiling_params_t& params) {
+        set_temporal_tiling_enabled(params.temporal_tiling);
+    };
 };
 
 struct FakeVAE : public VAE {