Deprecate Flax support (#12151) · blizmax/diffusers@532f41c · GitHub
Skip to content

Commit 532f41c

Browse files
sayakpaulDN6
andauthored
Deprecate Flax support (huggingface#12151)
* start removing flax stuff. * add deprecation warning. * add warning messages. * more warnings. * remove dockerfiles. * remove more. * Update src/diffusers/models/attention_flax.py Co-authored-by: Dhruv Nair <dhruv.nair@gmail.com> * up --------- Co-authored-by: Dhruv Nair <dhruv.nair@gmail.com>
1 parent 5fcd5f5 commit 532f41c

21 files changed

Lines changed: 186 additions & 1848 deletions

‎.github/workflows/pr_flax_dependency_test.yml‎

Lines changed: 0 additions & 38 deletions
This file was deleted.

‎docker/diffusers-flax-cpu/Dockerfile‎

Lines changed: 0 additions & 49 deletions
This file was deleted.

‎docker/diffusers-flax-tpu/Dockerfile‎

Lines changed: 0 additions & 51 deletions
This file was deleted.

‎src/diffusers/models/attention_flax.py‎

Lines changed: 30 additions & 0 deletions

‎src/diffusers/models/controlnets/controlnet_flax.py‎

Lines changed: 14 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -20,7 +20,7 @@
2020
from flax.core.frozen_dict import FrozenDict
2121

2222
from ...configuration_utils import ConfigMixin, flax_register_to_config
23-
from ...utils import BaseOutput
23+
from ...utils import BaseOutput, logging
2424
from ..embeddings_flax import FlaxTimestepEmbedding, FlaxTimesteps
2525
from ..modeling_flax_utils import FlaxModelMixin
2626
from ..unets.unet_2d_blocks_flax import (
@@ -30,6 +30,9 @@
3030
)
3131

3232

33+
logger = logging.get_logger(__name__)
34+
35+
3336
@flax.struct.dataclass
3437
class FlaxControlNetOutput(BaseOutput):
3538
"""
@@ -50,6 +53,11 @@ class FlaxControlNetConditioningEmbedding(nn.Module):
5053
dtype: jnp.dtype = jnp.float32
5154

5255
def setup(self) -> None:
56+
logger.warning(
57+
"Flax classes are deprecated and will be removed in Diffusers v1.0.0. We "
58+
"recommend migrating to PyTorch classes or pinning your version of Diffusers."
59+
)
60+
5361
self.conv_in = nn.Conv(
5462
self.block_out_channels[0],
5563
kernel_size=(3, 3),
@@ -184,6 +192,11 @@ def init_weights(self, rng: jax.Array) -> FrozenDict:
184192
return self.init(rngs, sample, timesteps, encoder_hidden_states, controlnet_cond)["params"]
185193

186194
def setup(self) -> None:
195+
logger.warning(
196+
"Flax classes are deprecated and will be removed in Diffusers v1.0.0. We "
197+
"recommend migrating to PyTorch classes or pinning your version of Diffusers."
198+
)
199+
187200
block_out_channels = self.block_out_channels
188201
time_embed_dim = block_out_channels[0] * 4
189202

‎src/diffusers/models/embeddings_flax.py‎

Lines changed: 15 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -16,6 +16,11 @@
1616
import flax.linen as nn
1717
import jax.numpy as jnp
1818

19+
from ..utils import logging
20+
21+
22+
logger = logging.get_logger(__name__)
23+
1924

2025
def get_sinusoidal_embeddings(
2126
timesteps: jnp.ndarray,
@@ -76,6 +81,11 @@ class FlaxTimestepEmbedding(nn.Module):
7681
The data type for the embedding parameters.
7782
"""
7883

84+
logger.warning(
85+
"Flax classes are deprecated and will be removed in Diffusers v1.0.0. We "
86+
"recommend migrating to PyTorch classes or pinning your version of Diffusers."
87+
)
88+
7989
time_embed_dim: int = 32
8090
dtype: jnp.dtype = jnp.float32
8191

@@ -104,6 +114,11 @@ class FlaxTimesteps(nn.Module):
104114
flip_sin_to_cos: bool = False
105115
freq_shift: float = 1
106116

117+
logger.warning(
118+
"Flax classes are deprecated and will be removed in Diffusers v1.0.0. We "
119+
"recommend migrating to PyTorch classes or pinning your version of Diffusers."
120+
)
121+
107122
@nn.compact
108123
def __call__(self, timesteps):
109124
return get_sinusoidal_embeddings(

‎src/diffusers/models/modeling_flax_utils.py‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -290,6 +290,10 @@ def from_pretrained(
290290
You should probably TRAIN this model on a down-stream task to be able to use it for predictions and inference.
291291
```
292292
"""
293+
logger.warning(
294+
"Flax classes are deprecated and will be removed in Diffusers v1.0.0. We "
295+
"recommend migrating to PyTorch classes or pinning your version of Diffusers."
296+
)
293297
config = kwargs.pop("config", None)
294298
cache_dir = kwargs.pop("cache_dir", None)
295299
force_download = kwargs.pop("force_download", False)

‎src/diffusers/models/resnet_flax.py‎

Lines changed: 20 additions & 0 deletions

0 commit comments

Comments
 (0)