Skip to content

Update Flax removal version #13729

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
sayakpaul merged 2 commits into from
May 13, 2026
Merged

Update Flax removal version #13729

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
0397776
update
DN6 May 12, 2026
46f14a7
Merge branch 'main' into flax-dep
sayakpaul May 13, 2026
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
10 changes: 5 additions & 5 deletions src/diffusers/models/attention_flax.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -157,7 +157,7 @@ class FlaxAttention(nn.Module):

def setup(self):
logger.warning(
"Flax classes are deprecated and will be removed in Diffusers v1.0.0. We "
"Flax classes are deprecated and will be removed in Diffusers v0.40.0. We "
"recommend migrating to PyTorch classes or pinning your version of Diffusers."
)

Expand Down Expand Up @@ -288,7 +288,7 @@ class FlaxBasicTransformerBlock(nn.Module):

def setup(self):
logger.warning(
"Flax classes are deprecated and will be removed in Diffusers v1.0.0. We "
"Flax classes are deprecated and will be removed in Diffusers v0.40.0. We "
"recommend migrating to PyTorch classes or pinning your version of Diffusers."
)

Expand Down Expand Up @@ -381,7 +381,7 @@ class FlaxTransformer2DModel(nn.Module):

def setup(self):
logger.warning(
"Flax classes are deprecated and will be removed in Diffusers v1.0.0. We "
"Flax classes are deprecated and will be removed in Diffusers v0.40.0. We "
"recommend migrating to PyTorch classes or pinning your version of Diffusers."
)

Expand Down Expand Up @@ -475,7 +475,7 @@ class FlaxFeedForward(nn.Module):

def setup(self):
logger.warning(
"Flax classes are deprecated and will be removed in Diffusers v1.0.0. We "
"Flax classes are deprecated and will be removed in Diffusers v0.40.0. We "
"recommend migrating to PyTorch classes or pinning your version of Diffusers."
)

Expand Down Expand Up @@ -510,7 +510,7 @@ class FlaxGEGLU(nn.Module):

def setup(self):
logger.warning(
"Flax classes are deprecated and will be removed in Diffusers v1.0.0. We "
"Flax classes are deprecated and will be removed in Diffusers v0.40.0. We "
"recommend migrating to PyTorch classes or pinning your version of Diffusers."
)

Expand Down
4 changes: 2 additions & 2 deletions src/diffusers/models/controlnets/controlnet_flax.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -52,7 +52,7 @@ class FlaxControlNetConditioningEmbedding(nn.Module):

def setup(self) -> None:
logger.warning(
"Flax classes are deprecated and will be removed in Diffusers v1.0.0. We "
"Flax classes are deprecated and will be removed in Diffusers v0.40.0. We "
"recommend migrating to PyTorch classes or pinning your version of Diffusers."
)

Expand Down Expand Up @@ -191,7 +191,7 @@ def init_weights(self, rng: jax.Array) -> FrozenDict:

def setup(self) -> None:
logger.warning(
"Flax classes are deprecated and will be removed in Diffusers v1.0.0. We "
"Flax classes are deprecated and will be removed in Diffusers v0.40.0. We "
"recommend migrating to PyTorch classes or pinning your version of Diffusers."
)

Expand Down
4 changes: 2 additions & 2 deletions src/diffusers/models/embeddings_flax.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -82,7 +82,7 @@ class FlaxTimestepEmbedding(nn.Module):
"""

logger.warning(
"Flax classes are deprecated and will be removed in Diffusers v1.0.0. We "
"Flax classes are deprecated and will be removed in Diffusers v0.40.0. We "
"recommend migrating to PyTorch classes or pinning your version of Diffusers."
)

Expand Down Expand Up @@ -115,7 +115,7 @@ class FlaxTimesteps(nn.Module):
freq_shift: float = 1

logger.warning(
"Flax classes are deprecated and will be removed in Diffusers v1.0.0. We "
"Flax classes are deprecated and will be removed in Diffusers v0.40.0. We "
"recommend migrating to PyTorch classes or pinning your version of Diffusers."
)

Expand Down
2 changes: 1 addition & 1 deletion src/diffusers/models/modeling_flax_utils.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -285,7 +285,7 @@ def from_pretrained(
```
"""
logger.warning(
"Flax classes are deprecated and will be removed in Diffusers v1.0.0. We "
"Flax classes are deprecated and will be removed in Diffusers v0.40.0. We "
"recommend migrating to PyTorch classes or pinning your version of Diffusers."
)
config = kwargs.pop("config", None)
Expand Down
6 changes: 3 additions & 3 deletions src/diffusers/models/resnet_flax.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -27,7 +27,7 @@ class FlaxUpsample2D(nn.Module):

