Merge pull request #22 from black-forest-labs/fix-streamlit-step-modulation

Fix the step modulation in streamlit

demo_st.py

@@ -15,11 +15,18 @@ from transformers import pipeline
 
 from flux.cli import SamplingOptions
 from flux.sampling import denoise, get_noise, get_schedule, prepare, unpack
-from flux.util import (configs, embed_watermark, load_ae, load_clip,
-                       load_flow_model, load_t5)
+from flux.util import (
+    configs,
+    embed_watermark,
+    load_ae,
+    load_clip,
+    load_flow_model,
+    load_t5,
+)
 
 NSFW_THRESHOLD = 0.85
 
+
 @st.cache_resource()
 def get_models(name: str, device: torch.device, offload: bool, is_schnell: bool):
     t5 = load_t5(device, max_length=256 if is_schnell else 512)
@@ -187,7 +194,7 @@ def main(
         # divide pixel space by 16**2 to acocunt for latent space conversion
         timesteps = get_schedule(
             opts.num_steps,
-            x.shape[-1] * x.shape[-2] // (16 * 16),
+            (x.shape[-1] * x.shape[-2]) // 4,
             shift=(not is_schnell),
         )
         if init_image is not None: