Hugging Face's logo

KRadim
/
conditional-medium-unet-tiny-imagenet

Text-to-Image
JAX
English
diffusion-model
flax
flax-nnx
unet
tiny-imagenet
conditional-image-generation
from-scratch
Model card Files
xet
 Community
KRadim commited on
Commit
a9df365
·
verified ·
1 Parent(s): 861810c

Upload Conditional MediumUNet model (Orbax checkpoint + native MLIR + compiled StableHLO artifact)

Browse files
Files changed (22) hide show
  1. .gitattributes +4 -0
  2. README.md +2 -2
  3. conditional_medium_unet.mlir +0 -0
  4. conditional_medium_unet.stablehlo +1 -1
  5. jax_checkpoint/_CHECKPOINT_METADATA +1 -1
  6. jax_checkpoint/d/a993e20199472f7b7522e312335e77ae +0 -0
  7. jax_checkpoint/manifest.ocdbt +0 -0
  8. jax_checkpoint/ocdbt.process_0/d/4ab5ac241e5c8af33cf9b4ad61c51d4a +0 -0
  9. jax_checkpoint/ocdbt.process_0/d/71b8db327cb48547f11be67cb9653387 +0 -0
  10. jax_checkpoint/ocdbt.process_0/d/78f4a2af1b845f802755613edacfe788 +3 -0
  11. jax_checkpoint/ocdbt.process_0/d/7fb353f027d4e2bb19bf4f349a12ad2d +3 -0
  12. jax_checkpoint/ocdbt.process_0/d/ec59e563ec9537a187e576ba8738e931 +0 -0
  13. jax_checkpoint/ocdbt.process_0/manifest.ocdbt +0 -0
  14. model_checkpoint/_CHECKPOINT_METADATA +1 -1
  15. model_checkpoint/d/a993e20199472f7b7522e312335e77ae +0 -0
  16. model_checkpoint/manifest.ocdbt +0 -0
  17. model_checkpoint/ocdbt.process_0/d/4ab5ac241e5c8af33cf9b4ad61c51d4a +0 -0
  18. model_checkpoint/ocdbt.process_0/d/71b8db327cb48547f11be67cb9653387 +0 -0
  19. model_checkpoint/ocdbt.process_0/d/78f4a2af1b845f802755613edacfe788 +3 -0
  20. model_checkpoint/ocdbt.process_0/d/7fb353f027d4e2bb19bf4f349a12ad2d +3 -0
  21. model_checkpoint/ocdbt.process_0/d/ec59e563ec9537a187e576ba8738e931 +0 -0
  22. model_checkpoint/ocdbt.process_0/manifest.ocdbt +0 -0
.gitattributes CHANGED
@@ -48,3 +48,7 @@ jax_checkpoint/ocdbt.process_0/d/7a58b6daf29d428286ba3d936858b2ca filter=lfs dif
48
  jax_checkpoint/ocdbt.process_0/d/f2bc40d9bf782f380aea1a83a156299d filter=lfs diff=lfs merge=lfs -text
49
  model_checkpoint/ocdbt.process_0/d/7a58b6daf29d428286ba3d936858b2ca filter=lfs diff=lfs merge=lfs -text
50
  model_checkpoint/ocdbt.process_0/d/f2bc40d9bf782f380aea1a83a156299d filter=lfs diff=lfs merge=lfs -text
 
 
 
 
 
48
  jax_checkpoint/ocdbt.process_0/d/f2bc40d9bf782f380aea1a83a156299d filter=lfs diff=lfs merge=lfs -text
49
  model_checkpoint/ocdbt.process_0/d/7a58b6daf29d428286ba3d936858b2ca filter=lfs diff=lfs merge=lfs -text
50
  model_checkpoint/ocdbt.process_0/d/f2bc40d9bf782f380aea1a83a156299d filter=lfs diff=lfs merge=lfs -text
