fix: enable Wan/TAEHV video generation on the Metal backend (#850) (#1731)

This commit is contained in:
mmandelker-code authored and GitHub committed 2026-07-03 00:38:13 +08:00
1 parent 556f04bb3f
commit 7dab3668d0
4 files changed
+39 -12

No files matched your search

+2 -1
View File
@@ -214,7 +214,7 @@ namespace LTXV {
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);
x = ggml_ext_pad_ext(ctx, runner_ctx->backend, x, static_cast<int>(left_pad), 0, 0, 0, 0, 0, 0, 0);
}
auto frames = ggml_conv_1d(ctx, forward_basis, x, hop_length, 0, 1);
@@ -451,6 +451,7 @@ namespace LTXV {
int pad_h = kernel_size.first - 1;
int pad_w = kernel_size.second - 1;
x = ggml_ext_pad_ext(ctx->ggml_ctx,
ctx->backend,
x,
pad_w / 2,
pad_w - pad_w / 2,
+2 -2
View File
@@ -408,7 +408,7 @@ public:
h = conv->forward(ctx, h);
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);
auto mem = ggml_ext_pad_ext(ctx->ggml_ctx, ctx->backend, 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);
}
@@ -479,7 +479,7 @@ public:
int index = 3;
for (int i = 0; i < num_layers; i++) {
for (int j = 0; j < num_blocks; j++) {
auto mem = ggml_pad_ext(ctx->ggml_ctx, h, 0, 0, 0, 0, 0, 0, 1, 0);
auto mem = ggml_ext_pad_ext(ctx->ggml_ctx, ctx->backend, 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++)]);
+4 -4
View File
@@ -72,8 +72,8 @@ namespace WAN {
lp2 -= (int)cache_x->ne[2];
}
x = ggml_ext_pad_ext(ctx->ggml_ctx, x, lp0, rp0, lp1, rp1, lp2, rp2, 0, 0, ctx->circular_x_enabled, ctx->circular_y_enabled);
return ggml_ext_conv_3d(ctx->ggml_ctx, x, w, b, in_channels,
x = ggml_ext_pad_ext(ctx->ggml_ctx, ctx->backend, x, lp0, rp0, lp1, rp1, lp2, rp2, 0, 0, ctx->circular_x_enabled, ctx->circular_y_enabled);
return ggml_ext_conv_3d(ctx->ggml_ctx, ctx->backend, x, w, b, in_channels,
std::get<2>(stride), std::get<1>(stride), std::get<0>(stride),
0, 0, 0,
std::get<2>(dilation), std::get<1>(dilation), std::get<0>(dilation));
@@ -195,7 +195,7 @@ namespace WAN {
2);
}
if (chunk_idx == 1 && cache_x->ne[2] < 2) { // Rep
cache_x = ggml_pad_ext(ctx->ggml_ctx, cache_x, 0, 0, 0, 0, (int)cache_x->ne[2], 0, 0, 0);
cache_x = ggml_ext_pad_ext(ctx->ggml_ctx, ctx->backend, cache_x, 0, 0, 0, 0, (int)cache_x->ne[2], 0, 0, 0);
// aka cache_x = torch.cat([torch.zeros_like(cache_x).to(cache_x.device),cache_x],dim=2)
}
if (chunk_idx == 1) {
@@ -283,7 +283,7 @@ namespace WAN {
int pad_t = (factor_t - T % factor_t) % factor_t;
x = ggml_pad_ext(ctx->ggml_ctx, x, 0, 0, 0, 0, pad_t, 0, 0, 0);
x = ggml_ext_pad_ext(ctx->ggml_ctx, ctx->backend, x, 0, 0, 0, 0, pad_t, 0, 0, 0);
T = x->ne[2];
x = ggml_reshape_4d(ctx->ggml_ctx, x, W * H, factor_t, T / factor_t, C); // [C, T/factor_t, factor_t, H*W]