refactor: consolidate WAN VAE version checks (#1712)

.github/workflows/build.yml

@@ -207,6 +207,7 @@ jobs:
               UBUNTU_VERSION=24.04
               CUDA_ARCHITECTURES=121
               GGML_CUDA_FA_ALL_QUANTS=ON
+              GGML_CUDA_ENABLE_DYNAMIC_CPU_BACKENDS=OFF
 
     env:
       REGISTRY: ghcr.io

Dockerfile.cuda

@@ -23,18 +23,27 @@ COPY . .
 ARG CUDACXX=/usr/local/cuda/bin/nvcc
 ARG CUDA_ARCHITECTURES=""
 ARG GGML_CUDA_FA_ALL_QUANTS=""
+ARG GGML_CUDA_ENABLE_DYNAMIC_CPU_BACKENDS=ON
 
-RUN cmake . -B ./build \
+RUN set -- \
-    -DSD_CUDA=ON \
+        -DSD_CUDA=ON; \
-    -DSD_BUILD_SHARED_LIBS=ON \
+    if [ "${GGML_CUDA_ENABLE_DYNAMIC_CPU_BACKENDS}" = "ON" ]; then \
-    -DGGML_NATIVE=OFF \
+        set -- "$@" \
-    -DSD_BUILD_SHARED_GGML_LIB=ON \
+            -DSD_BUILD_SHARED_LIBS=ON \
-    -DGGML_BACKEND_DL=ON \
+            -DGGML_NATIVE=OFF \
-    -DGGML_CPU_ALL_VARIANTS=ON \
+            -DSD_BUILD_SHARED_GGML_LIB=ON \
-    -DCMAKE_BUILD_WITH_INSTALL_RPATH=ON \
+            -DGGML_BACKEND_DL=ON \
-    -DCMAKE_INSTALL_RPATH='$ORIGIN' \
+            -DGGML_CPU_ALL_VARIANTS=ON \
-    ${CUDA_ARCHITECTURES:+-DCMAKE_CUDA_ARCHITECTURES="${CUDA_ARCHITECTURES}"} \
+            -DCMAKE_BUILD_WITH_INSTALL_RPATH=ON \
-    ${GGML_CUDA_FA_ALL_QUANTS:+-DGGML_CUDA_FA_ALL_QUANTS=${GGML_CUDA_FA_ALL_QUANTS}}
+            "-DCMAKE_INSTALL_RPATH=\$ORIGIN"; \
+    fi; \
+    if [ -n "${CUDA_ARCHITECTURES}" ]; then \
+        set -- "$@" "-DCMAKE_CUDA_ARCHITECTURES=${CUDA_ARCHITECTURES}"; \
+    fi; \
+    if [ -n "${GGML_CUDA_FA_ALL_QUANTS}" ]; then \
+        set -- "$@" "-DGGML_CUDA_FA_ALL_QUANTS=${GGML_CUDA_FA_ALL_QUANTS}"; \
+    fi; \
+    cmake . -B ./build "$@"
 RUN cmake --build ./build --config Release -j$(nproc)
 
 FROM nvidia/cuda:${CUDA_VERSION}-cudnn-runtime-ubuntu${UBUNTU_VERSION} AS runtime

Dockerfile.vulkan

@@ -28,7 +28,7 @@ RUN cmake . -B ./build \
     -DGGML_CPU_ALL_VARIANTS=ON \
     -DCMAKE_BUILD_WITH_INSTALL_RPATH=ON \
     -DCMAKE_INSTALL_RPATH='$ORIGIN'
-RUN cmake --build ./build --config Release --parallel
+RUN cmake --build ./build --config Release -j$(nproc)
 
 FROM ubuntu:$UBUNTU_VERSION AS runtime
 

ggml

@@ -1 +1 @@
-Subproject commit 3af5f5760e19a96427f5f7a93b79cbdf3d4b265b
+Subproject commit eced84c86f8b012c752c016f7fe789adea168e1e

src/model.h

@@ -208,6 +208,13 @@ static inline bool sd_version_uses_flux2_vae(SDVersion version) {
     return false;
 }
 
+static inline bool sd_version_uses_wan_vae(SDVersion version) {
+    if (sd_version_is_wan(version) || sd_version_is_qwen_image(version) || sd_version_is_krea2(version) || sd_version_is_anima(version)) {
+        return true;
+    }
+    return false;
+}
+
 static inline bool sd_version_is_inpaint(SDVersion version) {
     if (version == VERSION_SD1_INPAINT ||
         version == VERSION_SD2_INPAINT ||

src/model/diffusion/anima.hpp

@@ -227,7 +227,6 @@ namespace Anima {
             k4 = k_norm->forward(ctx, k4);
 
             ggml_tensor* attn_out = nullptr;
-            float scale           = (sd_backend_is(ctx->backend, "Vulkan") && ctx->flash_attn_enabled) ? 1.0f / 32.0f : 1.0f;
             if (pe_q != nullptr || pe_k != nullptr) {
                 if (pe_q == nullptr) {
                     pe_q = pe_k;
@@ -245,8 +244,7 @@ namespace Anima {
                                                      num_heads,
                                                      nullptr,
                                                      true,
-                                                     ctx->flash_attn_enabled,
+                                                     ctx->flash_attn_enabled);
-                                                     scale);
             } else {
                 auto q_flat = ggml_reshape_3d(ctx->ggml_ctx, q4, head_dim * num_heads, L_q, N);
                 auto k_flat = ggml_reshape_3d(ctx->ggml_ctx, k4, head_dim * num_heads, L_k, N);
@@ -258,8 +256,7 @@ namespace Anima {
                                                      num_heads,
                                                      nullptr,
                                                      false,
-                                                     ctx->flash_attn_enabled,
+                                                     ctx->flash_attn_enabled);
-                                                     scale);
             }
 
             return out_proj->forward(ctx, attn_out);

src/model/diffusion/ernie_image.hpp

@@ -162,8 +162,6 @@ namespace ErnieImage {
             int64_t S = x->ne[1];
             int64_t N = x->ne[2];
 
-            float scale = (sd_backend_is(ctx->backend, "Vulkan") && ctx->flash_attn_enabled) ? 1.0f / 32.0f : 1.0f;
-
             auto q = to_q->forward(ctx, x);
             auto k = to_k->forward(ctx, x);
             auto v = to_v->forward(ctx, x);
@@ -184,7 +182,7 @@ namespace ErnieImage {
             k = ggml_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, k, 0, 2, 1, 3));  // [N, heads, S, head_dim]
             k = ggml_reshape_3d(ctx->ggml_ctx, k, k->ne[0], k->ne[1], k->ne[2] * k->ne[3]);
 
-            x = ggml_ext_attention_ext(ctx->ggml_ctx, ctx->backend, q, k, v, num_heads, attention_mask, true, ctx->flash_attn_enabled, scale);  // [N, S, hidden_size]
+            x = ggml_ext_attention_ext(ctx->ggml_ctx, ctx->backend, q, k, v, num_heads, attention_mask, true, ctx->flash_attn_enabled);  // [N, S, hidden_size]
             x = to_out_0->forward(ctx, x);
             return x;
         }

src/model/vae/tae.hpp

@@ -548,7 +548,7 @@ public:
         }
         auto result = decoder->forward(ctx, z);
         if (sd_version_is_wan(version) || sd_version_is_ltxav(version)) {
-            // (W, H, C, T) -> (W, H, T, C)
+            // (W, H, T, C) -> (W, H, C, T)
             result = ggml_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, result, 0, 1, 3, 2));
         }
         return result;
