Update README.md

README.md

@@ -6,31 +6,40 @@ tags:
 - diffusion
 library_name: diffuse
 ---
+---
 
-# mnist
+## Mnist Generation
+Flow matching diffusion model trained for mnist generation.
+Use with [**diffuse**](https://github.com/jcopo/diffuse), a JAX/Flax sampling library.
+Light enough to run on CPU
 
-Diffusion model for use with [diffuse](https://github.com/jcopo/diffuse), a JAX/Flax sampling library.
+---
 
 ## Model Details
+* **Framework:** JAX/Flax (NNX)
+* **Format:** msgpack
+* **Prediction Type:** Velocity (Flow Matching)
 
-- **Model Type**: CondUNet2D
-- **Checkpoint Step**: 45,000
-- **Precision**: float32
-- **Framework**: JAX/Flax (NNX)
-- **Format**: msgpack
+---
 
 ## Usage
 
+### Download and Load Model
+
 ```python
+import os
+
+import jax
+import jax.numpy as jnp
 from flax import nnx, serialization
 from huggingface_hub import hf_hub_download
-import importlib.util
 
 # Download model weights and config
-model_path = hf_hub_download(repo_id="jcopo/mnist", filename="model.msgpack")
-config_path = hf_hub_download(repo_id="jcopo/mnist", filename="config.py")
+model_path = hf_hub_download(repo_id="{jcopo/mnist}", filename="model.msgpack")
+config_path = hf_hub_download(repo_id="{jcopo/mnist}", filename="config.py")
 
 # Load config to get model architecture
+import importlib.util
 spec = importlib.util.spec_from_file_location("model_config", config_path)
 config_module = importlib.util.module_from_spec(spec)
 spec.loader.exec_module(config_module)
@@ -43,13 +52,10 @@ with open(model_path, "rb") as f:
     state_dict = serialization.from_bytes(None, f.read())
 
 # Restore weights into model
-nnx.update(model, state_dict)
+graphdef, state = nnx.split(model)
+state.replace_by_pure_dict(state_dict)
+model = nnx.merge(graphdef, state)
 model.eval()  # Set to evaluation mode
 
-# Now use the model for inference
-# output = model(input_data)
+print("✅ Model loaded successfully!")
 ```
-
-## Training Configuration
-
-This model was trained with the Triax framework using the configuration saved in the checkpoint.