Upload model at step 45000

README.md

@@ -3,17 +3,18 @@ tags:
 - jax
 - flax
 - flax-nnx
-library_name: triax
+- diffusion
+library_name: diffuse
 ---
 
 # mnist
 
-Model trained using [Triax](https://github.com/your-org/triax), a JAX/Flax training framework.
+Diffusion model for use with [diffuse](https://github.com/jcopo/diffuse), a JAX/Flax sampling library.
 
 ## Model Details
 
 - **Model Type**: CondUNet2D
-- **Training Step**: 45,000
+- **Checkpoint Step**: 45,000
 - **Precision**: float32
 - **Framework**: JAX/Flax (NNX)
 - **Format**: msgpack

config.py

@@ -7,7 +7,7 @@ Precision: float32
 import jax
 import jax.numpy as jnp
 from flax import nnx
-from triax.models import CondUNet2D
+from diffuse.models import CondUNet2D
 
 # Model definition
 model = CondUNet2D(
@@ -17,8 +17,8 @@ model = CondUNet2D(
         num_res_blocks=variant["num_res_blocks"],
         attention_resolutions=variant["attention_resolutions"],
         num_heads=variant["num_heads"],
-        dropout=DROPOUT,
-        dropout_rate=DROPOUT_RATE,
+        dropout=True,
+        dropout_rate=0.1,
         activation=nnx.swish,
         param_dtype=jnp.float32,
         dtype=compute_dtype,