Diffusion model case

EricLBuehler · Sep 18, 2024 · 9153aff · 9153aff
1 parent 8dd0000
commit 9153aff
Showing 1 changed file with 4 additions and 1 deletion.
diff --git a/mistralrs-core/src/engine/mod.rs b/mistralrs-core/src/engine/mod.rs
@@ -77,7 +77,10 @@ impl Engine {
         let device = get_mut_arcmutex!(pipeline).device().clone();
         let is_xlora = get_mut_arcmutex!(pipeline).get_metadata().is_xlora;
         let has_no_kv_cache = get_mut_arcmutex!(pipeline).get_metadata().has_no_kv_cache;
-        assert_eq!(has_no_kv_cache, no_kv_cache);
+        if no_kv_cache {
+            // Diffusion models...
+            assert_eq!(has_no_kv_cache, no_kv_cache);
+        }
         // Prefix caching is always disabled if using PagedAttention for now.
         // TODO
         let no_prefix_cache = matches!(config, SchedulerConfig::PagedAttentionMeta { .. })