flaxdiff 0.1.35.5__py3-none-any.whl → 0.1.36__py3-none-any.whl

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (23) hide show
  1. flaxdiff/data/dataset_map.py +71 -0
  2. flaxdiff/data/datasets.py +169 -0
  3. flaxdiff/data/online_loader.py +69 -42
  4. flaxdiff/models/attention.py +1 -0
  5. flaxdiff/models/simple_unet.py +11 -11
  6. flaxdiff/models/simple_vit.py +1 -1
  7. flaxdiff/samplers/common.py +72 -20
  8. flaxdiff/samplers/ddim.py +5 -5
  9. flaxdiff/samplers/ddpm.py +5 -11
  10. flaxdiff/samplers/euler.py +7 -10
  11. flaxdiff/samplers/heun_sampler.py +3 -4
  12. flaxdiff/samplers/multistep_dpm.py +2 -3
  13. flaxdiff/samplers/rk4_sampler.py +9 -9
  14. flaxdiff/trainer/autoencoder_trainer.py +1 -1
  15. flaxdiff/trainer/diffusion_trainer.py +124 -32
  16. flaxdiff/trainer/simple_trainer.py +187 -91
  17. flaxdiff/trainer/video_diffusion_trainer.py +62 -0
  18. flaxdiff/utils.py +105 -2
  19. {flaxdiff-0.1.35.5.dist-info → flaxdiff-0.1.36.dist-info}/METADATA +11 -5
  20. flaxdiff-0.1.36.dist-info/RECORD +43 -0
  21. {flaxdiff-0.1.35.5.dist-info → flaxdiff-0.1.36.dist-info}/WHEEL +1 -1
  22. flaxdiff-0.1.35.5.dist-info/RECORD +0 -40
  23. {flaxdiff-0.1.35.5.dist-info → flaxdiff-0.1.36.dist-info}/top_level.txt +0 -0
