Small fixed for DFE, polar guidance, and other things

toolkit/sampler.py

@@ -88,6 +88,18 @@ flux_config = {
   "use_dynamic_shifting": True
 }
 
+sd_flow_config = {
+  "_class_name": "FlowMatchEulerDiscreteScheduler",
+  "_diffusers_version": "0.30.0.dev0",
+  "base_image_seq_len": 256,
+  "base_shift": 0.5,
+  "max_image_seq_len": 4096,
+  "max_shift": 1.15,
+  "num_train_timesteps": 1000,
+  "shift": 3.0,
+  "use_dynamic_shifting": False
+}
+
 
 def get_sampler(
         sampler: str,
@@ -133,6 +145,8 @@ def get_sampler(
     elif sampler == "flowmatch":
         scheduler_cls = CustomFlowMatchEulerDiscreteScheduler
         config_to_use = copy.deepcopy(flux_config)
+        if arch == "sd":
+            config_to_use = copy.deepcopy(sd_flow_config)
     else:
         raise ValueError(f"Sampler {sampler} not supported")