calculate_shift

Author: hlky

src/diffusers/pipelines/flux/pipeline_flux.py

@@ -34,6 +34,7 @@
 from ...utils import (
     USE_PEFT_BACKEND,
     is_torch_xla_available,
+    deprecate,
     logging,
     replace_example_docstring,
     scale_lora_layers,
@@ -73,12 +74,24 @@
 
 def calculate_shift(
     image_seq_len,
-    scheduler,
+    base_seq_len: Optional[int] = 256,
+    max_seq_len: Optional[int] = 4096,
+    base_shift: Optional[float] = 0.5,
+    max_shift: Optional[float] = 1.16,
+    scheduler: Optional[FlowMatchEulerDiscreteScheduler] = None,
 ):
-    base_seq_len = scheduler.config.get("base_image_seq_len", 256)
-    max_seq_len = scheduler.config.get("max_image_seq_len", 4096)
-    base_shift = scheduler.config.get("base_shift", 0.5)
-    max_shift = scheduler.config.get("max_shift", 1.16)
+    if base_seq_len or max_seq_len or base_shift or max_shift or scheduler is None:
+        deprecation_message = "Pass `scheduler` to `calculate_shift`."
+        deprecate(
+            "calculate_shift scheduler",
+            "1.0.0",
+            deprecation_message,
+            standard_warn=False,
+        )
+    base_seq_len = base_seq_len or scheduler.config.get("base_image_seq_len", 256)
+    max_seq_len = max_seq_len or scheduler.config.get("max_image_seq_len", 4096)
+    base_shift = base_shift or scheduler.config.get("base_shift", 0.5)
+    max_shift = max_shift or scheduler.config.get("max_shift", 1.16)
     m = (max_shift - base_shift) / (max_seq_len - base_seq_len)
     b = base_shift - m * base_seq_len
     mu = image_seq_len * m + b