def setup(self):
logger.warning(
"Flax classes are deprecated and will be removed in Diffusers v1.0.0. We "
"Flax classes are deprecated and will be removed in Diffusers v0.40.0. We "
"recommend migrating to PyTorch classes or pinning your version of Diffusers."
)

Expand Down Expand Up @@ -56,7 +56,7 @@ class FlaxDownsample2D(nn.Module):

def setup(self):
logger.warning(
"Flax classes are deprecated and will be removed in Diffusers v1.0.0. We "
"Flax classes are deprecated and will be removed in Diffusers v0.40.0. We "
"recommend migrating to PyTorch classes or pinning your version of Diffusers."
)

Expand Down Expand Up @@ -84,7 +84,7 @@ class FlaxResnetBlock2D(nn.Module):

def setup(self):
logger.warning(
"Flax classes are deprecated and will be removed in Diffusers v1.0.0. We "
"Flax classes are deprecated and will be removed in Diffusers v0.40.0. We "
"recommend migrating to PyTorch classes or pinning your version of Diffusers."
)

Expand Down
10 changes: 5 additions & 5 deletions src/diffusers/models/unets/unet_2d_blocks_flax.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -65,7 +65,7 @@ class FlaxCrossAttnDownBlock2D(nn.Module):

def setup(self):
logger.warning(
"Flax classes are deprecated and will be removed in Diffusers v1.0.0. We "
"Flax classes are deprecated and will be removed in Diffusers v0.40.0. We "
"recommend migrating to PyTorch classes or pinning your version of Diffusers."
)

Expand Down Expand Up @@ -145,7 +145,7 @@ class FlaxDownBlock2D(nn.Module):

def setup(self):
logger.warning(
"Flax classes are deprecated and will be removed in Diffusers v1.0.0. We "
"Flax classes are deprecated and will be removed in Diffusers v0.40.0. We "
"recommend migrating to PyTorch classes or pinning your version of Diffusers."
)

Expand Down Expand Up @@ -223,7 +223,7 @@ class FlaxCrossAttnUpBlock2D(nn.Module):

def setup(self):
logger.warning(
"Flax classes are deprecated and will be removed in Diffusers v1.0.0. We "
"Flax classes are deprecated and will be removed in Diffusers v0.40.0. We "
"recommend migrating to PyTorch classes or pinning your version of Diffusers."
)

Expand Down Expand Up @@ -308,7 +308,7 @@ class FlaxUpBlock2D(nn.Module):

def setup(self):
logger.warning(
"Flax classes are deprecated and will be removed in Diffusers v1.0.0. We "
"Flax classes are deprecated and will be removed in Diffusers v0.40.0. We "
"recommend migrating to PyTorch classes or pinning your version of Diffusers."
)

Expand Down Expand Up @@ -381,7 +381,7 @@ class FlaxUNetMidBlock2DCrossAttn(nn.Module):

def setup(self):
logger.warning(
"Flax classes are deprecated and will be removed in Diffusers v1.0.0. We "
"Flax classes are deprecated and will be removed in Diffusers v0.40.0. We "
"recommend migrating to PyTorch classes or pinning your version of Diffusers."
)

Expand Down
2 changes: 1 addition & 1 deletion src/diffusers/models/unets/unet_2d_condition_flax.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -166,7 +166,7 @@ def init_weights(self, rng: jax.Array) -> FrozenDict:

def setup(self) -> None:
logger.warning(
"Flax classes are deprecated and will be removed in Diffusers v1.0.0. We "
"Flax classes are deprecated and will be removed in Diffusers v0.40.0. We "
"recommend migrating to PyTorch classes or pinning your version of Diffusers."
)

Expand Down
20 changes: 10 additions & 10 deletions src/diffusers/models/vae_flax.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -76,7 +76,7 @@ class FlaxUpsample2D(nn.Module):

