Small fixed for DFE, polar guidance, and other things

2026-04-25 16:59:22 +00:00 · 2025-02-12 09:27:44 -07:00
parent 10aa7e9d5e
commit 787bb37e76
7 changed files with 87 additions and 43 deletions
--- a/toolkit/stable_diffusion_model.py
+++ b/toolkit/stable_diffusion_model.py
@@ -974,12 +974,17 @@ class StableDiffusion:
                            "prediction_type": self.prediction_type,
                        })
                else:
+                    arch = 'sd'
+                    if self.model_config.is_pixart:
+                        arch = 'pixart'
+                    if self.model_config.is_flux:
+                        arch = 'flux'
                    noise_scheduler = get_sampler(
                        sampler,
                        {
                            "prediction_type": self.prediction_type,
                        },
-                        'sd' if not self.is_pixart else 'pixart'
+                        arch
                    )

                try: