change preprocess vae check

nod-ai · Jul 22, 2024 · 32fc090 · 32fc090
1 parent fbaecd6
commit 32fc090
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/models/turbine_models/custom_models/sd_inference/utils.py b/models/turbine_models/custom_models/sd_inference/utils.py
@@ -252,7 +252,7 @@ def compile_to_vmfb(
             flags.extend(MI_flags["pad_attention"])
         elif "punet" in flagset_keywords:
             flags.extend(MI_flags["punet"])
-        elif "vae_preprocess" in flagset_keywords:
+        elif "vae" in safe_name:
             flags.extend(MI_flags["vae_preprocess"])
         else:
             flags.extend(MI_flags["preprocess_default"])