@@ -556,8 +556,10 @@ public:
 
     ggml_tensor* encode(GGMLRunnerContext* ctx, ggml_tensor* x) {
         auto encoder = std::dynamic_pointer_cast<TinyVideoEncoder>(blocks["encoder"]);
-        // (W, H, T, C) -> (W, H, C, T)
+        if (sd_version_is_wan(version) || sd_version_is_ltxav(version)) {
-        x                  = ggml_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, x, 0, 1, 3, 2));
+            // (W, H, T, C) -> (W, H, C, T)
+            x = ggml_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, x, 0, 1, 3, 2));
+        }
         int64_t num_frames = x->ne[3];
         if (num_frames % encoder->t_downscale) {
             // pad to multiple of encoder->t_downscale at the end
@@ -567,7 +569,10 @@ public:
             }
         }
         x = encoder->forward(ctx, x);
-        x = ggml_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, x, 0, 1, 3, 2));
+        if (sd_version_is_wan(version) || sd_version_is_ltxav(version)) {
+            // (W, H, C, T) -> (W, H, T, C)
+            x = ggml_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, x, 0, 1, 3, 2));
+        }
         return x;
     }
 };

src/runtime/denoiser.hpp

@@ -616,7 +616,6 @@ struct LogitNormalScheduler : SigmaScheduler {
         one_minus_t_min = sigmoid(0.5f * logsnr_max);
         // t_max = 1.0f / (1.0f + std::exp(0.5f * logsnr_min));
         one_minus_t_max = sigmoid(0.5f * logsnr_min);
-
     }
 
     LogitNormalScheduler(int image_seq_len = 0, const char* extra_sample_args = nullptr) {

src/stable-diffusion.cpp

@@ -892,11 +892,7 @@ public:
             }
 
             auto create_tae = [&](bool decode_only) -> std::shared_ptr<VAE> {
-                if (sd_version_is_wan(version) ||
+                if (sd_version_uses_wan_vae(version) || sd_version_is_ltxav(version)) {
-                    sd_version_is_qwen_image(version) ||
-                    sd_version_is_krea2(version) ||
-                    sd_version_is_anima(version) ||
-                    sd_version_is_ltxav(version)) {
                     return std::make_shared<TinyVideoAutoEncoder>(backend_for(SDBackendModule::VAE),
                                                                   tensor_storage_map,
                                                                   "decoder",
@@ -933,10 +929,7 @@ public:
                                                          false,
                                                          version,
                                                          model_manager);
-                } else if (sd_version_is_wan(version) ||
+                } else if (sd_version_uses_wan_vae(version)) {
-                           sd_version_is_qwen_image(version) ||
-                           sd_version_is_krea2(version) ||
-                           sd_version_is_anima(version)) {
                     return std::make_shared<WAN::WanVAERunner>(backend_for(SDBackendModule::VAE),
                                                                tensor_storage_map,
                                                                "first_stage_model",
@@ -1742,7 +1735,7 @@ public:
                 } else if (sd_version_uses_flux_vae(version)) {
                     latent_rgb_proj = flux_latent_rgb_proj;
                     latent_rgb_bias = flux_latent_rgb_bias;
-                } else if (sd_version_is_wan(version) || sd_version_is_qwen_image(version) || sd_version_is_anima(version) || sd_version_is_krea2(version)) {
+                } else if (sd_version_uses_wan_vae(version)) {
                     latent_rgb_proj = wan_21_latent_rgb_proj;
                     latent_rgb_bias = wan_21_latent_rgb_bias;
                 } else {
@@ -3156,7 +3149,7 @@ enum scheduler_t sd_get_default_scheduler(const sd_ctx_t* sd_ctx, enum sample_me
         return SIMPLE_SCHEDULER;
     } else if (sd_ctx != nullptr && sd_ctx->sd != nullptr && sd_version_is_ltxav(sd_ctx->sd->version)) {
         return LTX2_SCHEDULER;
-    } else if(sd_ctx != nullptr && sd_ctx->sd != nullptr && sd_version_is_ideogram4(sd_ctx->sd->version)) {
+    } else if (sd_ctx != nullptr && sd_ctx->sd != nullptr && sd_version_is_ideogram4(sd_ctx->sd->version)) {
         return LOGIT_NORMAL_SCHEDULER;
     }
     return DISCRETE_SCHEDULER;