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

flaxdiff/data/dataset_map.py

@@ -0,0 +1,71 @@
+from .sources.tfds import data_source_tfds, tfds_augmenters
+from .sources.gcs import data_source_gcs, data_source_combined_gcs, gcs_augmenters
+
+# Configure the following for your datasets
+datasetMap = {
+    "oxford_flowers102": {
+        "source": data_source_tfds("oxford_flowers102", use_tf=False),
+        "augmenter": tfds_augmenters,
+    },
+    "cc12m": {
+        "source": data_source_gcs('arrayrecord2/cc12m'),
+        "augmenter": gcs_augmenters,
+    },
+    "laiona_coco": {
+        "source": data_source_gcs('arrayrecord2/laion-aesthetics-12m+mscoco-2017'),
+        "augmenter": gcs_augmenters,
+    },
+    "aesthetic_coyo": {
+        "source": data_source_gcs('arrayrecords/aestheticCoyo_0.25clip_6aesthetic'),
+        "augmenter": gcs_augmenters,
+    },
+    "combined_aesthetic": {
+        "source": data_source_combined_gcs([
+                'arrayrecord2/laion-aesthetics-12m+mscoco-2017',
+                'arrayrecords/aestheticCoyo_0.25clip_6aesthetic',
+                'arrayrecord2/cc12m',
+                'arrayrecords/aestheticCoyo_0.25clip_6aesthetic',
+            ]),
+        "augmenter": gcs_augmenters,
+    },
+    "laiona_coco_coyo": {
+        "source": data_source_combined_gcs([
+                'arrayrecords/aestheticCoyo_0.25clip_6aesthetic',
+                'arrayrecord2/laion-aesthetics-12m+mscoco-2017',
+                'arrayrecords/aestheticCoyo_0.25clip_6aesthetic',
+            ]),
+        "augmenter": gcs_augmenters,
+    },
+    "combined_30m": {
+        "source": data_source_combined_gcs([
+                'arrayrecord2/laion-aesthetics-12m+mscoco-2017',
+                'arrayrecord2/cc12m',
+                'arrayrecord2/aestheticCoyo_0.26_clip_5.5aesthetic_256plus',
+                "arrayrecord2/playground+leonardo_x4+cc3m.parquet",
+            ]),
+        "augmenter": gcs_augmenters,
+    }
+}
+
+onlineDatasetMap = {
+    "combined_online": {
+        "source": [
+            # "gs://flaxdiff-datasets-regional/datasets/laion-aesthetics-12m+mscoco-2017.parquet"
+            # "ChristophSchuhmann/MS_COCO_2017_URL_TEXT",
+            # "dclure/laion-aesthetics-12m-umap",
+            "gs://flaxdiff-datasets-regional/datasets/laion-aesthetics-12m+mscoco-2017",
+            "gs://flaxdiff-datasets-regional/datasets/coyo700m-aesthetic-5.4_25M",
+            "gs://flaxdiff-datasets-regional/datasets/leonardo-liked-1.8m",
+            "gs://flaxdiff-datasets-regional/datasets/leonardo-liked-1.8m",
+            "gs://flaxdiff-datasets-regional/datasets/leonardo-liked-1.8m",
+            "gs://flaxdiff-datasets-regional/datasets/cc12m",
+            "gs://flaxdiff-datasets-regional/datasets/playground-liked",
+            "gs://flaxdiff-datasets-regional/datasets/leonardo-liked-1.8m",
+            "gs://flaxdiff-datasets-regional/datasets/leonardo-liked-1.8m",
+            "gs://flaxdiff-datasets-regional/datasets/cc3m",
+            "gs://flaxdiff-datasets-regional/datasets/cc3m",
+            "gs://flaxdiff-datasets-regional/datasets/laion2B-en-aesthetic-4.2_37M",
+            # "gs://flaxdiff-datasets-regional/datasets/laiion400m-185M"
+        ]
+    }
+}

flaxdiff/data/datasets.py

