flaxdiff 0.1.37.1__py3-none-any.whl → 0.1.37.3__py3-none-any.whl

flaxdiff/samplers/__init__.py

@@ -1,7 +1,7 @@
-from .common import *
-from .ddim import *
-from .ddpm import *
-from .euler import *
-from .heun_sampler import *
-from .rk4_sampler import *
-from .multistep_dpm import *
+from .common import DiffusionSampler
+from .ddim import DDIMSampler
+from .ddpm import DDPMSampler, SimpleDDPMSampler
+from .euler import EulerSampler, SimplifiedEulerSampler, EulerAncestralSampler
+from .heun_sampler import HeunSampler
+from .rk4_sampler import RK4Sampler
+from .multistep_dpm import MultiStepDPM

flaxdiff/samplers/common.py

@@ -2,20 +2,15 @@ from flax import linen as nn
 import jax
 import jax.numpy as jnp
 import tqdm
-from typing import Union
+from typing import Union, Type
 from ..schedulers import NoiseScheduler
 from ..utils import RandomMarkovState, MarkovState, clip_images
 from ..predictors import DiffusionPredictionTransform, EpsilonPredictionTransform
 
 class DiffusionSampler():
-    model:nn.Module
-    noise_schedule:NoiseScheduler
-    params:dict
-    model_output_transform:DiffusionPredictionTransform
-
     def __init__(self, model:nn.Module, params:dict,  
                  noise_schedule:NoiseScheduler, 
-                 model_output_transform:DiffusionPredictionTransform=EpsilonPredictionTransform(),
+                 model_output_transform:DiffusionPredictionTransform,
                  guidance_scale:float = 0.0,
                  null_labels_seq:jax.Array=None,
                  autoencoder=None,
@@ -122,9 +117,11 @@ class DiffusionSampler():
                         end_step:int = 0,
                         steps_override=None,
                         priors=None, 
-                        rngstate:RandomMarkovState=RandomMarkovState(jax.random.PRNGKey(42)),
+                        rngstate:RandomMarkovState=None,
                         model_conditioning_inputs:tuple=()
                         ) -> jnp.ndarray:
+        if rngstate is None:
+            rngstate = RandomMarkovState(jax.random.PRNGKey(42))
         if priors is None:
             rngstate, newrngs = rngstate.get_random_key()
             samples = self.get_initial_samples(num_images, newrngs, start_step)
@@ -169,4 +166,5 @@ class DiffusionSampler():
         if self.autoencoder is not None:
             samples = self.autoencoder.decode(samples)
         samples = clip_images(samples)
-        return samples
+        return samples
+    

