kanshan/diffusers
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Check shape and remove deprecated APIs in scheduling_ddpm_flax.py (#7703)

Browse Source 
`model_output.shape` may only have rank 1.

There are warnings related to use of random keys.

```
tests/schedulers/test_scheduler_flax.py: 13 warnings
  /Users/phillypham/diffusers/src/diffusers/schedulers/scheduling_ddpm_flax.py:268: FutureWarning: normal accepts a single key, but was given a key array of shape (1, 2) != (). Use jax.vmap for batching. In a future JAX version, this will be an error.
    noise = jax.random.normal(split_key, shape=model_output.shape, dtype=self.dtype)

tests/schedulers/test_scheduler_flax.py::FlaxDDPMSchedulerTest::test_betas
  /Users/phillypham/virtualenv/diffusers/lib/python3.9/site-packages/jax/_src/random.py:731: FutureWarning: uniform accepts a single key, but was given a key array of shape (1,) != (). Use jax.vmap for batching. In a future JAX version, this will be an error.
    u = uniform(key, shape, dtype, lo, hi)  # type: ignore[arg-type]
```
This commit is contained in:
Philip Pham
2024-05-08 07:57:19 -04:00
committed by GitHub
parent d50baf0c63
commit f29b93488d
1 changed files with 7 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/diffusers/schedulers/scheduling_ddpm_flax.py
+7 -3
View File
@@ -222,9 +222,13 @@ class FlaxDDPMScheduler(FlaxSchedulerMixin, ConfigMixin):
t = timestep
if key is None:
key = jax.random.PRNGKey(0)
key = jax.random.key(0)
if model_output.shape[1] == sample.shape[1] * 2 and self.config.variance_type in ["learned", "learned_range"]:
if (
len(model_output.shape) > 1
and model_output.shape[1] == sample.shape[1] * 2
and self.config.variance_type in ["learned", "learned_range"]
):
model_output, predicted_variance = jnp.split(model_output, sample.shape[1], axis=1)
else:
predicted_variance = None
@@ -264,7 +268,7 @@ class FlaxDDPMScheduler(FlaxSchedulerMixin, ConfigMixin):
# 6. Add noise
def random_variance():
split_key = jax.random.split(key, num=1)
split_key = jax.random.split(key, num=1)[0]
noise = jax.random.normal(split_key, shape=model_output.shape, dtype=self.dtype)
return (self._get_variance(state, t, predicted_variance=predicted_variance) ** 0.5) * noise