51
+ jax_checkpoint/ocdbt.process_0/d/78f4a2af1b845f802755613edacfe788 filter=lfs diff=lfs merge=lfs -text
52
+ jax_checkpoint/ocdbt.process_0/d/7fb353f027d4e2bb19bf4f349a12ad2d filter=lfs diff=lfs merge=lfs -text
53
+ model_checkpoint/ocdbt.process_0/d/78f4a2af1b845f802755613edacfe788 filter=lfs diff=lfs merge=lfs -text
54
+ model_checkpoint/ocdbt.process_0/d/7fb353f027d4e2bb19bf4f349a12ad2d filter=lfs diff=lfs merge=lfs -text
README.md CHANGED
@@ -18,7 +18,7 @@ tags:
18
 
19
  # Conditional MediumUNet Diffusion Model (Flax NNX, JAX)
20
 
21
- A class-conditional, from-scratch DDPM-style diffusion model trained on **Tiny ImageNet ($64 imes 64 imes 3$, 200 classes)**, built using **JAX** and **Flax NNX**.
22
  The model architecture (`ConditionalUNet`) predicts the clean image from a noisy input, conditioned simultaneously on a sinusoidally embedded continuous noise level and a class label embedding.
23
 
24
  ## Model Details
@@ -30,7 +30,7 @@ The model architecture (`ConditionalUNet`) predicts the clean image from a noisy
30
  - **Framework:** JAX / Flax NNX
31
  - **Training data:** Tiny ImageNet (200 classes) — full 100,000-image Tiny ImageNet training set
32
  - **Optimizer:** AdamW with warmup-cosine-decay LR schedule, gradient clipping (global norm 1.0)
33
- - **Epochs:** 120
34
  - **Batch size:** 256
35
 
36
  ## Exported Formats
 
18
 
19
  # Conditional MediumUNet Diffusion Model (Flax NNX, JAX)
20
 
21
+ A class-conditional, from-scratch DDPM-style diffusion model trained on **Tiny ImageNet (64 X 64 X 3, 200 classes)**, built using **JAX** and **Flax NNX**.
22
  The model architecture (`ConditionalUNet`) predicts the clean image from a noisy input, conditioned simultaneously on a sinusoidally embedded continuous noise level and a class label embedding.
23
 
24
  ## Model Details
 
30
  - **Framework:** JAX / Flax NNX
31
  - **Training data:** Tiny ImageNet (200 classes) — full 100,000-image Tiny ImageNet training set
32
  - **Optimizer:** AdamW with warmup-cosine-decay LR schedule, gradient clipping (global norm 1.0)
33
+ - **Epochs:** 160
34
  - **Batch size:** 256
35
 
36
  ## Exported Formats
conditional_medium_unet.mlir CHANGED
The diff for this file is too large to render. See raw diff
 
conditional_medium_unet.stablehlo CHANGED
@@ -1,3 +1,3 @@
1
  version https://git-lfs.github.com/spec/v1
2
- oid sha256:2b18f4c65ef7cc85ba1d07e71bb7b1abdf91fad83ae5bdeb18fb3c63bedaa6ed
3
  size 3320456
 
1
  version https://git-lfs.github.com/spec/v1
2
+ oid sha256:be37014bf8a47c77a6b5a4332872ddbc8dff320e87d947eba1f065a2a88b36d5
3
  size 3320456
jax_checkpoint/_CHECKPOINT_METADATA CHANGED
@@ -1 +1 @@
1
- {"item_handlers": "orbax.checkpoint._src.handlers.standard_checkpoint_handler.StandardCheckpointHandler", "metrics": {}, "performance_metrics": {}, "init_timestamp_nsecs": 1785668722549381192, "commit_timestamp_nsecs": 1785668722781355156, "custom_metadata": {}}
 
1
+ {"item_handlers": "orbax.checkpoint._src.handlers.standard_checkpoint_handler.StandardCheckpointHandler", "metrics": {}, "performance_metrics": {}, "init_timestamp_nsecs": 1786051103177585413, "commit_timestamp_nsecs": 1786051103860121302, "custom_metadata": {}}
jax_checkpoint/d/a993e20199472f7b7522e312335e77ae ADDED
Binary file (12.8 kB). View file
 
