Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

flux: make scheduler config params optional #10384

Open
wants to merge 3 commits into
base: main
Choose a base branch
from
Open

flux: make scheduler config params optional #10384

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
66806e3
dont assume scheduler has optional config params
vladmandic Dec 25, 2024
b3c6787
make style, make fix-copies
hlky Dec 25, 2024
64aee6e
calculate_shift
hlky Dec 25, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
27 changes: 19 additions & 8 deletions src/diffusers/pipelines/flux/pipeline_flux.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -34,6 +34,7 @@
from ...utils import (
USE_PEFT_BACKEND,
is_torch_xla_available,
deprecate,
logging,
replace_example_docstring,
scale_lora_layers,
Expand Down Expand Up @@ -73,11 +74,24 @@

def calculate_shift(
image_seq_len,
base_seq_len: int = 256,
max_seq_len: int = 4096,
base_shift: float = 0.5,
max_shift: float = 1.16,
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,
):
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)
Comment on lines +77 to +94
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks for the contribution @vladmandic, it will be fun to try to those scheduler changes with Flux.

As you mentioned CI, I've ran make fix-copies to demonstrate. Every occurrence of # Copied from diffusers.pipelines.flux.pipeline_flux.calculate_shift is updated, this highlights the effects in our own codebase, we'll need to update the call of calculate_shift in every pipeline and there can be effects downstream, so I've pushed this version which adds a deprecation message. We'll run make fix-copies again when the changes here are finalised.

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

thanks!