{flaxdiff-0.1.37.1.dist-info → flaxdiff-0.1.37.3.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: flaxdiff
-Version: 0.1.37.1
+Version: 0.1.37.3
 Summary: A versatile and easy to understand Diffusion library
 Author-email: Ashish Kumar Singh <ashishkmr472@gmail.com>
 License-Expression: MIT

{flaxdiff-0.1.37.1.dist-info → flaxdiff-0.1.37.3.dist-info}/RECORD

@@ -7,6 +7,8 @@ flaxdiff/data/online_loader.py,sha256=1Fi_QRixxRzbt602nORINcDeHEccvCrBpagrz4PURY
 flaxdiff/data/sources/gcs.py,sha256=11ZuQhvMyJRLg21DgVdzO5qEuae7zgzTXGNOskF-cbs,3380
 flaxdiff/data/sources/tfds.py,sha256=7n-uobG_UvkD5mU_1ovPd9kb6xJrbEKFFXdVEHDunts,2781
 flaxdiff/metrics/inception.py,sha256=a5kjMCPMT9gB88c_HCKiek-2vsAyoE35K7nDt4h4pVI,31843
+flaxdiff/metrics/psnr.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+flaxdiff/metrics/ssim.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 flaxdiff/metrics/utils.py,sha256=YuuOfqvqgIjsceupwNeJ59vQ2TnGeNMIyKdkIqOmoNg,1702
 flaxdiff/models/__init__.py,sha256=FAivVYXxM2JrCFIXf-C3374RB2Hth25dBrzOeNFhH1U,26
 flaxdiff/models/attention.py,sha256=JvrP7-09MV6IfRLRBhqjPmNUU-lkEMk9TOnJSBKcar8,13289
@@ -19,8 +21,8 @@ flaxdiff/models/autoencoder/autoencoder.py,sha256=27_hYl0yXAdH9Mx4Xu9J79mSNo-FEK
 flaxdiff/models/autoencoder/diffusers.py,sha256=JHeFLCxiHhu-QHwhKiCuKsQJn4AZumquiuxgZkiYGQ0,3643
 flaxdiff/models/autoencoder/simple_autoenc.py,sha256=UXHPgDmwGTnv3Uts6Zj3p9R9nJXnEiEXbllgarwDfXM,805
 flaxdiff/predictors/__init__.py,sha256=SKkYYRF9Wfgk2zhtZw4vCXOdOeRlrm2Mk6cvuaEvAzc,4403
-flaxdiff/samplers/__init__.py,sha256=dsK_a6HwMpxNuDbl5SzmvKJHYC_DGbTxxS3gfQYwcUY,166
-flaxdiff/samplers/common.py,sha256=8SabB-lgcbJZtNBfS96cn2xakeOj6-HC0Evhzeh2NlY,8921
+flaxdiff/samplers/__init__.py,sha256=EY9v1pgwEoR64Kiz9K8fAR-4_ir9c03mYeY3hrpUNhE,308
+flaxdiff/samplers/common.py,sha256=7gKNY4mWVnLjtcioGLFD_Vwmxg9zJovUb8EcYWlc_GE,8833
 flaxdiff/samplers/ddim.py,sha256=hTjDm0SmIj-Tkc80QRATMcN_sKVhHbqZQboRQCAn4mY,569
 flaxdiff/samplers/ddpm.py,sha256=JgkNSo7fp7Jm-8rCy4eu5m4YIzXTWzxv-iHf3EQ0z5w,2243
 flaxdiff/samplers/euler.py,sha256=QSkttB4DYnepDGwhWq3EGXYjMAqj4qLOdh7u98HttzY,2791
@@ -41,7 +43,7 @@ flaxdiff/trainer/autoencoder_trainer.py,sha256=hxihkRL9WCIQVGOP-pc1jjjIUaRXDLcNo
 flaxdiff/trainer/diffusion_trainer.py,sha256=KVeXJ9ZQKcvD-O_hCJnxro0dQRuQe5ZVGGMEL4Lgm9k,12814
 flaxdiff/trainer/simple_trainer.py,sha256=lmRo8N0bMupIyS3ejPvPtxoskY_3GLC8iyJE6u4TIWc,21990
 flaxdiff/trainer/video_diffusion_trainer.py,sha256=gMkKpnKNTo8QhTx5ptEEkc7W5-7rzXIr9queU53hXyQ,2197
-flaxdiff-0.1.37.1.dist-info/METADATA,sha256=ThP8nRYIdfvEE2ytSgTbulqbCDBLaWk22t4u2Mk-7-8,23985
-flaxdiff-0.1.37.1.dist-info/WHEEL,sha256=CmyFI0kx5cdEMTLiONQRbGQwjIoR1aIYB7eCAQ4KPJ0,91
-flaxdiff-0.1.37.1.dist-info/top_level.txt,sha256=-2-nXnfkJgSfkki1tjm5Faw6Dso7vhtdn2szwCdX5CQ,9
-flaxdiff-0.1.37.1.dist-info/RECORD,,
+flaxdiff-0.1.37.3.dist-info/METADATA,sha256=7U7SINGO_ZzsUeuTPi2LTMqxrj93Cvglyh1Q7D39zRM,23985
+flaxdiff-0.1.37.3.dist-info/WHEEL,sha256=CmyFI0kx5cdEMTLiONQRbGQwjIoR1aIYB7eCAQ4KPJ0,91
+flaxdiff-0.1.37.3.dist-info/top_level.txt,sha256=-2-nXnfkJgSfkki1tjm5Faw6Dso7vhtdn2szwCdX5CQ,9
+flaxdiff-0.1.37.3.dist-info/RECORD,,