@@ -0,0 +1,169 @@
+import jax.numpy as jnp
+import grain.python as pygrain
+from typing import Dict
+import numpy as np
+import jax
+from flaxdiff.utils import convert_to_global_tree, AutoTextTokenizer
+from .dataset_map import datasetMap, onlineDatasetMap
+import traceback
+from .online_loader import OnlineStreamingDataLoader
+import queue
+from jax.sharding import Mesh
+import threading
+
+def batch_mesh_map(mesh):
+    class augmenters(pygrain.MapTransform):
+        def __init__(self, *args, **kwargs):
+            super().__init__(*args, **kwargs)
+
+        def map(self, batch) -> Dict[str, jnp.array]:
+            return convert_to_global_tree(mesh, batch)
+    return augmenters
+
+def get_dataset_grain(
+    data_name="cc12m",
+    batch_size=64,
+    image_scale=256,
+    count=None,
+    num_epochs=None,
+    method=jax.image.ResizeMethod.LANCZOS3,
+    worker_count=32,
+    read_thread_count=64,
+    read_buffer_size=50,
+    worker_buffer_size=20,
+    seed=0,
+    dataset_source="/mnt/gcs_mount/arrayrecord2/cc12m/",
+):
+    dataset = datasetMap[data_name]
+    data_source = dataset["source"](dataset_source)
+    augmenter = dataset["augmenter"](image_scale, method)
+
+    local_batch_size = batch_size // jax.process_count()
+
+    sampler = pygrain.IndexSampler(
+        num_records=len(data_source) if count is None else count,
+        shuffle=True,
+        seed=seed,
+        num_epochs=num_epochs,
+        shard_options=pygrain.ShardByJaxProcess(),
+    )
+
+    def get_trainset():
+        transformations = [
+            augmenter(),
+            pygrain.Batch(local_batch_size, drop_remainder=True),
+        ]
+        
+        # if mesh != None:
+        #     transformations += [batch_mesh_map(mesh)]
+
+        loader = pygrain.DataLoader(
+            data_source=data_source,
+            sampler=sampler,
+            operations=transformations,
+            worker_count=worker_count,
+            read_options=pygrain.ReadOptions(
+                read_thread_count, read_buffer_size
+            ),
+            worker_buffer_size=worker_buffer_size,
+        )
+        return loader
+    
+
+    return {
+        "train": get_trainset,
+        "train_len": len(data_source),
+        "local_batch_size": local_batch_size,
+        "global_batch_size": batch_size,
+        # "null_labels": null_labels,
+        # "null_labels_full": null_labels_full,
+        # "model": model,
+        # "tokenizer": tokenizer,
+    }
+
+def generate_collate_fn():
+    auto_tokenize = AutoTextTokenizer(tensor_type="np")
+    def default_collate(batch):
+        try:
+            # urls = [sample["url"] for sample in batch]
+            captions = [sample["caption"] for sample in batch]
+            results = auto_tokenize(captions)
+            images = np.stack([sample["image"] for sample in batch], axis=0)
+            return {
+                "image": images,
+                "input_ids": results['input_ids'],
+                "attention_mask": results['attention_mask'],
+            }
+        except Exception as e:
+            print("Error in collate function", e, [sample["image"].shape for sample in batch])
+            traceback.print_exc()
+            
+    return default_collate
+    
+def get_dataset_online(
+        data_name="combined_online",
+        batch_size=64,
+        image_scale=256,
+        count=None,
+        num_epochs=None,
+        method=jax.image.ResizeMethod.LANCZOS3,
+        worker_count=32,
+        read_thread_count=64,
+        read_buffer_size=50,
+        worker_buffer_size=20,
+        seed=0,
+        dataset_source="/mnt/gcs_mount/arrayrecord2/cc12m/",
+    ):
+    local_batch_size = batch_size // jax.process_count()
+    
+    sources = onlineDatasetMap[data_name]["source"]
+    dataloader = OnlineStreamingDataLoader(
+            sources, 
+            batch_size=local_batch_size,
+            num_workers=worker_count,
+            num_threads=read_thread_count,
+            image_shape=(image_scale, image_scale),
+            global_process_count=jax.process_count(),
+            global_process_index=jax.process_index(),
+            prefetch=worker_buffer_size,
+            collate_fn=generate_collate_fn(),
+            default_split="train",
+        )
+    
+    def get_trainset(mesh: Mesh = None):
+        if mesh != None:
+            class dataLoaderWithMesh:
+                def __init__(self, dataloader, mesh):
+                    self.dataloader = dataloader
+                    self.mesh = mesh
+                    self.tmp_queue = queue.Queue(worker_buffer_size)
+                    def batch_loader():
+                        for batch in self.dataloader:
+                            try:
+                                self.tmp_queue.put(convert_to_global_tree(mesh, batch))
+                            except Exception as e:
+                                print("Error processing batch", e)
+                    self.loader_thread = threading.Thread(target=batch_loader)
+                    self.loader_thread.start()
+                    
+                def __iter__(self):
+                    return self
+                
+                def __next__(self):
+                    return self.tmp_queue.get()
+                
+            dataloader_with_mesh = dataLoaderWithMesh(dataloader, mesh)
+                    
+            return dataloader_with_mesh
+        return dataloader
+    
+    return {
+        "train": get_trainset,
+        "train_len": len(dataloader) * jax.process_count(),
+        "local_batch_size": local_batch_size,
+        "global_batch_size": batch_size,
+        # "null_labels": null_labels,
+        # "null_labels_full": null_labels_full,
+        # "model": model,
+        # "tokenizer": tokenizer,
+    }