m = (max_shift - base_shift) / (max_seq_len - base_seq_len)
b = base_shift - m * base_seq_len
mu = image_seq_len * m + b
Expand Down Expand Up @@ -824,10 +838,7 @@ def __call__(
image_seq_len = latents.shape[1]
mu = calculate_shift(
image_seq_len,
self.scheduler.config.base_image_seq_len,
self.scheduler.config.max_image_seq_len,
self.scheduler.config.base_shift,
self.scheduler.config.max_shift,
self.scheduler,
)
timesteps, num_inference_steps = retrieve_timesteps(
self.scheduler,
Expand Down
9 changes: 5 additions & 4 deletions src/diffusers/pipelines/flux/pipeline_flux_control_img2img.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -90,11 +90,12 @@
# Copied from diffusers.pipelines.flux.pipeline_flux.calculate_shift
def calculate_shift(
image_seq_len,
base_seq_len: int = 256,
max_seq_len: int = 4096,
base_shift: float = 0.5,
max_shift: float = 1.16,
scheduler,
):
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)
m = (max_shift - base_shift) / (max_seq_len - base_seq_len)
b = base_shift - m * base_seq_len
mu = image_seq_len * m + b
Expand Down
9 changes: 5 additions & 4 deletions src/diffusers/pipelines/flux/pipeline_flux_control_inpaint.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -116,11 +116,12 @@
# Copied from diffusers.pipelines.flux.pipeline_flux.calculate_shift
def calculate_shift(
image_seq_len,
base_seq_len: int = 256,
max_seq_len: int = 4096,
base_shift: float = 0.5,
max_shift: float = 1.16,
scheduler,
):
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)
m = (max_shift - base_shift) / (max_seq_len - base_seq_len)
b = base_shift - m * base_seq_len
mu = image_seq_len * m + b
Expand Down
9 changes: 5 additions & 4 deletions src/diffusers/pipelines/flux/pipeline_flux_controlnet.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -86,11 +86,12 @@
# Copied from diffusers.pipelines.flux.pipeline_flux.calculate_shift
def calculate_shift(
image_seq_len,
base_seq_len: int = 256,
max_seq_len: int = 4096,
base_shift: float = 0.5,
max_shift: float = 1.16,
scheduler,
):
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)
m = (max_shift - base_shift) / (max_seq_len - base_seq_len)
b = base_shift - m * base_seq_len
mu = image_seq_len * m + b
Expand Down
9 changes: 5 additions & 4 deletions src/diffusers/pipelines/flux/pipeline_flux_controlnet_image_to_image.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -84,11 +84,12 @@
# Copied from diffusers.pipelines.flux.pipeline_flux.calculate_shift
def calculate_shift(
image_seq_len,
base_seq_len: int = 256,
max_seq_len: int = 4096,
base_shift: float = 0.5,
max_shift: float = 1.16,
scheduler,
):
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)
m = (max_shift - base_shift) / (max_seq_len - base_seq_len)
b = base_shift - m * base_seq_len
mu = image_seq_len * m + b
Expand Down
9 changes: 5 additions & 4 deletions src/diffusers/pipelines/flux/pipeline_flux_controlnet_inpainting.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -86,11 +86,12 @@
# Copied from diffusers.pipelines.flux.pipeline_flux.calculate_shift
def calculate_shift(
image_seq_len,
base_seq_len: int = 256,
max_seq_len: int = 4096,
base_shift: float = 0.5,
max_shift: float = 1.16,
scheduler,
):
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)
m = (max_shift - base_shift) / (max_seq_len - base_seq_len)
b = base_shift - m * base_seq_len
mu = image_seq_len * m + b
Expand Down
9 changes: 5 additions & 4 deletions src/diffusers/pipelines/flux/pipeline_flux_fill.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -79,11 +79,12 @@
# Copied from diffusers.pipelines.flux.pipeline_flux.calculate_shift
def calculate_shift(
image_seq_len,
base_seq_len: int = 256,
max_seq_len: int = 4096,
base_shift: float = 0.5,
max_shift: float = 1.16,
scheduler,
):
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)
m = (max_shift - base_shift) / (max_seq_len - base_seq_len)
b = base_shift - m * base_seq_len
mu = image_seq_len * m + b
Expand Down
9 changes: 5 additions & 4 deletions src/diffusers/pipelines/flux/pipeline_flux_img2img.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -74,11 +74,12 @@
# Copied from diffusers.pipelines.flux.pipeline_flux.calculate_shift
def calculate_shift(
image_seq_len,
base_seq_len: int = 256,
max_seq_len: int = 4096,
base_shift: float = 0.5,
max_shift: float = 1.16,
scheduler,
):
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)
m = (max_shift - base_shift) / (max_seq_len - base_seq_len)
b = base_shift - m * base_seq_len
mu = image_seq_len * m + b
Expand Down
9 changes: 5 additions & 4 deletions src/diffusers/pipelines/flux/pipeline_flux_inpaint.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -71,11 +71,12 @@
# Copied from diffusers.pipelines.flux.pipeline_flux.calculate_shift
def calculate_shift(
image_seq_len,
base_seq_len: int = 256,
max_seq_len: int = 4096,
base_shift: float = 0.5,
max_shift: float = 1.16,
scheduler,
):
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)
m = (max_shift - base_shift) / (max_seq_len - base_seq_len)
b = base_shift - m * base_seq_len
mu = image_seq_len * m + b
Expand Down
9 changes: 5 additions & 4 deletions src/diffusers/pipelines/ltx/pipeline_ltx.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -69,11 +69,12 @@
# Copied from diffusers.pipelines.flux.pipeline_flux.calculate_shift
def calculate_shift(
image_seq_len,
base_seq_len: int = 256,
max_seq_len: int = 4096,
base_shift: float = 0.5,
max_shift: float = 1.16,
scheduler,
):
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)
m = (max_shift - base_shift) / (max_seq_len - base_seq_len)
b = base_shift - m * base_seq_len
mu = image_seq_len * m + b
Expand Down
9 changes: 5 additions & 4 deletions src/diffusers/pipelines/ltx/pipeline_ltx_image2video.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -74,11 +74,12 @@
# Copied from diffusers.pipelines.flux.pipeline_flux.calculate_shift
def calculate_shift(
image_seq_len,
base_seq_len: int = 256,
max_seq_len: int = 4096,
base_shift: float = 0.5,
max_shift: float = 1.16,
scheduler,
):
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)
m = (max_shift - base_shift) / (max_seq_len - base_seq_len)
b = base_shift - m * base_seq_len
mu = image_seq_len * m + b
Expand Down
9 changes: 5 additions & 4 deletions src/diffusers/pipelines/stable_diffusion_3/pipeline_stable_diffusion_3.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -73,11 +73,12 @@
# Copied from diffusers.pipelines.flux.pipeline_flux.calculate_shift
def calculate_shift(
image_seq_len,
base_seq_len: int = 256,
max_seq_len: int = 4096,
base_shift: float = 0.5,
max_shift: float = 1.16,
scheduler,
):
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)
m = (max_shift - base_shift) / (max_seq_len - base_seq_len)
b = base_shift - m * base_seq_len
mu = image_seq_len * m + b
Expand Down
9 changes: 5 additions & 4 deletions src/diffusers/pipelines/stable_diffusion_3/pipeline_stable_diffusion_3_img2img.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -78,11 +78,12 @@
# Copied from diffusers.pipelines.flux.pipeline_flux.calculate_shift
def calculate_shift(
image_seq_len,
base_seq_len: int = 256,
max_seq_len: int = 4096,
base_shift: float = 0.5,
max_shift: float = 1.16,
scheduler,
):
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)
m = (max_shift - base_shift) / (max_seq_len - base_seq_len)
b = base_shift - m * base_seq_len
mu = image_seq_len * m + b
Expand Down
9 changes: 5 additions & 4 deletions src/diffusers/pipelines/stable_diffusion_3/pipeline_stable_diffusion_3_inpaint.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -77,11 +77,12 @@
# Copied from diffusers.pipelines.flux.pipeline_flux.calculate_shift
def calculate_shift(
image_seq_len,
base_seq_len: int = 256,
max_seq_len: int = 4096,
base_shift: float = 0.5,
max_shift: float = 1.16,
scheduler,
):
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)
m = (max_shift - base_shift) / (max_seq_len - base_seq_len)
b = base_shift - m * base_seq_len
mu = image_seq_len * m + b
Expand Down