flaxdiff-0.1.35.5.dist-info/RECORD DELETED
@@ -1,40 +0,0 @@
1
- flaxdiff/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
2
- flaxdiff/utils.py,sha256=B0GcHlzlVYDNEIdh2v5qmP4u0neIT-FqexNohuyuCvg,2452
3
- flaxdiff/data/__init__.py,sha256=PM3PkHihyohT5SHVYKc8vQ4IeVfGPpCktkSVwvqMjQ4,52
4
- flaxdiff/data/online_loader.py,sha256=DoHrMZCi5gMd9tmkCpZIUU9lGxvfYtuaz58943_lCRc,11315
5
- flaxdiff/models/__init__.py,sha256=FAivVYXxM2JrCFIXf-C3374RB2Hth25dBrzOeNFhH1U,26
6
- flaxdiff/models/attention.py,sha256=Oz0-F0jllo3cIqQQJwYbYMrkhVzZWf5bq1UC6_RU1r8,13234
7
- flaxdiff/models/common.py,sha256=hWsSs2BP2J-JN1s4qLRr-h-KYkcVyl2hOp1Wsm_L-h8,10994
8
- flaxdiff/models/favor_fastattn.py,sha256=79Ew1nqarsNLPzZaBSd1ILORzJr74CupYeqGiCQK5E4,27689
9
- flaxdiff/models/simple_unet.py,sha256=EmTmf3wLIVULGpL1cSmAPhvg7RW4U8Ff3E8Qrtt8RLY,11397
10
- flaxdiff/models/simple_vit.py,sha256=YQiK5AHMjNxH5DAF6Oy4kjRPqddrvUP4i-CuavASPPI,7930
11
- flaxdiff/models/autoencoder/__init__.py,sha256=qY-7MldZpsfkF-_T2LqlRK7VHbqfmosz0NmvzDlBkOk,78
12
- flaxdiff/models/autoencoder/autoencoder.py,sha256=27_hYl0yXAdH9Mx4Xu9J79mSNo-FEKr9SxhVaS3ffn4,591
13
- flaxdiff/models/autoencoder/diffusers.py,sha256=JHeFLCxiHhu-QHwhKiCuKsQJn4AZumquiuxgZkiYGQ0,3643
14
- flaxdiff/models/autoencoder/simple_autoenc.py,sha256=UXHPgDmwGTnv3Uts6Zj3p9R9nJXnEiEXbllgarwDfXM,805
15
- flaxdiff/predictors/__init__.py,sha256=SKkYYRF9Wfgk2zhtZw4vCXOdOeRlrm2Mk6cvuaEvAzc,4403
16
- flaxdiff/samplers/__init__.py,sha256=_S-9TwDeshrI0VmapV-J2hqjTByOa0-oOeUs_IdovjU,285
17
- flaxdiff/samplers/common.py,sha256=_an5h5Niz9Joz_-ppridLrGHpu8X0VVvhNGknPu6AUY,5272
18
- flaxdiff/samplers/ddim.py,sha256=XHMBX06S5hMTnKMaGh6fmq189pQcaGkA6fnX6YPbHP0,511
19
- flaxdiff/samplers/ddpm.py,sha256=d_58hfVShJHsRPQf5h1-4YKrD43-HGjWz9Vd8hltZBg,2627
20
- flaxdiff/samplers/euler.py,sha256=Epf7LBKUky7B8b-1ZyIlLWdRMgjmP08BQraGSKmr_3I,2726
21
- flaxdiff/samplers/heun_sampler.py,sha256=hhWnSM26OfOIFAcsuWYa1z-2QPjASuoYTop2byLWqzE,1388
22
- flaxdiff/samplers/multistep_dpm.py,sha256=ocmEq2sCvsULy6oTFaD5BhTU4c8VHsge4bdg6tfxW80,2724
23
- flaxdiff/samplers/rk4_sampler.py,sha256=BF-dMV1KauO-SYShqrCfm3U3V-1n4clqQXBeoG8RWQo,1728
24
- flaxdiff/schedulers/__init__.py,sha256=3id390WEfdf-MN-oLSPAhlRFIXrFWr6ioAHPAwURJyE,375
25
- flaxdiff/schedulers/common.py,sha256=b-W4iI-aqScpVE8VZbBpiYvAVI6rqDkUP-C_hEVBwCI,4151
26
- flaxdiff/schedulers/continuous.py,sha256=5c_niOA20fxJ5oJDi09FfayIRogBGwtfG0XThW2IUZk,334
27
- flaxdiff/schedulers/cosine.py,sha256=9ban0dFHLMm35wQvaBT4nCQwPGmzNsXwQ1xI0oppmJI,2005
28
- flaxdiff/schedulers/discrete.py,sha256=O54wH2HVu3olJA71NxgAXFW9cr6B6Gl-DR_uZeytpds,3319
29
- flaxdiff/schedulers/exp.py,sha256=cPTnUJpYdzJRRZqMLYQz0rRUCpEmaP2tXhRumLx94jA,605
30
- flaxdiff/schedulers/karras.py,sha256=4GN120kGwdxxU-h2mVdhBVy9IORkUMm_vvz3XjthBcI,3355
31
- flaxdiff/schedulers/linear.py,sha256=6003F5ISq1Wc0h6UAzY95MJgsDIKGMhBzbiVALpea0k,581
32
- flaxdiff/schedulers/sqrt.py,sha256=1F84ZgQPuoNMhe6yxGTR2G0h7dPOZtm4UDQOakbSsEU,445
33
- flaxdiff/trainer/__init__.py,sha256=T-vUVq4zHcMK6kpCsG4Gu8vn71q6lZD-lg-Ul7yKfEk,128
34
- flaxdiff/trainer/autoencoder_trainer.py,sha256=al7AsZ7yeDMEiDD-gbcXf0ADq_xfk1VMxvg24GfA-XQ,7008
35
- flaxdiff/trainer/diffusion_trainer.py,sha256=wKkg63DWZjx2MoM3VQNCDIr40rWN8fUGxH9jWWxfZao,9373
36
- flaxdiff/trainer/simple_trainer.py,sha256=cawm6fZNQoLLATMneAU2gQ9j7kefqHnBPHuaIj3i_a4,18237
37
- flaxdiff-0.1.35.5.dist-info/METADATA,sha256=J8ikdMf9LIzdAuk6E3E77sI8NPegdgaKf5mrDHhpdXc,22085
38
- flaxdiff-0.1.35.5.dist-info/WHEEL,sha256=GV9aMThwP_4oNCtvEC2ec3qUYutgWeAzklro_0m4WJQ,91
39
- flaxdiff-0.1.35.5.dist-info/top_level.txt,sha256=-2-nXnfkJgSfkki1tjm5Faw6Dso7vhtdn2szwCdX5CQ,9
40
- flaxdiff-0.1.35.5.dist-info/RECORD,,