flaxdiff/data/online_loader.py

@@ -45,36 +45,43 @@ def fetch_single_image(image_url, timeout=None, retries=0):
 
 
 def default_image_processor(
-    image, image_shape,
+    image, image_shape, 
     min_image_shape=(128, 128),
     upscale_interpolation=cv2.INTER_CUBIC,
     downscale_interpolation=cv2.INTER_AREA,
 ):
-    image = np.array(image)
-    original_height, original_width = image.shape[:2]
-    # check if the image is too small
-    if min(original_height, original_width) < min(min_image_shape):
-        return None, original_height, original_width
-    # check if wrong aspect ratio
-    if max(original_height, original_width) / min(original_height, original_width) > 2.4:
-        return None, original_height, original_width
-    # check if the variance is too low
-    if np.std(image) < 1e-5:
-        return None, original_height, original_width
-    # image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
-    downscale = max(original_width, original_height) > max(image_shape)
-    interpolation = downscale_interpolation if downscale else upscale_interpolation
-
-    image = A.longest_max_size(image, max(
-        image_shape), interpolation=interpolation)
-    image = A.pad(
-        image,
-        min_height=image_shape[0],
-        min_width=image_shape[1],
-        border_mode=cv2.BORDER_CONSTANT,
-        value=[255, 255, 255],
-    )
-    return image, original_height, original_width
+    try:
+        image = np.array(image)
+        if len(image.shape) != 3 or image.shape[2] != 3:
+            return None, 0, 0
+        original_height, original_width = image.shape[:2]
+        # check if the image is too small
+        if min(original_height, original_width) < min(min_image_shape):
+            return None, original_height, original_width
+        # check if wrong aspect ratio
+        if max(original_height, original_width) / min(original_height, original_width) > 2.4:
+            return None, original_height, original_width
+        # check if the variance is too low
+        if np.std(image) < 1e-5:
+            return None, original_height, original_width
+        # image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
+        downscale = max(original_width, original_height) > max(image_shape)
+        interpolation = downscale_interpolation if downscale else upscale_interpolation
+
+        image = A.longest_max_size(image, max(
+            image_shape), interpolation=interpolation)
+        image = A.pad(
+            image,
+            min_height=image_shape[0],
+            min_width=image_shape[1],
+            border_mode=cv2.BORDER_CONSTANT,
+            value=[255, 255, 255],
+        )
+        return image, original_height, original_width
+    except Exception as e:
+        # print("Error processing image", e, image_shape, interpolation)
+        # traceback.print_exc()
+        return None, 0, 0
 
 
 def map_sample(
@@ -120,14 +127,36 @@ def map_sample(
         # })
         pass
 
-
 def default_feature_extractor(sample):
+    url = None
+    if "url" in sample:
+        url = sample["url"]
+    elif "URL" in sample:
+        url = sample["URL"]
+    elif "image_url" in sample:
+        url = sample["image_url"]
+    else:
+        print("No url found in sample, skipping", sample.keys())
+    
+    caption = None
+    if "caption" in sample:
+        caption = sample["caption"]
+    elif "CAPTION" in sample:
+        caption = sample["CAPTION"]
+    elif "txt" in sample:
+        caption = sample["txt"]
+    elif "TEXT" in sample:
+        caption = sample["TEXT"]
+    elif "text" in sample:
+        caption = sample["text"]
+    else:
+        print("No caption found in sample, skipping", sample.keys())
+        
     return {
-        "url": sample["url"],
-        "caption": sample["caption"],
+        "url": url,
+        "caption": caption,
     }
-
-
+    
 def map_batch(
     batch, num_threads=256, image_shape=(256, 256),
     min_image_shape=(128, 128),
@@ -301,15 +330,13 @@ class OnlineStreamingDataLoader():
         self.dataset = dataset.shard(
             num_shards=global_process_count, index=global_process_index)
         print(f"Dataset length: {len(dataset)}")
-        self.iterator = ImageBatchIterator(
-            self.dataset, image_shape=image_shape,
-            min_image_shape=min_image_shape,
-            num_workers=num_workers, batch_size=batch_size, num_threads=num_threads,
-            timeout=timeout, retries=retries, image_processor=image_processor,
-            upscale_interpolation=upscale_interpolation,
-            downscale_interpolation=downscale_interpolation,
-            feature_extractor=feature_extractor
-        )
+        self.iterator = ImageBatchIterator(self.dataset, image_shape=image_shape,
+                                           min_image_shape=min_image_shape,
+                                           num_workers=num_workers, batch_size=batch_size, num_threads=num_threads,
+                                            timeout=timeout, retries=retries, image_processor=image_processor,
+                                             upscale_interpolation=upscale_interpolation,
+                                             downscale_interpolation=downscale_interpolation,
+                                             feature_extractor=feature_extractor)
         self.batch_size = batch_size
 
         # Launch a thread to load batches in the background
@@ -320,7 +347,7 @@ class OnlineStreamingDataLoader():
                 try:
                     self.batch_queue.put(collate_fn(batch))
                 except Exception as e:
-                    print("Error processing batch", e)
+                    print("Error collating batch", e)
 
         self.loader_thread = threading.Thread(target=batch_loader)
         self.loader_thread.start()
@@ -333,4 +360,4 @@ class OnlineStreamingDataLoader():
         # return self.collate_fn(next(self.iterator))
 
     def __len__(self):
-        return len(self.dataset)
+        return len(self.dataset)

flaxdiff/models/attention.py

@@ -11,6 +11,7 @@ import einops
 import functools
 import math
 from .common import kernel_init
+import jax.experimental.pallas.ops.tpu.flash_attention
 
 class EfficientAttention(nn.Module):
     """

flaxdiff/models/simple_unet.py

@@ -50,7 +50,7 @@ class Unet(nn.Module):
             features=self.feature_depths[0],
             kernel_size=(3, 3),
             strides=(1, 1),
-            kernel_init=self.kernel_init(1.0),
+            kernel_init=self.kernel_init(scale=1.0),
             dtype=self.dtype,
             precision=self.precision
         )(x)
@@ -65,7 +65,7 @@ class Unet(nn.Module):
                     down_conv_type,
                     name=f"down_{i}_residual_{j}",
                     features=dim_in,
-                    kernel_init=self.kernel_init(1.0),
+                    kernel_init=self.kernel_init(scale=1.0),
                     kernel_size=(3, 3),
                     strides=(1, 1),
                     activation=self.activation,
@@ -85,7 +85,7 @@ class Unet(nn.Module):
                                         force_fp32_for_softmax=attention_config.get("force_fp32_for_softmax", False),
                                         norm_inputs=attention_config.get("norm_inputs", True),
                                         explicitly_add_residual=attention_config.get("explicitly_add_residual", True),
-                                        kernel_init=self.kernel_init(1.0),
+                                        kernel_init=self.kernel_init(scale=1.0),
                                         name=f"down_{i}_attention_{j}")(x, textcontext)
                 # print("down residual for feature level", i, "is of shape", x.shape, "features", dim_in)
                 downs.append(x)
@@ -108,7 +108,7 @@ class Unet(nn.Module):
                 middle_conv_type,
                 name=f"middle_res1_{j}",
                 features=middle_dim_out,
-                kernel_init=self.kernel_init(1.0),
+                kernel_init=self.kernel_init(scale=1.0),
                 kernel_size=(3, 3),
                 strides=(1, 1),
                 activation=self.activation,
@@ -129,13 +129,13 @@ class Unet(nn.Module):
                                     force_fp32_for_softmax=middle_attention.get("force_fp32_for_softmax", False),
                                     norm_inputs=middle_attention.get("norm_inputs", True),
                                     explicitly_add_residual=middle_attention.get("explicitly_add_residual", True),
-                                    kernel_init=self.kernel_init(1.0),
+                                    kernel_init=self.kernel_init(scale=1.0),
                                     name=f"middle_attention_{j}")(x, textcontext)
             x = ResidualBlock(
                 middle_conv_type,
                 name=f"middle_res2_{j}",
                 features=middle_dim_out,
-                kernel_init=self.kernel_init(1.0),
+                kernel_init=self.kernel_init(scale=1.0),
                 kernel_size=(3, 3),
                 strides=(1, 1),
                 activation=self.activation,
@@ -157,7 +157,7 @@ class Unet(nn.Module):
                     up_conv_type,# if j == 0 else "separable",
                     name=f"up_{i}_residual_{j}",
                     features=dim_out,
-                    kernel_init=self.kernel_init(1.0),
+                    kernel_init=self.kernel_init(scale=1.0),
                     kernel_size=kernel_size,
                     strides=(1, 1),
                     activation=self.activation,
@@ -177,7 +177,7 @@ class Unet(nn.Module):
                                         force_fp32_for_softmax=middle_attention.get("force_fp32_for_softmax", False),
                                         norm_inputs=attention_config.get("norm_inputs", True),
                                         explicitly_add_residual=attention_config.get("explicitly_add_residual", True),
-                                        kernel_init=self.kernel_init(1.0),
+                                        kernel_init=self.kernel_init(scale=1.0),
                                         name=f"up_{i}_attention_{j}")(x, textcontext)
             # print("Upscaling ", i, x.shape)
             if i != len(feature_depths) - 1:
@@ -196,7 +196,7 @@ class Unet(nn.Module):
             features=self.feature_depths[0],
             kernel_size=(3, 3),
             strides=(1, 1),
-            kernel_init=self.kernel_init(1.0),
+            kernel_init=self.kernel_init(scale=1.0),
             dtype=self.dtype,
             precision=self.precision
         )(x)
@@ -207,7 +207,7 @@ class Unet(nn.Module):
             conv_type,
             name="final_residual",
             features=self.feature_depths[0],
-            kernel_init=self.kernel_init(1.0),
+            kernel_init=self.kernel_init(scale=1.0),
             kernel_size=(3,3),
             strides=(1, 1),
             activation=self.activation,
@@ -226,7 +226,7 @@ class Unet(nn.Module):
             kernel_size=(3, 3),
             strides=(1, 1),
             # activation=jax.nn.mish
-            kernel_init=self.kernel_init(0.0),
+            kernel_init=self.kernel_init(scale=0.0),
             dtype=self.dtype,
             precision=self.precision
         )(x)

flaxdiff/models/simple_vit.py

@@ -70,7 +70,7 @@ class UViT(nn.Module):
     kernel_init: Callable = partial(kernel_init, scale=1.0)
     add_residualblock_output: bool = False
     norm_inputs: bool = False
-    explicitly_add_residual: bool = False
+    explicitly_add_residual: bool = True
 
     def setup(self):
         if self.norm_groups > 0:

flaxdiff/samplers/common.py

@@ -15,36 +15,76 @@ class DiffusionSampler():
 
     def __init__(self, model:nn.Module, params:dict,  
                  noise_schedule:NoiseScheduler, 
-                 model_output_transform:DiffusionPredictionTransform=EpsilonPredictionTransform()):
+                 model_output_transform:DiffusionPredictionTransform=EpsilonPredictionTransform(),
+                 guidance_scale:float = 0.0,
+                 null_labels_seq:jax.Array=None,
+                 autoencoder=None,
+                 image_size=256,
+                 autoenc_scale_reduction=8,
+                 autoenc_latent_channels=4,
+                 ):
         self.model = model
         self.noise_schedule = noise_schedule
         self.params = params
         self.model_output_transform = model_output_transform
-
-        @jax.jit
-        def sample_model(x_t, t):
-            rates = self.noise_schedule.get_rates(t)
-            c_in = self.model_output_transform.get_input_scale(rates)
-            model_output = self.model.apply(self.params, *self.noise_schedule.transform_inputs(x_t * c_in, t))
-            x_0, eps = self.model_output_transform(x_t, model_output, t, self.noise_schedule)
-            return x_0, eps, model_output
+        self.guidance_scale = guidance_scale
+        self.image_size = image_size
+        self.autoenc_scale_reduction = autoenc_scale_reduction
+        self.autoencoder = autoencoder
+        self.autoenc_latent_channels = autoenc_latent_channels
         
+        if self.guidance_scale > 0:
+            # Classifier free guidance
+            assert null_labels_seq is not None, "Null labels sequence is required for classifier-free guidance"
+            print("Using classifier-free guidance")
+            def sample_model(x_t, t, *additional_inputs):
+                # Concatenate unconditional and conditional inputs
+                x_t_cat = jnp.concatenate([x_t] * 2, axis=0)
+                t_cat = jnp.concatenate([t] * 2, axis=0)
+                rates_cat = self.noise_schedule.get_rates(t_cat)
+                c_in_cat = self.model_output_transform.get_input_scale(rates_cat)
+                
+                text_labels_seq, = additional_inputs
+                text_labels_seq = jnp.concatenate([text_labels_seq, jnp.broadcast_to(null_labels_seq, text_labels_seq.shape)], axis=0)
+                model_output = self.model.apply(self.params, *self.noise_schedule.transform_inputs(x_t_cat * c_in_cat, t_cat), text_labels_seq)
+                # Split model output into unconditional and conditional parts
+                model_output_cond, model_output_uncond = jnp.split(model_output, 2, axis=0)
+                model_output = model_output_uncond + guidance_scale * (model_output_cond - model_output_uncond)
+                
+                x_0, eps = self.model_output_transform(x_t, model_output, t, self.noise_schedule)
+                return x_0, eps, model_output
+        else:
+            # Unconditional sampling
+            def sample_model(x_t, t, *additional_inputs):
+                rates = self.noise_schedule.get_rates(t)
+                c_in = self.model_output_transform.get_input_scale(rates)
+                model_output = self.model.apply(self.params, *self.noise_schedule.transform_inputs(x_t * c_in, t), *additional_inputs)
+                x_0, eps = self.model_output_transform(x_t, model_output, t, self.noise_schedule)
+                return x_0, eps, model_output
+            
+        # if jax.device_count() > 1:
+        #     mesh = jax.sharding.Mesh(jax.devices(), 'data')
+        #     sample_model = shard_map(sample_model, mesh=mesh, in_specs=(P('data'), P('data'), P('data')),
+        #                              out_specs=(P('data'), P('data'), P('data')))
+        sample_model = jax.jit(sample_model)
         self.sample_model = sample_model
 
     # Used to sample from the diffusion model
-    def sample_step(self, current_samples:jnp.ndarray, current_step, next_step=None, state:MarkovState=None) -> tuple[jnp.ndarray, MarkovState]:
+    def sample_step(self, current_samples:jnp.ndarray, current_step, model_conditioning_inputs, next_step=None, state:MarkovState=None) -> tuple[jnp.ndarray, MarkovState]:
         # First clip the noisy images
-        # pred_images = clip_images(pred_images)
         step_ones = jnp.ones((current_samples.shape[0], ), dtype=jnp.int32)
         current_step = step_ones * current_step
         next_step = step_ones * next_step
-        pred_images, pred_noise, _ = self.sample_model(current_samples, current_step)
+        pred_images, pred_noise, _ = self.sample_model(current_samples, current_step, *model_conditioning_inputs)
         # plotImages(pred_images)
+        # pred_images = clip_images(pred_images)
         new_samples, state =  self.take_next_step(current_samples=current_samples, reconstructed_samples=pred_images, 
-                             pred_noise=pred_noise, current_step=current_step, next_step=next_step, state=state)
+                             pred_noise=pred_noise, current_step=current_step, next_step=next_step, state=state,
+                             model_conditioning_inputs=model_conditioning_inputs
+                             )
         return new_samples, state
 
-    def take_next_step(self, current_samples, reconstructed_samples, 
+    def take_next_step(self, current_samples, reconstructed_samples, model_conditioning_inputs,
                  pred_noise, current_step, state:RandomMarkovState, next_step=1) -> tuple[jnp.ndarray, RandomMarkovState]:
         # estimate the q(x_{t-1} | x_t, x_0). 
         # pred_images is x_0, noisy_images is x_t, steps is t
@@ -62,11 +102,16 @@ class DiffusionSampler():
         steps = jnp.linspace(end_step, start_step, diffusion_steps, dtype=jnp.int16)[::-1]
         return steps
     
-    def get_initial_samples(self, num_images, rngs:jax.random.PRNGKey, start_step, image_size=64):
+    def get_initial_samples(self, num_images, rngs:jax.random.PRNGKey, start_step):
         start_step = self.scale_steps(start_step)
         alpha_n, sigma_n = self.noise_schedule.get_rates(start_step)
         variance = jnp.sqrt(alpha_n ** 2 + sigma_n ** 2) 
-        return jax.random.normal(rngs, (num_images, image_size, image_size, 3)) * variance
+        image_size = self.image_size
+        image_channels = 3
+        if self.autoencoder is not None:
+            image_size = image_size // self.autoenc_scale_reduction
+            image_channels = self.autoenc_latent_channels
+        return jax.random.normal(rngs, (num_images, image_size, image_size, image_channels)) * variance
 
     def generate_images(self,
                         num_images=16, 
@@ -75,18 +120,23 @@ class DiffusionSampler():
                         end_step:int = 0,
                         steps_override=None,
                         priors=None, 
-                        rngstate:RandomMarkovState=RandomMarkovState(jax.random.PRNGKey(42))) -> jnp.ndarray:
+                        rngstate:RandomMarkovState=RandomMarkovState(jax.random.PRNGKey(42)),
+                        model_conditioning_inputs:tuple=()
+                        ) -> jnp.ndarray:
         if priors is None:
             rngstate, newrngs = rngstate.get_random_key()
             samples = self.get_initial_samples(num_images, newrngs, start_step)
         else:
             print("Using priors")
+            if self.autoencoder is not None:
+                priors = self.autoencoder.encode(priors)
             samples = priors
 
-        @jax.jit
+        # @jax.jit
         def sample_step(state:RandomMarkovState, samples, current_step, next_step):
             samples, state = self.sample_step(current_samples=samples,
                                               current_step=current_step, 
+                                              model_conditioning_inputs=model_conditioning_inputs,
                                               state=state, next_step=next_step)
             return samples, state
 
@@ -108,6 +158,8 @@ class DiffusionSampler():
             else:
                 # print("last step")
                 step_ones = jnp.ones((num_images, ), dtype=jnp.int32)
-                samples, _, _ = self.sample_model(samples, current_step * step_ones)
+                samples, _, _ = self.sample_model(samples, current_step * step_ones, *model_conditioning_inputs)
+        if self.autoencoder is not None:
+            samples = self.autoencoder.decode(samples)
         samples = clip_images(samples)
-        return samples
+        return samples

flaxdiff/samplers/ddim.py

@@ -1,10 +1,10 @@
 import jax.numpy as jnp
 from .common import DiffusionSampler
-from ..utils import MarkovState
+from ..utils import MarkovState, RandomMarkovState
 
 class DDIMSampler(DiffusionSampler):
-    def take_next_step(self, 
-                 current_samples, reconstructed_samples, 
-                 pred_noise, current_step, state:MarkovState, next_step=None) -> tuple[jnp.ndarray, MarkovState]:
+    def take_next_step(self, current_samples, reconstructed_samples, model_conditioning_inputs,
+                 pred_noise, current_step, state:RandomMarkovState, next_step=1) -> tuple[jnp.ndarray, RandomMarkovState]:
         next_signal_rate, next_noise_rate = self.noise_schedule.get_rates(next_step)
-        return reconstructed_samples * next_signal_rate + pred_noise * next_noise_rate, state
+        return reconstructed_samples * next_signal_rate + pred_noise * next_noise_rate, state
+