jax_checkpoint/manifest.ocdbt CHANGED
Binary files a/jax_checkpoint/manifest.ocdbt and b/jax_checkpoint/manifest.ocdbt differ
 
jax_checkpoint/ocdbt.process_0/d/4ab5ac241e5c8af33cf9b4ad61c51d4a ADDED
Binary file (48.5 kB). View file
 
jax_checkpoint/ocdbt.process_0/d/71b8db327cb48547f11be67cb9653387 ADDED
Binary file (195 Bytes). View file
 
jax_checkpoint/ocdbt.process_0/d/78f4a2af1b845f802755613edacfe788 ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:b225c8f1bb95e29bfbb520e55d53a3e57e0190dbff354b90dbf9b03b72387b7a
3
+ size 2513103
jax_checkpoint/ocdbt.process_0/d/7fb353f027d4e2bb19bf4f349a12ad2d ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:3098b45fa5d8be19a3cf693bd7dbeb0603e6b5c36b56faaadb64e732a19386e2
3
+ size 501131
jax_checkpoint/ocdbt.process_0/d/ec59e563ec9537a187e576ba8738e931 ADDED
Binary file (877 Bytes). View file
 
jax_checkpoint/ocdbt.process_0/manifest.ocdbt CHANGED
Binary files a/jax_checkpoint/ocdbt.process_0/manifest.ocdbt and b/jax_checkpoint/ocdbt.process_0/manifest.ocdbt differ
 
model_checkpoint/_CHECKPOINT_METADATA CHANGED
@@ -1 +1 @@
1
- {"item_handlers": "orbax.checkpoint._src.handlers.standard_checkpoint_handler.StandardCheckpointHandler", "metrics": {}, "performance_metrics": {}, "init_timestamp_nsecs": 1785668722549381192, "commit_timestamp_nsecs": 1785668722781355156, "custom_metadata": {}}
 
1
+ {"item_handlers": "orbax.checkpoint._src.handlers.standard_checkpoint_handler.StandardCheckpointHandler", "metrics": {}, "performance_metrics": {}, "init_timestamp_nsecs": 1786051103177585413, "commit_timestamp_nsecs": 1786051103860121302, "custom_metadata": {}}
model_checkpoint/d/a993e20199472f7b7522e312335e77ae ADDED
Binary file (12.8 kB). View file
 
model_checkpoint/manifest.ocdbt CHANGED
Binary files a/model_checkpoint/manifest.ocdbt and b/model_checkpoint/manifest.ocdbt differ
 
model_checkpoint/ocdbt.process_0/d/4ab5ac241e5c8af33cf9b4ad61c51d4a ADDED
Binary file (48.5 kB). View file
 
model_checkpoint/ocdbt.process_0/d/71b8db327cb48547f11be67cb9653387 ADDED
Binary file (195 Bytes). View file
 
model_checkpoint/ocdbt.process_0/d/78f4a2af1b845f802755613edacfe788 ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:b225c8f1bb95e29bfbb520e55d53a3e57e0190dbff354b90dbf9b03b72387b7a
3
+ size 2513103
model_checkpoint/ocdbt.process_0/d/7fb353f027d4e2bb19bf4f349a12ad2d ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:3098b45fa5d8be19a3cf693bd7dbeb0603e6b5c36b56faaadb64e732a19386e2
3
+ size 501131
model_checkpoint/ocdbt.process_0/d/ec59e563ec9537a187e576ba8738e931 ADDED
Binary file (877 Bytes). View file
 
model_checkpoint/ocdbt.process_0/manifest.ocdbt CHANGED
Binary files a/model_checkpoint/ocdbt.process_0/manifest.ocdbt and b/model_checkpoint/ocdbt.process_0/manifest.ocdbt differ