def setup(self):
logger.warning(
"Flax classes are deprecated and will be removed in Diffusers v1.0.0. We "
"Flax classes are deprecated and will be removed in Diffusers v0.40.0. We "
"recommend migrating to PyTorch classes or pinning your version of Diffusers."
)
self.conv = nn.Conv(
Expand Down Expand Up @@ -114,7 +114,7 @@ class FlaxDownsample2D(nn.Module):

def setup(self):
logger.warning(
"Flax classes are deprecated and will be removed in Diffusers v1.0.0. We "
"Flax classes are deprecated and will be removed in Diffusers v0.40.0. We "
"recommend migrating to PyTorch classes or pinning your version of Diffusers."
)

Expand Down Expand Up @@ -161,7 +161,7 @@ class FlaxResnetBlock2D(nn.Module):

def setup(self):
logger.warning(
"Flax classes are deprecated and will be removed in Diffusers v1.0.0. We "
"Flax classes are deprecated and will be removed in Diffusers v0.40.0. We "
"recommend migrating to PyTorch classes or pinning your version of Diffusers."
)

Expand Down Expand Up @@ -238,7 +238,7 @@ class FlaxAttentionBlock(nn.Module):

def setup(self):
logger.warning(
"Flax classes are deprecated and will be removed in Diffusers v1.0.0. We "
"Flax classes are deprecated and will be removed in Diffusers v0.40.0. We "
"recommend migrating to PyTorch classes or pinning your version of Diffusers."
)

Expand Down Expand Up @@ -324,7 +324,7 @@ class FlaxDownEncoderBlock2D(nn.Module):

def setup(self):
logger.warning(
"Flax classes are deprecated and will be removed in Diffusers v1.0.0. We "
"Flax classes are deprecated and will be removed in Diffusers v0.40.0. We "
"recommend migrating to PyTorch classes or pinning your version of Diffusers."
)

Expand Down Expand Up @@ -386,7 +386,7 @@ class FlaxUpDecoderBlock2D(nn.Module):

def setup(self):
logger.warning(
"Flax classes are deprecated and will be removed in Diffusers v1.0.0. We "
"Flax classes are deprecated and will be removed in Diffusers v0.40.0. We "
"recommend migrating to PyTorch classes or pinning your version of Diffusers."
)

Expand Down Expand Up @@ -445,7 +445,7 @@ class FlaxUNetMidBlock2D(nn.Module):

def setup(self):
logger.warning(
"Flax classes are deprecated and will be removed in Diffusers v1.0.0. We "
"Flax classes are deprecated and will be removed in Diffusers v0.40.0. We "
"recommend migrating to PyTorch classes or pinning your version of Diffusers."
)

Expand Down Expand Up @@ -541,7 +541,7 @@ class FlaxEncoder(nn.Module):

def setup(self):
logger.warning(
"Flax classes are deprecated and will be removed in Diffusers v1.0.0. We "
"Flax classes are deprecated and will be removed in Diffusers v0.40.0. We "
"recommend migrating to PyTorch classes or pinning your version of Diffusers."
)

Expand Down Expand Up @@ -658,7 +658,7 @@ class FlaxDecoder(nn.Module):

def setup(self):
logger.warning(
"Flax classes are deprecated and will be removed in Diffusers v1.0.0. We "
"Flax classes are deprecated and will be removed in Diffusers v0.40.0. We "
"recommend migrating to PyTorch classes or pinning your version of Diffusers."
)

Expand Down Expand Up @@ -835,7 +835,7 @@ class FlaxAutoencoderKL(nn.Module, FlaxModelMixin, ConfigMixin):

def setup(self):
logger.warning(
"Flax classes are deprecated and will be removed in Diffusers v1.0.0. We "
"Flax classes are deprecated and will be removed in Diffusers v0.40.0. We "
"recommend migrating to PyTorch classes or pinning your version of Diffusers."
)

Expand Down
2 changes: 1 addition & 1 deletion src/diffusers/pipelines/pipeline_flax_utils.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -309,7 +309,7 @@ def from_pretrained(cls, pretrained_model_name_or_path: str | os.PathLike | None
```
"""
logger.warning(
"Flax classes are deprecated and will be removed in Diffusers v1.0.0. We "
"Flax classes are deprecated and will be removed in Diffusers v0.40.0. We "
"recommend migrating to PyTorch classes or pinning your version of Diffusers."
)

Expand Down
2 changes: 1 addition & 1 deletion src/diffusers/schedulers/scheduling_ddim_flax.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -129,7 +129,7 @@ def __init__(
dtype: jnp.dtype = jnp.float32,
):
logger.warning(
"Flax classes are deprecated and will be removed in Diffusers v1.0.0. We "
"Flax classes are deprecated and will be removed in Diffusers v0.40.0. We "
"recommend migrating to PyTorch classes or pinning your version of Diffusers."
)
self.dtype = dtype
Expand Down
2 changes: 1 addition & 1 deletion src/diffusers/schedulers/scheduling_ddpm_flax.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -114,7 +114,7 @@ def __init__(
dtype: jnp.dtype = jnp.float32,
):
logger.warning(
"Flax classes are deprecated and will be removed in Diffusers v1.0.0. We "
"Flax classes are deprecated and will be removed in Diffusers v0.40.0. We "
"recommend migrating to PyTorch classes or pinning your version of Diffusers."
)
self.dtype = dtype
Expand Down
2 changes: 1 addition & 1 deletion src/diffusers/schedulers/scheduling_dpmsolver_multistep_flax.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -175,7 +175,7 @@ def __init__(
dtype: jnp.dtype = jnp.float32,
):
logger.warning(
"Flax classes are deprecated and will be removed in Diffusers v1.0.0. We "
"Flax classes are deprecated and will be removed in Diffusers v0.40.0. We "
"recommend migrating to PyTorch classes or pinning your version of Diffusers."
)
self.dtype = dtype
Expand Down
2 changes: 1 addition & 1 deletion src/diffusers/schedulers/scheduling_euler_discrete_flax.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -112,7 +112,7 @@ def __init__(
dtype: jnp.dtype = jnp.float32,
):
logger.warning(
"Flax classes are deprecated and will be removed in Diffusers v1.0.0. We "
"Flax classes are deprecated and will be removed in Diffusers v0.40.0. We "
"recommend migrating to PyTorch classes or pinning your version of Diffusers."
)
self.dtype = dtype
Expand Down
2 changes: 1 addition & 1 deletion src/diffusers/schedulers/scheduling_karras_ve_flax.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -105,7 +105,7 @@ def __init__(
s_max: float = 50,
):
logger.warning(
"Flax classes are deprecated and will be removed in Diffusers v1.0.0. We "
"Flax classes are deprecated and will be removed in Diffusers v0.40.0. We "
"recommend migrating to PyTorch classes or pinning your version of Diffusers."
)

Expand Down
2 changes: 1 addition & 1 deletion src/diffusers/schedulers/scheduling_lms_discrete_flax.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -114,7 +114,7 @@ def __init__(
dtype: jnp.dtype = jnp.float32,
):
logger.warning(
"Flax classes are deprecated and will be removed in Diffusers v1.0.0. We "
"Flax classes are deprecated and will be removed in Diffusers v0.40.0. We "
"recommend migrating to PyTorch classes or pinning your version of Diffusers."
)
self.dtype = dtype
Expand Down
2 changes: 1 addition & 1 deletion src/diffusers/schedulers/scheduling_pndm_flax.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -135,7 +135,7 @@ def __init__(
dtype: jnp.dtype = jnp.float32,
):
logger.warning(
"Flax classes are deprecated and will be removed in Diffusers v1.0.0. We "
"Flax classes are deprecated and will be removed in Diffusers v0.40.0. We "
"recommend migrating to PyTorch classes or pinning your version of Diffusers."
)
self.dtype = dtype
Expand Down
2 changes: 1 addition & 1 deletion src/diffusers/schedulers/scheduling_sde_ve_flax.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -103,7 +103,7 @@ def __init__(
correct_steps: int = 1,
):
logger.warning(
"Flax classes are deprecated and will be removed in Diffusers v1.0.0. We "
"Flax classes are deprecated and will be removed in Diffusers v0.40.0. We "
"recommend migrating to PyTorch classes or pinning your version of Diffusers."
)

Expand Down
2 changes: 1 addition & 1 deletion src/diffusers/schedulers/scheduling_utils_flax.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -128,7 +128,7 @@ def from_pretrained(

"""
logger.warning(
"Flax classes are deprecated and will be removed in Diffusers v1.0.0. We "
"Flax classes are deprecated and will be removed in Diffusers v0.40.0. We "
"recommend migrating to PyTorch classes or pinning your version of Diffusers."
)
config, kwargs = cls.load_config(
Expand Down
Loading