flaxdiff 0.1.37.2__tar.gz → 0.1.37.3__tar.gz

{flaxdiff-0.1.37.2 → flaxdiff-0.1.37.3}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: flaxdiff
-Version: 0.1.37.2
+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.2 → flaxdiff-0.1.37.3}/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.2 → flaxdiff-0.1.37.3}/flaxdiff.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: flaxdiff
-Version: 0.1.37.2
+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.2 → flaxdiff-0.1.37.3}/flaxdiff.egg-info/SOURCES.txt

@@ -14,6 +14,8 @@ flaxdiff/data/online_loader.py
 flaxdiff/data/sources/gcs.py
 flaxdiff/data/sources/tfds.py
 flaxdiff/metrics/inception.py
+flaxdiff/metrics/psnr.py
+flaxdiff/metrics/ssim.py
 flaxdiff/metrics/utils.py
 flaxdiff/models/__init__.py
 flaxdiff/models/attention.py

{flaxdiff-0.1.37.2 → flaxdiff-0.1.37.3}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "flaxdiff"
-version = "0.1.37.2"
+version = "0.1.37.3"
 description = "A versatile and easy to understand Diffusion library"
 readme = "README.md"
 authors = [