Skip to content

flux: make scheduler config params optional #10384

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

Merged
merged 7 commits into from
Jan 9, 2025
Merged

flux: make scheduler config params optional #10384

Show file tree
Hide file tree
Changes from 5 commits
Commits
Show all changes
7 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
20614bd
fix-copies, usage in pipelines
hlky Dec 30, 2024
aa1abb0
Merge branch 'main' into flux
hlky Dec 30, 2024
88123a3
hlky Jan 9, 2025
f77efec
Merge branch 'main' into flux
hlky Jan 9, 2025
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 examples/community/pipeline_flux_differential_img2img.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -29,6 +29,7 @@
from diffusers.schedulers import FlowMatchEulerDiscreteScheduler
from diffusers.utils import (
USE_PEFT_BACKEND,
deprecate,
is_torch_xla_available,
logging,
replace_example_docstring,
Expand Down Expand Up @@ -84,11 +85,24 @@
# 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,
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)
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 @@ -876,10 +890,7 @@ def __call__(
image_seq_len = (int(height) // self.vae_scale_factor) * (int(width) // self.vae_scale_factor)
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,
scheduler=self.scheduler,
)
timesteps, num_inference_steps = retrieve_timesteps(
self.scheduler,
Expand Down
32 changes: 20 additions & 12 deletions examples/community/pipeline_flux_rf_inversion.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -30,6 +30,7 @@
from diffusers.schedulers import FlowMatchEulerDiscreteScheduler
from diffusers.utils import (
USE_PEFT_BACKEND,
deprecate,
is_torch_xla_available,
logging,
replace_example_docstring,
Expand Down Expand Up @@ -91,11 +92,24 @@
# 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,
base_seq_len: Optional[int] = 256,
Copy link
Collaborator

Choose a reason for hiding this comment

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

I think we don't need to change this function, we can just change how the parameters are passed from the pipeline

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)
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 @@ -822,10 +836,7 @@ def __call__(
image_seq_len = (int(height) // self.vae_scale_factor // 2) * (int(width) // self.vae_scale_factor // 2)
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,
scheduler=self.scheduler,
)
timesteps, num_inference_steps = retrieve_timesteps(
self.scheduler,
Expand Down Expand Up @@ -992,10 +1003,7 @@ def invert(
image_seq_len = (int(height) // self.vae_scale_factor // 2) * (int(width) // self.vae_scale_factor // 2)
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,
scheduler=self.scheduler,
Copy link
Collaborator

Choose a reason for hiding this comment

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

Suggested change
scheduler=self.scheduler,
scheduler.config.get("base_image_seq_len", 256)
scheduler.config.get("max_image_seq_len", 4096)
scheduler.config.get("base_shift", 0.5)
scheduler.config.get("max_shift", 1.16)

)
timesteps, num_inversion_steps = retrieve_timesteps(
self.scheduler,
Expand Down
28 changes: 20 additions & 8 deletions examples/community/pipeline_flux_with_cfg.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -28,6 +28,7 @@
from diffusers.schedulers import FlowMatchEulerDiscreteScheduler
from diffusers.utils import (
USE_PEFT_BACKEND,
deprecate,
is_torch_xla_available,
logging,
replace_example_docstring,
Expand Down Expand Up @@ -64,13 +65,27 @@
"""


# 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,
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)
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 @@ -757,10 +772,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,
scheduler=self.scheduler,
)
timesteps, num_inference_steps = retrieve_timesteps(
self.scheduler,
Expand Down
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 @@ -33,6 +33,7 @@
from ...schedulers import FlowMatchEulerDiscreteScheduler
from ...utils import (
USE_PEFT_BACKEND,
deprecate,
is_torch_xla_available,
logging,
replace_example_docstring,
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)
Copy link
Contributor

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,
scheduler=self.scheduler,
)
timesteps, num_inference_steps = retrieve_timesteps(
self.scheduler,
Expand Down
28 changes: 20 additions & 8 deletions src/diffusers/pipelines/flux/pipeline_flux_control.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -26,6 +26,7 @@
from ...schedulers import FlowMatchEulerDiscreteScheduler
from ...utils import (
USE_PEFT_BACKEND,
deprecate,
is_torch_xla_available,
logging,
replace_example_docstring,
Expand Down Expand Up @@ -82,13 +83,27 @@
"""


# 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,
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)
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 @@ -802,10 +817,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,
scheduler=self.scheduler,
)
timesteps, num_inference_steps = retrieve_timesteps(
self.scheduler,
Expand Down
27 changes: 19 additions & 8 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 @@ -26,6 +26,7 @@
from ...schedulers import FlowMatchEulerDiscreteScheduler
from ...utils import (
USE_PEFT_BACKEND,
deprecate,
is_torch_xla_available,
logging,
replace_example_docstring,
Expand Down Expand Up @@ -90,11 +91,24 @@
# 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,
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)
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 @@ -809,10 +823,7 @@ def __call__(
image_seq_len = (int(height) // self.vae_scale_factor // 2) * (int(width) // self.vae_scale_factor // 2)
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,
scheduler=self.scheduler,
)
timesteps, num_inference_steps = retrieve_timesteps(
self.scheduler,
Expand Down
27 changes: 19 additions & 8 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 @@ -35,6 +35,7 @@
from ...schedulers import FlowMatchEulerDiscreteScheduler
from ...utils import (
USE_PEFT_BACKEND,
deprecate,
is_torch_xla_available,
logging,
replace_example_docstring,
Expand Down Expand Up @@ -116,11 +117,24 @@
# 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,
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)
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 @@ -985,10 +999,7 @@ def __call__(
image_seq_len = (int(height) // self.vae_scale_factor // 2) * (int(width) // self.vae_scale_factor // 2)
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,
scheduler=self.scheduler,
)
timesteps, num_inference_steps = retrieve_timesteps(
self.scheduler,
Expand Down
Loading
Loading