Upload model at step 45000

README.md

@@ -0,0 +1,54 @@
+---
+tags:
+- jax
+- flax
+- flax-nnx
+library_name: triax
+---
+
+# mnist
+
+Model trained using [Triax](https://github.com/your-org/triax), a JAX/Flax training framework.
+
+## Model Details
+
+- **Model Type**: CondUNet2D
+- **Training Step**: 45,000
+- **Precision**: float32
+- **Framework**: JAX/Flax (NNX)
+- **Format**: msgpack
+
+## Usage
+
+```python
+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")
+
+# Load config to get model architecture
+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)
+
+# Initialize model from config
+model = config_module.model
+
+# Load weights
+with open(model_path, "rb") as f:
+    state_dict = serialization.from_bytes(None, f.read())
+
+# Restore weights into model
+nnx.update(model, state_dict)
+model.eval()  # Set to evaluation mode
+
+# Now use the model for inference
+# output = model(input_data)
+```
+
+## Training Configuration
+
+This model was trained with the Triax framework using the configuration saved in the checkpoint.

config.py

@@ -0,0 +1,23 @@
+"""Model configuration for jcopo/mnist
+
+This file contains the model architecture definition.
+Training step: 45000
+Precision: float32
+"""
+
+from triax.models.nn.condUNet import CondUNet2D
+
+# Model architecture
+# TODO: Fill in the actual initialization parameters from your training config
+model = CondUNet2D(
+    # Add your model parameters here
+    # Example:
+    # hidden_dim=256,
+    # num_layers=4,
+    # etc.
+)
+
+# Metadata
+STEP = 45000
+PRECISION = "float32"
+MODEL_TYPE = "CondUNet2D"