Add files using upload-large-folder tool

MindEye_Imagery/src/generative_models/scripts/demo/detect.py

@@ -0,0 +1,156 @@
+import argparse
+
+import cv2
+import numpy as np
+
+try:
+    from imwatermark import WatermarkDecoder
+except ImportError as e:
+    try:
+        # Assume some of the other dependencies such as torch are not fulfilled
+        # import file without loading unnecessary libraries.
+        import importlib.util
+        import sys
+
+        spec = importlib.util.find_spec("imwatermark.maxDct")
+        assert spec is not None
+        maxDct = importlib.util.module_from_spec(spec)
+        sys.modules["maxDct"] = maxDct
+        spec.loader.exec_module(maxDct)
+
+        class WatermarkDecoder(object):
+            """A minimal version of
+            https://github.com/ShieldMnt/invisible-watermark/blob/main/imwatermark/watermark.py
+            to only reconstruct bits using dwtDct"""
+
+            def __init__(self, wm_type="bytes", length=0):
+                assert wm_type == "bits", "Only bits defined in minimal import"
+                self._wmType = wm_type
+                self._wmLen = length
+
+            def reconstruct(self, bits):
+                if len(bits) != self._wmLen:
+                    raise RuntimeError("bits are not matched with watermark length")
+
+                return bits
+
+            def decode(self, cv2Image, method="dwtDct", **configs):
+                (r, c, channels) = cv2Image.shape
+                if r * c < 256 * 256:
+                    raise RuntimeError("image too small, should be larger than 256x256")
+
+                bits = []
+                assert method == "dwtDct"
+                embed = maxDct.EmbedMaxDct(watermarks=[], wmLen=self._wmLen, **configs)
+                bits = embed.decode(cv2Image)
+                return self.reconstruct(bits)
+
+    except:
+        raise e
+
+
+# A fixed 48-bit message that was choosen at random
+# WATERMARK_MESSAGE = 0xB3EC907BB19E
+WATERMARK_MESSAGE = 0b101100111110110010010000011110111011000110011110
+# bin(x)[2:] gives bits of x as str, use int to convert them to 0/1
+WATERMARK_BITS = [int(bit) for bit in bin(WATERMARK_MESSAGE)[2:]]
+MATCH_VALUES = [
+    [27, "No watermark detected"],
+    [33, "Partial watermark match. Cannot determine with certainty."],
+    [
+        35,
+        (
+            "Likely watermarked. In our test 0.02% of real images were "
+            'falsely detected as "Likely watermarked"'
+        ),
+    ],
+    [
+        49,
+        (
+            "Very likely watermarked. In our test no real images were "
+            'falsely detected as "Very likely watermarked"'
+        ),
+    ],
+]
+
+
+class GetWatermarkMatch:
+    def __init__(self, watermark):
+        self.watermark = watermark
+        self.num_bits = len(self.watermark)
+        self.decoder = WatermarkDecoder("bits", self.num_bits)
+
+    def __call__(self, x: np.ndarray) -> np.ndarray:
+        """
+        Detects the number of matching bits the predefined watermark with one
+        or multiple images. Images should be in cv2 format, e.g. h x w x c BGR.
+
+        Args:
+            x: ([B], h w, c) in range [0, 255]
+
+        Returns:
+           number of matched bits ([B],)
+        """
+        squeeze = len(x.shape) == 3
+        if squeeze:
+            x = x[None, ...]
+
+        bs = x.shape[0]
+        detected = np.empty((bs, self.num_bits), dtype=bool)
+        for k in range(bs):
+            detected[k] = self.decoder.decode(x[k], "dwtDct")
+        result = np.sum(detected == self.watermark, axis=-1)
+        if squeeze:
+            return result[0]
+        else:
+            return result
+
+
+get_watermark_match = GetWatermarkMatch(WATERMARK_BITS)
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "filename",
+        nargs="+",
+        type=str,
+        help="Image files to check for watermarks",
+    )
+    opts = parser.parse_args()
+
+    print(
+        """
+        This script tries to detect watermarked images. Please be aware of
+        the following:
+        - As the watermark is supposed to be invisible, there is the risk that
+          watermarked images may not be detected.
+        - To maximize the chance of detection make sure that the image has the same
+          dimensions as when the watermark was applied (most likely 1024x1024
+          or 512x512).
+        - Specific image manipulation may drastically decrease the chance that
+          watermarks can be detected.
+        - There is also the chance that an image has the characteristics of the
+          watermark by chance.
+        - The watermark script is public, anybody may watermark any images, and
+          could therefore claim it to be generated.
+        - All numbers below are based on a test using 10,000 images without any
+          modifications after applying the watermark.
+        """
+    )
+
+    for fn in opts.filename:
+        image = cv2.imread(fn)
+        if image is None:
+            print(f"Couldn't read {fn}. Skipping")
+            continue
+
+        num_bits = get_watermark_match(image)
+        k = 0
+        while num_bits > MATCH_VALUES[k][0]:
+            k += 1
+        print(
+            f"{fn}: {MATCH_VALUES[k][1]}",
+            f"Bits that matched the watermark {num_bits} from {len(WATERMARK_BITS)}\n",
+            sep="\n\t",
+        )

MindEye_Imagery/src/generative_models/scripts/demo/discretization.py

@@ -0,0 +1,59 @@
+import torch
+
+from sgm.modules.diffusionmodules.discretizer import Discretization
+
+
+class Img2ImgDiscretizationWrapper:
+    """
+    wraps a discretizer, and prunes the sigmas
+    params:
+        strength: float between 0.0 and 1.0. 1.0 means full sampling (all sigmas are returned)
+    """
+
+    def __init__(self, discretization: Discretization, strength: float = 1.0):
+        self.discretization = discretization
+        self.strength = strength
+        assert 0.0 <= self.strength <= 1.0
+
+    def __call__(self, *args, **kwargs):
+        # sigmas start large first, and decrease then
+        sigmas = self.discretization(*args, **kwargs)
+        print(f"sigmas after discretization, before pruning img2img: ", sigmas)
+        sigmas = torch.flip(sigmas, (0,))
+        sigmas = sigmas[: max(int(self.strength * len(sigmas)), 1)]
+        print("prune index:", max(int(self.strength * len(sigmas)), 1))
+        sigmas = torch.flip(sigmas, (0,))
+        print(f"sigmas after pruning: ", sigmas)
+        return sigmas
+
+
+class Txt2NoisyDiscretizationWrapper:
+    """
+    wraps a discretizer, and prunes the sigmas
+    params:
+        strength: float between 0.0 and 1.0. 0.0 means full sampling (all sigmas are returned)
+    """
+
+    def __init__(
+        self, discretization: Discretization, strength: float = 0.0, original_steps=None
+    ):
+        self.discretization = discretization
+        self.strength = strength
+        self.original_steps = original_steps
+        assert 0.0 <= self.strength <= 1.0
+
+    def __call__(self, *args, **kwargs):
+        # sigmas start large first, and decrease then
+        sigmas = self.discretization(*args, **kwargs)
+        print(f"sigmas after discretization, before pruning img2img: ", sigmas)
+        sigmas = torch.flip(sigmas, (0,))
+        if self.original_steps is None:
+            steps = len(sigmas)
+        else:
+            steps = self.original_steps + 1
+        prune_index = max(min(int(self.strength * steps) - 1, steps - 1), 0)
+        sigmas = sigmas[prune_index:]
+        print("prune index:", prune_index)
+        sigmas = torch.flip(sigmas, (0,))
+        print(f"sigmas after pruning: ", sigmas)
+        return sigmas

MindEye_Imagery/src/generative_models/scripts/demo/sampling.py

@@ -0,0 +1,364 @@
+from pytorch_lightning import seed_everything
+
+from scripts.demo.streamlit_helpers import *
+
+SAVE_PATH = "outputs/demo/txt2img/"
+
+SD_XL_BASE_RATIOS = {
+    "0.5": (704, 1408),
+    "0.52": (704, 1344),
+    "0.57": (768, 1344),
+    "0.6": (768, 1280),
+    "0.68": (832, 1216),
+    "0.72": (832, 1152),
+    "0.78": (896, 1152),
+    "0.82": (896, 1088),
+    "0.88": (960, 1088),
+    "0.94": (960, 1024),
+    "1.0": (1024, 1024),
+    "1.07": (1024, 960),
+    "1.13": (1088, 960),
+    "1.21": (1088, 896),
+    "1.29": (1152, 896),
+    "1.38": (1152, 832),
+    "1.46": (1216, 832),
+    "1.67": (1280, 768),
+    "1.75": (1344, 768),
+    "1.91": (1344, 704),
+    "2.0": (1408, 704),
+    "2.09": (1472, 704),
+    "2.4": (1536, 640),
+    "2.5": (1600, 640),
+    "2.89": (1664, 576),
+    "3.0": (1728, 576),
+}
+
+VERSION2SPECS = {
+    "SDXL-base-1.0": {
+        "H": 1024,
+        "W": 1024,
+        "C": 4,
+        "f": 8,
+        "is_legacy": False,
+        "config": "configs/inference/sd_xl_base.yaml",
+        "ckpt": "checkpoints/sd_xl_base_1.0.safetensors",
+    },
+    "SDXL-base-0.9": {
+        "H": 1024,
+        "W": 1024,
+        "C": 4,
+        "f": 8,
+        "is_legacy": False,
+        "config": "configs/inference/sd_xl_base.yaml",
+        "ckpt": "checkpoints/sd_xl_base_0.9.safetensors",
+    },
+    "SD-2.1": {
+        "H": 512,
+        "W": 512,
+        "C": 4,
+        "f": 8,
+        "is_legacy": True,
+        "config": "configs/inference/sd_2_1.yaml",
+        "ckpt": "checkpoints/v2-1_512-ema-pruned.safetensors",
+    },
+    "SD-2.1-768": {
+        "H": 768,
+        "W": 768,
+        "C": 4,
+        "f": 8,
+        "is_legacy": True,
+        "config": "configs/inference/sd_2_1_768.yaml",
+        "ckpt": "checkpoints/v2-1_768-ema-pruned.safetensors",
+    },
+    "SDXL-refiner-0.9": {
+        "H": 1024,
+        "W": 1024,
+        "C": 4,
+        "f": 8,
+        "is_legacy": True,
+        "config": "configs/inference/sd_xl_refiner.yaml",
+        "ckpt": "checkpoints/sd_xl_refiner_0.9.safetensors",
+    },
+    "SDXL-refiner-1.0": {
+        "H": 1024,
+        "W": 1024,
+        "C": 4,
+        "f": 8,
+        "is_legacy": True,
+        "config": "configs/inference/sd_xl_refiner.yaml",
+        "ckpt": "checkpoints/sd_xl_refiner_1.0.safetensors",
+    },
+}
+
+
+def load_img(display=True, key=None, device="cuda"):
+    image = get_interactive_image(key=key)
+    if image is None:
+        return None
+    if display:
+        st.image(image)
+    w, h = image.size
+    print(f"loaded input image of size ({w}, {h})")
+    width, height = map(
+        lambda x: x - x % 64, (w, h)
+    )  # resize to integer multiple of 64
+    image = image.resize((width, height))
+    image = np.array(image.convert("RGB"))
+    image = image[None].transpose(0, 3, 1, 2)
+    image = torch.from_numpy(image).to(dtype=torch.float32) / 127.5 - 1.0
+    return image.to(device)
+
+
+def run_txt2img(
+    state,
+    version,
+    version_dict,
+    is_legacy=False,
+    return_latents=False,
+    filter=None,
+    stage2strength=None,
+):
+    if version.startswith("SDXL-base"):
+        W, H = st.selectbox("Resolution:", list(SD_XL_BASE_RATIOS.values()), 10)
+    else:
+        H = st.number_input("H", value=version_dict["H"], min_value=64, max_value=2048)
+        W = st.number_input("W", value=version_dict["W"], min_value=64, max_value=2048)
+    C = version_dict["C"]
+    F = version_dict["f"]
+
+    init_dict = {
+        "orig_width": W,
+        "orig_height": H,
+        "target_width": W,
+        "target_height": H,
+    }
+    value_dict = init_embedder_options(
+        get_unique_embedder_keys_from_conditioner(state["model"].conditioner),
+        init_dict,
+        prompt=prompt,
+        negative_prompt=negative_prompt,
+    )
+    sampler, num_rows, num_cols = init_sampling(stage2strength=stage2strength)
+    num_samples = num_rows * num_cols
+
+    if st.button("Sample"):
+        st.write(f"**Model I:** {version}")
+        out = do_sample(
+            state["model"],
+            sampler,
+            value_dict,
+            num_samples,
+            H,
+            W,
+            C,
+            F,
+            force_uc_zero_embeddings=["txt"] if not is_legacy else [],
+            return_latents=return_latents,
+            filter=filter,
+        )
+        return out
+
+
+def run_img2img(
+    state,
+    version_dict,
+    is_legacy=False,
+    return_latents=False,
+    filter=None,
+    stage2strength=None,
+):
+    img = load_img()
+    if img is None:
+        return None
+    H, W = img.shape[2], img.shape[3]
+
+    init_dict = {
+        "orig_width": W,
+        "orig_height": H,
+        "target_width": W,
+        "target_height": H,
+    }
+    value_dict = init_embedder_options(
+        get_unique_embedder_keys_from_conditioner(state["model"].conditioner),
+        init_dict,
+        prompt=prompt,
+        negative_prompt=negative_prompt,
+    )
+    strength = st.number_input(
+        "**Img2Img Strength**", value=0.75, min_value=0.0, max_value=1.0
+    )
+    sampler, num_rows, num_cols = init_sampling(
+        img2img_strength=strength,
+        stage2strength=stage2strength,
+    )
+    num_samples = num_rows * num_cols
+
+    if st.button("Sample"):
+        out = do_img2img(
+            repeat(img, "1 ... -> n ...", n=num_samples),
+            state["model"],
+            sampler,
+            value_dict,
+            num_samples,
+            force_uc_zero_embeddings=["txt"] if not is_legacy else [],
+            return_latents=return_latents,
+            filter=filter,
+        )
+        return out
+
+
+def apply_refiner(
+    input,
+    state,
+    sampler,
+    num_samples,
+    prompt,
+    negative_prompt,
+    filter=None,
+    finish_denoising=False,
+):
+    init_dict = {
+        "orig_width": input.shape[3] * 8,
+        "orig_height": input.shape[2] * 8,
+        "target_width": input.shape[3] * 8,
+        "target_height": input.shape[2] * 8,
+    }
+
+    value_dict = init_dict
+    value_dict["prompt"] = prompt
+    value_dict["negative_prompt"] = negative_prompt
+
+    value_dict["crop_coords_top"] = 0
+    value_dict["crop_coords_left"] = 0
+
+    value_dict["aesthetic_score"] = 6.0
+    value_dict["negative_aesthetic_score"] = 2.5
+
+    st.warning(f"refiner input shape: {input.shape}")
+    samples = do_img2img(
+        input,
+        state["model"],
+        sampler,
+        value_dict,
+        num_samples,
+        skip_encode=True,
+        filter=filter,
+        add_noise=not finish_denoising,
+    )
+
+    return samples
+
+
+if __name__ == "__main__":
+    st.title("Stable Diffusion")
+    version = st.selectbox("Model Version", list(VERSION2SPECS.keys()), 0)
+    version_dict = VERSION2SPECS[version]
+    if st.checkbox("Load Model"):
+        mode = st.radio("Mode", ("txt2img", "img2img"), 0)
+    else:
+        mode = "skip"
+    st.write("__________________________")
+
+    set_lowvram_mode(st.checkbox("Low vram mode", True))
+
+    if version.startswith("SDXL-base"):
+        add_pipeline = st.checkbox("Load SDXL-refiner?", False)
+        st.write("__________________________")
+    else:
+        add_pipeline = False
+
+    seed = st.sidebar.number_input("seed", value=42, min_value=0, max_value=int(1e9))
+    seed_everything(seed)
+
+    save_locally, save_path = init_save_locally(os.path.join(SAVE_PATH, version))
+
+    if mode != "skip":
+        state = init_st(version_dict, load_filter=True)
+        if state["msg"]:
+            st.info(state["msg"])
+        model = state["model"]
+
+    is_legacy = version_dict["is_legacy"]
+
+    prompt = st.text_input(
+        "prompt",
+        "Astronaut in a jungle, cold color palette, muted colors, detailed, 8k",
+    )
+    if is_legacy:
+        negative_prompt = st.text_input("negative prompt", "")
+    else:
+        negative_prompt = ""  # which is unused
+
+    stage2strength = None
+    finish_denoising = False
+
+    if add_pipeline:
+        st.write("__________________________")
+        version2 = st.selectbox("Refiner:", ["SDXL-refiner-1.0", "SDXL-refiner-0.9"])
+        st.warning(
+            f"Running with {version2} as the second stage model. Make sure to provide (V)RAM :) "
+        )
+        st.write("**Refiner Options:**")
+
+        version_dict2 = VERSION2SPECS[version2]
+        state2 = init_st(version_dict2, load_filter=False)
+        st.info(state2["msg"])
+
+        stage2strength = st.number_input(
+            "**Refinement strength**", value=0.15, min_value=0.0, max_value=1.0
+        )
+
+        sampler2, *_ = init_sampling(
+            key=2,
+            img2img_strength=stage2strength,
+            specify_num_samples=False,
+        )
+        st.write("__________________________")
+        finish_denoising = st.checkbox("Finish denoising with refiner.", True)
+        if not finish_denoising:
+            stage2strength = None
+
+    if mode == "txt2img":
+        out = run_txt2img(
+            state,
+            version,
+            version_dict,
+            is_legacy=is_legacy,
+            return_latents=add_pipeline,
+            filter=state.get("filter"),
+            stage2strength=stage2strength,
+        )
+    elif mode == "img2img":
+        out = run_img2img(
+            state,
+            version_dict,
+            is_legacy=is_legacy,
+            return_latents=add_pipeline,
+            filter=state.get("filter"),
+            stage2strength=stage2strength,
+        )
+    elif mode == "skip":
+        out = None
+    else:
+        raise ValueError(f"unknown mode {mode}")
+    if isinstance(out, (tuple, list)):
+        samples, samples_z = out
+    else:
+        samples = out
+        samples_z = None
+
+    if add_pipeline and samples_z is not None:
+        st.write("**Running Refinement Stage**")
+        samples = apply_refiner(
+            samples_z,
+            state2,
+            sampler2,
+            samples_z.shape[0],
+            prompt=prompt,
+            negative_prompt=negative_prompt if is_legacy else "",
+            filter=state.get("filter"),
+            finish_denoising=finish_denoising,
+        )
+
+    if save_locally and samples is not None:
+        perform_save_locally(save_path, samples)

MindEye_Imagery/src/generative_models/scripts/demo/streamlit_helpers.py

@@ -0,0 +1,887 @@
+import copy
+import math
+import os
+from glob import glob
+from typing import Dict, List, Optional, Tuple, Union
+
+import cv2
+import numpy as np
+import streamlit as st
+import torch
+import torch.nn as nn
+import torchvision.transforms as TT
+from einops import rearrange, repeat
+from imwatermark import WatermarkEncoder
+from omegaconf import ListConfig, OmegaConf
+from PIL import Image
+from safetensors.torch import load_file as load_safetensors
+from torch import autocast
+from torchvision import transforms
+from torchvision.utils import make_grid, save_image
+
+from scripts.demo.discretization import (Img2ImgDiscretizationWrapper,
+                                         Txt2NoisyDiscretizationWrapper)
+from scripts.util.detection.nsfw_and_watermark_dectection import \
+    DeepFloydDataFiltering
+from sgm.inference.helpers import embed_watermark
+from sgm.modules.diffusionmodules.guiders import (LinearPredictionGuider,
+                                                  VanillaCFG)
+from sgm.modules.diffusionmodules.sampling import (DPMPP2MSampler,
+                                                   DPMPP2SAncestralSampler,
+                                                   EulerAncestralSampler,
+                                                   EulerEDMSampler,
+                                                   HeunEDMSampler,
+                                                   LinearMultistepSampler)
+from sgm.util import append_dims, default, instantiate_from_config
+
+
+@st.cache_resource()
+def init_st(version_dict, load_ckpt=True, load_filter=True):
+    state = dict()
+    if not "model" in state:
+        config = version_dict["config"]
+        ckpt = version_dict["ckpt"]
+
+        config = OmegaConf.load(config)
+        model, msg = load_model_from_config(config, ckpt if load_ckpt else None)
+
+        state["msg"] = msg
+        state["model"] = model
+        state["ckpt"] = ckpt if load_ckpt else None
+        state["config"] = config
+        if load_filter:
+            state["filter"] = DeepFloydDataFiltering(verbose=False)
+    return state
+
+
+def load_model(model):
+    model.cuda()
+
+
+lowvram_mode = False
+
+
+def set_lowvram_mode(mode):
+    global lowvram_mode
+    lowvram_mode = mode
+
+
+def initial_model_load(model):
+    global lowvram_mode
+    if lowvram_mode:
+        model.model.half()
+    else:
+        model.cuda()
+    return model
+
+
+def unload_model(model):
+    global lowvram_mode
+    if lowvram_mode:
+        model.cpu()
+        torch.cuda.empty_cache()
+
+
+def load_model_from_config(config, ckpt=None, verbose=True):
+    model = instantiate_from_config(config.model)
+
+    if ckpt is not None:
+        print(f"Loading model from {ckpt}")
+        if ckpt.endswith("ckpt"):
+            pl_sd = torch.load(ckpt, map_location="cpu")
+            if "global_step" in pl_sd:
+                global_step = pl_sd["global_step"]
+                st.info(f"loaded ckpt from global step {global_step}")
+                print(f"Global Step: {pl_sd['global_step']}")
+            sd = pl_sd["state_dict"]
+        elif ckpt.endswith("safetensors"):
+            sd = load_safetensors(ckpt)
+        else:
+            raise NotImplementedError
+
+        msg = None
+
+        m, u = model.load_state_dict(sd, strict=False)
+
+        if len(m) > 0 and verbose:
+            print("missing keys:")
+            print(m)
+        if len(u) > 0 and verbose:
+            print("unexpected keys:")
+            print(u)
+    else:
+        msg = None
+
+    model = initial_model_load(model)
+    model.eval()
+    return model, msg
+
+
+def get_unique_embedder_keys_from_conditioner(conditioner):
+    return list(set([x.input_key for x in conditioner.embedders]))
+
+
+def init_embedder_options(keys, init_dict, prompt=None, negative_prompt=None):
+    # Hardcoded demo settings; might undergo some changes in the future
+
+    value_dict = {}
+    for key in keys:
+        if key == "txt":
+            if prompt is None:
+                prompt = "A professional photograph of an astronaut riding a pig"
+            if negative_prompt is None:
+                negative_prompt = ""
+
+            prompt = st.text_input("Prompt", prompt)
+            negative_prompt = st.text_input("Negative prompt", negative_prompt)
+
+            value_dict["prompt"] = prompt
+            value_dict["negative_prompt"] = negative_prompt
+
+        if key == "original_size_as_tuple":
+            orig_width = st.number_input(
+                "orig_width",
+                value=init_dict["orig_width"],
+                min_value=16,
+            )
+            orig_height = st.number_input(
+                "orig_height",
+                value=init_dict["orig_height"],
+                min_value=16,
+            )
+
+            value_dict["orig_width"] = orig_width
+            value_dict["orig_height"] = orig_height
+
+        if key == "crop_coords_top_left":
+            crop_coord_top = st.number_input("crop_coords_top", value=0, min_value=0)
+            crop_coord_left = st.number_input("crop_coords_left", value=0, min_value=0)
+
+            value_dict["crop_coords_top"] = crop_coord_top
+            value_dict["crop_coords_left"] = crop_coord_left
+
+        if key == "aesthetic_score":
+            value_dict["aesthetic_score"] = 6.0
+            value_dict["negative_aesthetic_score"] = 2.5
+
+        if key == "target_size_as_tuple":
+            value_dict["target_width"] = init_dict["target_width"]
+            value_dict["target_height"] = init_dict["target_height"]
+
+        if key in ["fps_id", "fps"]:
+            fps = st.number_input("fps", value=6, min_value=1)
+
+            value_dict["fps"] = fps
+            value_dict["fps_id"] = fps - 1
+
+        if key == "motion_bucket_id":
+            mb_id = st.number_input("motion bucket id", 0, 511, value=127)
+            value_dict["motion_bucket_id"] = mb_id
+
+        if key == "pool_image":
+            st.text("Image for pool conditioning")
+            image = load_img(
+                key="pool_image_input",
+                size=224,
+                center_crop=True,
+            )
+            if image is None:
+                st.info("Need an image here")
+                image = torch.zeros(1, 3, 224, 224)
+            value_dict["pool_image"] = image
+
+    return value_dict
+
+
+def perform_save_locally(save_path, samples):
+    os.makedirs(os.path.join(save_path), exist_ok=True)
+    base_count = len(os.listdir(os.path.join(save_path)))
+    samples = embed_watermark(samples)
+    for sample in samples:
+        sample = 255.0 * rearrange(sample.cpu().numpy(), "c h w -> h w c")
+        Image.fromarray(sample.astype(np.uint8)).save(
+            os.path.join(save_path, f"{base_count:09}.png")
+        )
+        base_count += 1
+
+
+def init_save_locally(_dir, init_value: bool = False):
+    save_locally = st.sidebar.checkbox("Save images locally", value=init_value)
+    if save_locally:
+        save_path = st.text_input("Save path", value=os.path.join(_dir, "samples"))
+    else:
+        save_path = None
+
+    return save_locally, save_path
+
+
+def get_guider(options, key):
+    guider = st.sidebar.selectbox(
+        f"Discretization #{key}",
+        [
+            "VanillaCFG",
+            "IdentityGuider",
+            "LinearPredictionGuider",
+        ],
+        options.get("guider", 0),
+    )
+
+    additional_guider_kwargs = options.pop("additional_guider_kwargs", {})
+
+    if guider == "IdentityGuider":
+        guider_config = {
+            "target": "sgm.modules.diffusionmodules.guiders.IdentityGuider"
+        }
+    elif guider == "VanillaCFG":
+        scale = st.number_input(
+            f"cfg-scale #{key}",
+            value=options.get("cfg", 5.0),
+            min_value=0.0,
+        )
+
+        guider_config = {
+            "target": "sgm.modules.diffusionmodules.guiders.VanillaCFG",
+            "params": {
+                "scale": scale,
+                **additional_guider_kwargs,
+            },
+        }
+    elif guider == "LinearPredictionGuider":
+        max_scale = st.number_input(
+            f"max-cfg-scale #{key}",
+            value=options.get("cfg", 1.5),
+            min_value=1.0,
+        )
+        min_scale = st.number_input(
+            f"min guidance scale",
+            value=options.get("min_cfg", 1.0),
+            min_value=1.0,
+            max_value=10.0,
+        )
+
+        guider_config = {
+            "target": "sgm.modules.diffusionmodules.guiders.LinearPredictionGuider",
+            "params": {
+                "max_scale": max_scale,
+                "min_scale": min_scale,
+                "num_frames": options["num_frames"],
+                **additional_guider_kwargs,
+            },
+        }
+    else:
+        raise NotImplementedError
+    return guider_config
+
+
+def init_sampling(
+    key=1,
+    img2img_strength: Optional[float] = None,
+    specify_num_samples: bool = True,
+    stage2strength: Optional[float] = None,
+    options: Optional[Dict[str, int]] = None,
+):
+    options = {} if options is None else options
+
+    num_rows, num_cols = 1, 1
+    if specify_num_samples:
+        num_cols = st.number_input(
+            f"num cols #{key}", value=num_cols, min_value=1, max_value=10
+        )
+
+    steps = st.sidebar.number_input(
+        f"steps #{key}", value=options.get("num_steps", 40), min_value=1, max_value=1000
+    )
+    sampler = st.sidebar.selectbox(
+        f"Sampler #{key}",
+        [
+            "EulerEDMSampler",
+            "HeunEDMSampler",
+            "EulerAncestralSampler",
+            "DPMPP2SAncestralSampler",
+            "DPMPP2MSampler",
+            "LinearMultistepSampler",
+        ],
+        options.get("sampler", 0),
+    )
+    discretization = st.sidebar.selectbox(
+        f"Discretization #{key}",
+        [
+            "LegacyDDPMDiscretization",
+            "EDMDiscretization",
+        ],
+        options.get("discretization", 0),
+    )
+
+    discretization_config = get_discretization(discretization, options=options, key=key)
+
+    guider_config = get_guider(options=options, key=key)
+
+    sampler = get_sampler(sampler, steps, discretization_config, guider_config, key=key)
+    if img2img_strength is not None:
+        st.warning(
+            f"Wrapping {sampler.__class__.__name__} with Img2ImgDiscretizationWrapper"
+        )
+        sampler.discretization = Img2ImgDiscretizationWrapper(
+            sampler.discretization, strength=img2img_strength
+        )
+    if stage2strength is not None:
+        sampler.discretization = Txt2NoisyDiscretizationWrapper(
+            sampler.discretization, strength=stage2strength, original_steps=steps
+        )
+    return sampler, num_rows, num_cols
+
+
+def get_discretization(discretization, options, key=1):
+    if discretization == "LegacyDDPMDiscretization":
+        discretization_config = {
+            "target": "sgm.modules.diffusionmodules.discretizer.LegacyDDPMDiscretization",
+        }
+    elif discretization == "EDMDiscretization":
+        sigma_min = st.number_input(
+            f"sigma_min #{key}", value=options.get("sigma_min", 0.03)
+        )  # 0.0292
+        sigma_max = st.number_input(
+            f"sigma_max #{key}", value=options.get("sigma_max", 14.61)
+        )  # 14.6146
+        rho = st.number_input(f"rho #{key}", value=options.get("rho", 3.0))
+        discretization_config = {
+            "target": "sgm.modules.diffusionmodules.discretizer.EDMDiscretization",
+            "params": {
+                "sigma_min": sigma_min,
+                "sigma_max": sigma_max,
+                "rho": rho,
+            },
+        }
+
+    return discretization_config
+
+
+def get_sampler(sampler_name, steps, discretization_config, guider_config, key=1):
+    if sampler_name == "EulerEDMSampler" or sampler_name == "HeunEDMSampler":
+        s_churn = st.sidebar.number_input(f"s_churn #{key}", value=0.0, min_value=0.0)
+        s_tmin = st.sidebar.number_input(f"s_tmin #{key}", value=0.0, min_value=0.0)
+        s_tmax = st.sidebar.number_input(f"s_tmax #{key}", value=999.0, min_value=0.0)
+        s_noise = st.sidebar.number_input(f"s_noise #{key}", value=1.0, min_value=0.0)
+
+        if sampler_name == "EulerEDMSampler":
+            sampler = EulerEDMSampler(
+                num_steps=steps,
+                discretization_config=discretization_config,
+                guider_config=guider_config,
+                s_churn=s_churn,
+                s_tmin=s_tmin,
+                s_tmax=s_tmax,
+                s_noise=s_noise,
+                verbose=True,
+            )
+        elif sampler_name == "HeunEDMSampler":
+            sampler = HeunEDMSampler(
+                num_steps=steps,
+                discretization_config=discretization_config,
+                guider_config=guider_config,
+                s_churn=s_churn,
+                s_tmin=s_tmin,
+                s_tmax=s_tmax,
+                s_noise=s_noise,
+                verbose=True,
+            )
+    elif (
+        sampler_name == "EulerAncestralSampler"
+        or sampler_name == "DPMPP2SAncestralSampler"
+    ):
+        s_noise = st.sidebar.number_input("s_noise", value=1.0, min_value=0.0)
+        eta = st.sidebar.number_input("eta", value=1.0, min_value=0.0)
+
+        if sampler_name == "EulerAncestralSampler":
+            sampler = EulerAncestralSampler(
+                num_steps=steps,
+                discretization_config=discretization_config,
+                guider_config=guider_config,
+                eta=eta,
+                s_noise=s_noise,
+                verbose=True,
+            )
+        elif sampler_name == "DPMPP2SAncestralSampler":
+            sampler = DPMPP2SAncestralSampler(
+                num_steps=steps,
+                discretization_config=discretization_config,
+                guider_config=guider_config,
+                eta=eta,
+                s_noise=s_noise,
+                verbose=True,
+            )
+    elif sampler_name == "DPMPP2MSampler":
+        sampler = DPMPP2MSampler(
+            num_steps=steps,
+            discretization_config=discretization_config,
+            guider_config=guider_config,
+            verbose=True,
+        )
+    elif sampler_name == "LinearMultistepSampler":
+        order = st.sidebar.number_input("order", value=4, min_value=1)
+        sampler = LinearMultistepSampler(
+            num_steps=steps,
+            discretization_config=discretization_config,
+            guider_config=guider_config,
+            order=order,
+            verbose=True,
+        )
+    else:
+        raise ValueError(f"unknown sampler {sampler_name}!")
+
+    return sampler
+
+
+def get_interactive_image() -> Image.Image:
+    image = st.file_uploader("Input", type=["jpg", "JPEG", "png"])
+    if image is not None:
+        image = Image.open(image)
+        if not image.mode == "RGB":
+            image = image.convert("RGB")
+        return image
+
+
+def load_img(
+    display: bool = True,
+    size: Union[None, int, Tuple[int, int]] = None,
+    center_crop: bool = False,
+):
+    image = get_interactive_image()
+    if image is None:
+        return None
+    if display:
+        st.image(image)
+    w, h = image.size
+    print(f"loaded input image of size ({w}, {h})")
+
+    transform = []
+    if size is not None:
+        transform.append(transforms.Resize(size))
+    if center_crop:
+        transform.append(transforms.CenterCrop(size))
+    transform.append(transforms.ToTensor())
+    transform.append(transforms.Lambda(lambda x: 2.0 * x - 1.0))
+
+    transform = transforms.Compose(transform)
+    img = transform(image)[None, ...]
+    st.text(f"input min/max/mean: {img.min():.3f}/{img.max():.3f}/{img.mean():.3f}")
+    return img
+
+
+def get_init_img(batch_size=1, key=None):
+    init_image = load_img(key=key).cuda()
+    init_image = repeat(init_image, "1 ... -> b ...", b=batch_size)
+    return init_image
+
+
+def do_sample(
+    model,
+    sampler,
+    value_dict,
+    num_samples,
+    H,
+    W,
+    C,
+    F,
+    force_uc_zero_embeddings: Optional[List] = None,
+    force_cond_zero_embeddings: Optional[List] = None,
+    batch2model_input: List = None,
+    return_latents=False,
+    filter=None,
+    T=None,
+    additional_batch_uc_fields=None,
+    decoding_t=None,
+):
+    force_uc_zero_embeddings = default(force_uc_zero_embeddings, [])
+    batch2model_input = default(batch2model_input, [])
+    additional_batch_uc_fields = default(additional_batch_uc_fields, [])
+
+    st.text("Sampling")
+
+    outputs = st.empty()
+    precision_scope = autocast
+    with torch.no_grad():
+        with precision_scope("cuda"):
+            with model.ema_scope():
+                if T is not None:
+                    num_samples = [num_samples, T]
+                else:
+                    num_samples = [num_samples]
+
+                load_model(model.conditioner)
+                batch, batch_uc = get_batch(
+                    get_unique_embedder_keys_from_conditioner(model.conditioner),
+                    value_dict,
+                    num_samples,
+                    T=T,
+                    additional_batch_uc_fields=additional_batch_uc_fields,
+                )
+
+                c, uc = model.conditioner.get_unconditional_conditioning(
+                    batch,
+                    batch_uc=batch_uc,
+                    force_uc_zero_embeddings=force_uc_zero_embeddings,
+                    force_cond_zero_embeddings=force_cond_zero_embeddings,
+                )
+                unload_model(model.conditioner)
+
+                for k in c:
+                    if not k == "crossattn":
+                        c[k], uc[k] = map(
+                            lambda y: y[k][: math.prod(num_samples)].to("cuda"), (c, uc)
+                        )
+                    if k in ["crossattn", "concat"] and T is not None:
+                        uc[k] = repeat(uc[k], "b ... -> b t ...", t=T)
+                        uc[k] = rearrange(uc[k], "b t ... -> (b t) ...", t=T)
+                        c[k] = repeat(c[k], "b ... -> b t ...", t=T)
+                        c[k] = rearrange(c[k], "b t ... -> (b t) ...", t=T)
+
+                additional_model_inputs = {}
+                for k in batch2model_input:
+                    if k == "image_only_indicator":
+                        assert T is not None
+
+                        if isinstance(
+                            sampler.guider, (VanillaCFG, LinearPredictionGuider)
+                        ):
+                            additional_model_inputs[k] = torch.zeros(
+                                num_samples[0] * 2, num_samples[1]
+                            ).to("cuda")
+                        else:
+                            additional_model_inputs[k] = torch.zeros(num_samples).to(
+                                "cuda"
+                            )
+                    else:
+                        additional_model_inputs[k] = batch[k]
+
+                shape = (math.prod(num_samples), C, H // F, W // F)
+                randn = torch.randn(shape).to("cuda")
+
+                def denoiser(input, sigma, c):
+                    return model.denoiser(
+                        model.model, input, sigma, c, **additional_model_inputs
+                    )
+
+                load_model(model.denoiser)
+                load_model(model.model)
+                samples_z = sampler(denoiser, randn, cond=c, uc=uc)
+                unload_model(model.model)
+                unload_model(model.denoiser)
+
+                load_model(model.first_stage_model)
+                model.en_and_decode_n_samples_a_time = (
+                    decoding_t  # Decode n frames at a time
+                )
+                samples_x = model.decode_first_stage(samples_z)
+                samples = torch.clamp((samples_x + 1.0) / 2.0, min=0.0, max=1.0)
+                unload_model(model.first_stage_model)
+
+                if filter is not None:
+                    samples = filter(samples)
+
+                if T is None:
+                    grid = torch.stack([samples])
+                    grid = rearrange(grid, "n b c h w -> (n h) (b w) c")
+                    outputs.image(grid.cpu().numpy())
+                else:
+                    as_vids = rearrange(samples, "(b t) c h w -> b t c h w", t=T)
+                    for i, vid in enumerate(as_vids):
+                        grid = rearrange(make_grid(vid, nrow=4), "c h w -> h w c")
+                        st.image(
+                            grid.cpu().numpy(),
+                            f"Sample #{i} as image",
+                        )
+
+                if return_latents:
+                    return samples, samples_z
+                return samples
+
+
+def get_batch(
+    keys,
+    value_dict: dict,
+    N: Union[List, ListConfig],
+    device: str = "cuda",
+    T: int = None,
+    additional_batch_uc_fields: List[str] = [],
+):
+    # Hardcoded demo setups; might undergo some changes in the future
+
+    batch = {}
+    batch_uc = {}
+
+    for key in keys:
+        if key == "txt":
+            batch["txt"] = [value_dict["prompt"]] * math.prod(N)
+
+            batch_uc["txt"] = [value_dict["negative_prompt"]] * math.prod(N)
+
+        elif key == "original_size_as_tuple":
+            batch["original_size_as_tuple"] = (
+                torch.tensor([value_dict["orig_height"], value_dict["orig_width"]])
+                .to(device)
+                .repeat(math.prod(N), 1)
+            )
+        elif key == "crop_coords_top_left":
+            batch["crop_coords_top_left"] = (
+                torch.tensor(
+                    [value_dict["crop_coords_top"], value_dict["crop_coords_left"]]
+                )
+                .to(device)
+                .repeat(math.prod(N), 1)
+            )
+        elif key == "aesthetic_score":
+            batch["aesthetic_score"] = (
+                torch.tensor([value_dict["aesthetic_score"]])
+                .to(device)
+                .repeat(math.prod(N), 1)
+            )
+            batch_uc["aesthetic_score"] = (
+                torch.tensor([value_dict["negative_aesthetic_score"]])
+                .to(device)
+                .repeat(math.prod(N), 1)
+            )
+
+        elif key == "target_size_as_tuple":
+            batch["target_size_as_tuple"] = (
+                torch.tensor([value_dict["target_height"], value_dict["target_width"]])
+                .to(device)
+                .repeat(math.prod(N), 1)
+            )
+        elif key == "fps":
+            batch[key] = (
+                torch.tensor([value_dict["fps"]]).to(device).repeat(math.prod(N))
+            )
+        elif key == "fps_id":
+            batch[key] = (
+                torch.tensor([value_dict["fps_id"]]).to(device).repeat(math.prod(N))
+            )
+        elif key == "motion_bucket_id":
+            batch[key] = (
+                torch.tensor([value_dict["motion_bucket_id"]])
+                .to(device)
+                .repeat(math.prod(N))
+            )
+        elif key == "pool_image":
+            batch[key] = repeat(value_dict[key], "1 ... -> b ...", b=math.prod(N)).to(
+                device, dtype=torch.half
+            )
+        elif key == "cond_aug":
+            batch[key] = repeat(
+                torch.tensor([value_dict["cond_aug"]]).to("cuda"),
+                "1 -> b",
+                b=math.prod(N),
+            )
+        elif key == "cond_frames":
+            batch[key] = repeat(value_dict["cond_frames"], "1 ... -> b ...", b=N[0])
+        elif key == "cond_frames_without_noise":
+            batch[key] = repeat(
+                value_dict["cond_frames_without_noise"], "1 ... -> b ...", b=N[0]
+            )
+        else:
+            batch[key] = value_dict[key]
+
+    if T is not None:
+        batch["num_video_frames"] = T
+
+    for key in batch.keys():
+        if key not in batch_uc and isinstance(batch[key], torch.Tensor):
+            batch_uc[key] = torch.clone(batch[key])
+        elif key in additional_batch_uc_fields and key not in batch_uc:
+            batch_uc[key] = copy.copy(batch[key])
+    return batch, batch_uc
+
+
+@torch.no_grad()
+def do_img2img(
+    img,
+    model,
+    sampler,
+    value_dict,
+    num_samples,
+    force_uc_zero_embeddings: Optional[List] = None,
+    force_cond_zero_embeddings: Optional[List] = None,
+    additional_kwargs={},
+    offset_noise_level: int = 0.0,
+    return_latents=False,
+    skip_encode=False,
+    filter=None,
+    add_noise=True,
+):
+    st.text("Sampling")
+
+    outputs = st.empty()
+    precision_scope = autocast
+    with torch.no_grad():
+        with precision_scope("cuda"):
+            with model.ema_scope():
+                load_model(model.conditioner)
+                batch, batch_uc = get_batch(
+                    get_unique_embedder_keys_from_conditioner(model.conditioner),
+                    value_dict,
+                    [num_samples],
+                )
+                c, uc = model.conditioner.get_unconditional_conditioning(
+                    batch,
+                    batch_uc=batch_uc,
+                    force_uc_zero_embeddings=force_uc_zero_embeddings,
+                    force_cond_zero_embeddings=force_cond_zero_embeddings,
+                )
+                unload_model(model.conditioner)
+                for k in c:
+                    c[k], uc[k] = map(lambda y: y[k][:num_samples].to("cuda"), (c, uc))
+
+                for k in additional_kwargs:
+                    c[k] = uc[k] = additional_kwargs[k]
+                if skip_encode:
+                    z = img
+                else:
+                    load_model(model.first_stage_model)
+                    z = model.encode_first_stage(img)
+                    unload_model(model.first_stage_model)
+
+                noise = torch.randn_like(z)
+
+                sigmas = sampler.discretization(sampler.num_steps).cuda()
+                sigma = sigmas[0]
+
+                st.info(f"all sigmas: {sigmas}")
+                st.info(f"noising sigma: {sigma}")
+                if offset_noise_level > 0.0:
+                    noise = noise + offset_noise_level * append_dims(
+                        torch.randn(z.shape[0], device=z.device), z.ndim
+                    )
+                if add_noise:
+                    noised_z = z + noise * append_dims(sigma, z.ndim).cuda()
+                    noised_z = noised_z / torch.sqrt(
+                        1.0 + sigmas[0] ** 2.0
+                    )  # Note: hardcoded to DDPM-like scaling. need to generalize later.
+                else:
+                    noised_z = z / torch.sqrt(1.0 + sigmas[0] ** 2.0)
+
+                def denoiser(x, sigma, c):
+                    return model.denoiser(model.model, x, sigma, c)
+
+                load_model(model.denoiser)
+                load_model(model.model)
+                samples_z = sampler(denoiser, noised_z, cond=c, uc=uc)
+                unload_model(model.model)
+                unload_model(model.denoiser)
+
+                load_model(model.first_stage_model)
+                samples_x = model.decode_first_stage(samples_z)
+                unload_model(model.first_stage_model)
+                samples = torch.clamp((samples_x + 1.0) / 2.0, min=0.0, max=1.0)
+
+                if filter is not None:
+                    samples = filter(samples)
+
+                grid = rearrange(grid, "n b c h w -> (n h) (b w) c")
+                outputs.image(grid.cpu().numpy())
+                if return_latents:
+                    return samples, samples_z
+                return samples
+
+
+def get_resizing_factor(
+    desired_shape: Tuple[int, int], current_shape: Tuple[int, int]
+) -> float:
+    r_bound = desired_shape[1] / desired_shape[0]
+    aspect_r = current_shape[1] / current_shape[0]
+    if r_bound >= 1.0:
+        if aspect_r >= r_bound:
+            factor = min(desired_shape) / min(current_shape)
+        else:
+            if aspect_r < 1.0:
+                factor = max(desired_shape) / min(current_shape)
+            else:
+                factor = max(desired_shape) / max(current_shape)
+    else:
+        if aspect_r <= r_bound:
+            factor = min(desired_shape) / min(current_shape)
+        else:
+            if aspect_r > 1:
+                factor = max(desired_shape) / min(current_shape)
+            else:
+                factor = max(desired_shape) / max(current_shape)
+
+    return factor
+
+
+def get_interactive_image(key=None) -> Image.Image:
+    image = st.file_uploader("Input", type=["jpg", "JPEG", "png"], key=key)
+    if image is not None:
+        image = Image.open(image)
+        if not image.mode == "RGB":
+            image = image.convert("RGB")
+        return image
+
+
+def load_img_for_prediction(
+    W: int, H: int, display=True, key=None, device="cuda"
+) -> torch.Tensor:
+    image = get_interactive_image(key=key)
+    if image is None:
+        return None
+    if display:
+        st.image(image)
+    w, h = image.size
+
+    image = np.array(image).transpose(2, 0, 1)
+    image = torch.from_numpy(image).to(dtype=torch.float32) / 255.0
+    image = image.unsqueeze(0)
+
+    rfs = get_resizing_factor((H, W), (h, w))
+    resize_size = [int(np.ceil(rfs * s)) for s in (h, w)]
+    top = (resize_size[0] - H) // 2
+    left = (resize_size[1] - W) // 2
+
+    image = torch.nn.functional.interpolate(
+        image, resize_size, mode="area", antialias=False
+    )
+    image = TT.functional.crop(image, top=top, left=left, height=H, width=W)
+
+    if display:
+        numpy_img = np.transpose(image[0].numpy(), (1, 2, 0))
+        pil_image = Image.fromarray((numpy_img * 255).astype(np.uint8))
+        st.image(pil_image)
+    return image.to(device) * 2.0 - 1.0
+
+
+def save_video_as_grid_and_mp4(
+    video_batch: torch.Tensor, save_path: str, T: int, fps: int = 5
+):
+    os.makedirs(save_path, exist_ok=True)
+    base_count = len(glob(os.path.join(save_path, "*.mp4")))
+
+    video_batch = rearrange(video_batch, "(b t) c h w -> b t c h w", t=T)
+    video_batch = embed_watermark(video_batch)
+    for vid in video_batch:
+        save_image(vid, fp=os.path.join(save_path, f"{base_count:06d}.png"), nrow=4)
+
+        video_path = os.path.join(save_path, f"{base_count:06d}.mp4")
+
+        writer = cv2.VideoWriter(
+            video_path,
+            cv2.VideoWriter_fourcc(*"MP4V"),
+            fps,
+            (vid.shape[-1], vid.shape[-2]),
+        )
+
+        vid = (
+            (rearrange(vid, "t c h w -> t h w c") * 255).cpu().numpy().astype(np.uint8)
+        )
+        for frame in vid:
+            frame = cv2.cvtColor(frame, cv2.COLOR_RGB2BGR)
+            writer.write(frame)
+
+        writer.release()
+
+        video_path_h264 = video_path[:-4] + "_h264.mp4"
+        os.system(f"ffmpeg -i {video_path} -c:v libx264 {video_path_h264}")
+
+        with open(video_path_h264, "rb") as f:
+            video_bytes = f.read()
+        st.video(video_bytes)
+
+        base_count += 1

MindEye_Imagery/src/generative_models/scripts/demo/turbo.py

@@ -0,0 +1,223 @@
+from streamlit_helpers import *
+from st_keyup import st_keyup
+from sgm.modules.diffusionmodules.sampling import EulerAncestralSampler
+
+VERSION2SPECS = {
+    "SDXL-Turbo": {
+        "H": 512,
+        "W": 512,
+        "C": 4,
+        "f": 8,
+        "is_legacy": False,
+        "config": "configs/inference/sd_xl_base.yaml",
+        "ckpt": "checkpoints/sd_xl_turbo_1.0.safetensors",
+    },
+    "SD-Turbo": {
+        "H": 512,
+        "W": 512,
+        "C": 4,
+        "f": 8,
+        "is_legacy": False,
+        "config": "configs/inference/sd_2_1.yaml",
+        "ckpt": "checkpoints/sd_turbo.safetensors",
+    },
+}
+
+
+class SubstepSampler(EulerAncestralSampler):
+    def __init__(self, n_sample_steps=1, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.n_sample_steps = n_sample_steps
+        self.steps_subset = [0, 100, 200, 300, 1000]
+
+    def prepare_sampling_loop(self, x, cond, uc=None, num_steps=None):
+        sigmas = self.discretization(
+            self.num_steps if num_steps is None else num_steps, device=self.device
+        )
+        sigmas = sigmas[
+            self.steps_subset[: self.n_sample_steps] + self.steps_subset[-1:]
+        ]
+        uc = cond
+        x *= torch.sqrt(1.0 + sigmas[0] ** 2.0)
+        num_sigmas = len(sigmas)
+        s_in = x.new_ones([x.shape[0]])
+        return x, s_in, sigmas, num_sigmas, cond, uc
+
+
+def seeded_randn(shape, seed):
+    randn = np.random.RandomState(seed).randn(*shape)
+    randn = torch.from_numpy(randn).to(device="cuda", dtype=torch.float32)
+    return randn
+
+
+class SeededNoise:
+    def __init__(self, seed):
+        self.seed = seed
+
+    def __call__(self, x):
+        self.seed = self.seed + 1
+        return seeded_randn(x.shape, self.seed)
+
+
+def init_embedder_options(keys, init_dict, prompt=None, negative_prompt=None):
+    value_dict = {}
+    for key in keys:
+        if key == "txt":
+            value_dict["prompt"] = prompt
+            value_dict["negative_prompt"] = ""
+
+        if key == "original_size_as_tuple":
+            orig_width = init_dict["orig_width"]
+            orig_height = init_dict["orig_height"]
+
+            value_dict["orig_width"] = orig_width
+            value_dict["orig_height"] = orig_height
+
+        if key == "crop_coords_top_left":
+            crop_coord_top = 0
+            crop_coord_left = 0
+
+            value_dict["crop_coords_top"] = crop_coord_top
+            value_dict["crop_coords_left"] = crop_coord_left
+
+        if key == "aesthetic_score":
+            value_dict["aesthetic_score"] = 6.0
+            value_dict["negative_aesthetic_score"] = 2.5
+
+        if key == "target_size_as_tuple":
+            value_dict["target_width"] = init_dict["target_width"]
+            value_dict["target_height"] = init_dict["target_height"]
+
+    return value_dict
+
+
+def sample(
+    model,
+    sampler,
+    prompt="A lush garden with oversized flowers and vibrant colors, inhabited by miniature animals.",
+    H=1024,
+    W=1024,
+    seed=0,
+    filter=None,
+):
+    F = 8
+    C = 4
+    shape = (1, C, H // F, W // F)
+
+    value_dict = init_embedder_options(
+        keys=get_unique_embedder_keys_from_conditioner(model.conditioner),
+        init_dict={
+            "orig_width": W,
+            "orig_height": H,
+            "target_width": W,
+            "target_height": H,
+        },
+        prompt=prompt,
+    )
+
+    if seed is None:
+        seed = torch.seed()
+    precision_scope = autocast
+    with torch.no_grad():
+        with precision_scope("cuda"):
+            batch, batch_uc = get_batch(
+                get_unique_embedder_keys_from_conditioner(model.conditioner),
+                value_dict,
+                [1],
+            )
+            c = model.conditioner(batch)
+            uc = None
+            randn = seeded_randn(shape, seed)
+
+            def denoiser(input, sigma, c):
+                return model.denoiser(
+                    model.model,
+                    input,
+                    sigma,
+                    c,
+                )
+
+            samples_z = sampler(denoiser, randn, cond=c, uc=uc)
+            samples_x = model.decode_first_stage(samples_z)
+            samples = torch.clamp((samples_x + 1.0) / 2.0, min=0.0, max=1.0)
+            if filter is not None:
+                samples = filter(samples)
+            samples = (
+                (255 * samples)
+                .to(dtype=torch.uint8)
+                .permute(0, 2, 3, 1)
+                .detach()
+                .cpu()
+                .numpy()
+            )
+    return samples
+
+
+def v_spacer(height) -> None:
+    for _ in range(height):
+        st.write("\n")
+
+
+if __name__ == "__main__":
+    st.title("Turbo")
+
+    head_cols = st.columns([1, 1, 1])
+    with head_cols[0]:
+        version = st.selectbox("Model Version", list(VERSION2SPECS.keys()), 0)
+        version_dict = VERSION2SPECS[version]
+
+    with head_cols[1]:
+        v_spacer(2)
+        if st.checkbox("Load Model"):
+            mode = "txt2img"
+        else:
+            mode = "skip"
+
+    if mode != "skip":
+        state = init_st(version_dict, load_filter=True)
+        if state["msg"]:
+            st.info(state["msg"])
+        model = state["model"]
+        load_model(model)
+
+    # seed
+    if "seed" not in st.session_state:
+        st.session_state.seed = 0
+
+    def increment_counter():
+        st.session_state.seed += 1
+
+    def decrement_counter():
+        if st.session_state.seed > 0:
+            st.session_state.seed -= 1
+
+    with head_cols[2]:
+        n_steps = st.number_input(label="number of steps", min_value=1, max_value=4)
+    
+    sampler = SubstepSampler(
+        n_sample_steps=1,
+        num_steps=1000,
+        eta=1.0,
+        discretization_config=dict(
+            target="sgm.modules.diffusionmodules.discretizer.LegacyDDPMDiscretization"
+        ),
+    )
+    sampler.n_sample_steps = n_steps
+    default_prompt = "A cinematic shot of a baby racoon wearing an intricate italian priest robe."
+    prompt = st_keyup("Enter a value", value=default_prompt, debounce=300, key="interactive_text")
+
+    cols = st.columns([1, 5, 1])
+    if mode != "skip":
+        with cols[0]:
+            v_spacer(14)
+            st.button("↩", on_click=decrement_counter)
+        with cols[2]:
+            v_spacer(14)
+            st.button("↪", on_click=increment_counter)
+
+        sampler.noise_sampler = SeededNoise(seed=st.session_state.seed)
+        out = sample(
+            model, sampler, H=512, W=512, seed=st.session_state.seed, prompt=prompt, filter=state.get("filter")
+        )
+        with cols[1]:
+            st.image(out[0])

MindEye_Imagery/src/generative_models/scripts/demo/video_sampling.py

@@ -0,0 +1,200 @@
+import os
+
+from pytorch_lightning import seed_everything
+
+from scripts.demo.streamlit_helpers import *
+
+SAVE_PATH = "outputs/demo/vid/"
+
+VERSION2SPECS = {
+    "svd": {
+        "T": 14,
+        "H": 576,
+        "W": 1024,
+        "C": 4,
+        "f": 8,
+        "config": "configs/inference/svd.yaml",
+        "ckpt": "checkpoints/svd.safetensors",
+        "options": {
+            "discretization": 1,
+            "cfg": 2.5,
+            "sigma_min": 0.002,
+            "sigma_max": 700.0,
+            "rho": 7.0,
+            "guider": 2,
+            "force_uc_zero_embeddings": ["cond_frames", "cond_frames_without_noise"],
+            "num_steps": 25,
+        },
+    },
+    "svd_image_decoder": {
+        "T": 14,
+        "H": 576,
+        "W": 1024,
+        "C": 4,
+        "f": 8,
+        "config": "configs/inference/svd_image_decoder.yaml",
+        "ckpt": "checkpoints/svd_image_decoder.safetensors",
+        "options": {
+            "discretization": 1,
+            "cfg": 2.5,
+            "sigma_min": 0.002,
+            "sigma_max": 700.0,
+            "rho": 7.0,
+            "guider": 2,
+            "force_uc_zero_embeddings": ["cond_frames", "cond_frames_without_noise"],
+            "num_steps": 25,
+        },
+    },
+    "svd_xt": {
+        "T": 25,
+        "H": 576,
+        "W": 1024,
+        "C": 4,
+        "f": 8,
+        "config": "configs/inference/svd.yaml",
+        "ckpt": "checkpoints/svd_xt.safetensors",
+        "options": {
+            "discretization": 1,
+            "cfg": 3.0,
+            "min_cfg": 1.5,
+            "sigma_min": 0.002,
+            "sigma_max": 700.0,
+            "rho": 7.0,
+            "guider": 2,
+            "force_uc_zero_embeddings": ["cond_frames", "cond_frames_without_noise"],
+            "num_steps": 30,
+            "decoding_t": 14,
+        },
+    },
+    "svd_xt_image_decoder": {
+        "T": 25,
+        "H": 576,
+        "W": 1024,
+        "C": 4,
+        "f": 8,
+        "config": "configs/inference/svd_image_decoder.yaml",
+        "ckpt": "checkpoints/svd_xt_image_decoder.safetensors",
+        "options": {
+            "discretization": 1,
+            "cfg": 3.0,
+            "min_cfg": 1.5,
+            "sigma_min": 0.002,
+            "sigma_max": 700.0,
+            "rho": 7.0,
+            "guider": 2,
+            "force_uc_zero_embeddings": ["cond_frames", "cond_frames_without_noise"],
+            "num_steps": 30,
+            "decoding_t": 14,
+        },
+    },
+}
+
+
+if __name__ == "__main__":
+    st.title("Stable Video Diffusion")
+    version = st.selectbox(
+        "Model Version",
+        [k for k in VERSION2SPECS.keys()],
+        0,
+    )
+    version_dict = VERSION2SPECS[version]
+    if st.checkbox("Load Model"):
+        mode = "img2vid"
+    else:
+        mode = "skip"
+
+    H = st.sidebar.number_input(
+        "H", value=version_dict["H"], min_value=64, max_value=2048
+    )
+    W = st.sidebar.number_input(
+        "W", value=version_dict["W"], min_value=64, max_value=2048
+    )
+    T = st.sidebar.number_input(
+        "T", value=version_dict["T"], min_value=0, max_value=128
+    )
+    C = version_dict["C"]
+    F = version_dict["f"]
+    options = version_dict["options"]
+
+    if mode != "skip":
+        state = init_st(version_dict, load_filter=True)
+        if state["msg"]:
+            st.info(state["msg"])
+        model = state["model"]
+
+        ukeys = set(
+            get_unique_embedder_keys_from_conditioner(state["model"].conditioner)
+        )
+
+        value_dict = init_embedder_options(
+            ukeys,
+            {},
+        )
+
+        value_dict["image_only_indicator"] = 0
+
+        if mode == "img2vid":
+            img = load_img_for_prediction(W, H)
+            cond_aug = st.number_input(
+                "Conditioning augmentation:", value=0.02, min_value=0.0
+            )
+            value_dict["cond_frames_without_noise"] = img
+            value_dict["cond_frames"] = img + cond_aug * torch.randn_like(img)
+            value_dict["cond_aug"] = cond_aug
+
+        seed = st.sidebar.number_input(
+            "seed", value=23, min_value=0, max_value=int(1e9)
+        )
+        seed_everything(seed)
+
+        save_locally, save_path = init_save_locally(
+            os.path.join(SAVE_PATH, version), init_value=True
+        )
+
+        options["num_frames"] = T
+
+        sampler, num_rows, num_cols = init_sampling(options=options)
+        num_samples = num_rows * num_cols
+
+        decoding_t = st.number_input(
+            "Decode t frames at a time (set small if you are low on VRAM)",
+            value=options.get("decoding_t", T),
+            min_value=1,
+            max_value=int(1e9),
+        )
+
+        if st.checkbox("Overwrite fps in mp4 generator", False):
+            saving_fps = st.number_input(
+                f"saving video at fps:", value=value_dict["fps"], min_value=1
+            )
+        else:
+            saving_fps = value_dict["fps"]
+
+        if st.button("Sample"):
+            out = do_sample(
+                model,
+                sampler,
+                value_dict,
+                num_samples,
+                H,
+                W,
+                C,
+                F,
+                T=T,
+                batch2model_input=["num_video_frames", "image_only_indicator"],
+                force_uc_zero_embeddings=options.get("force_uc_zero_embeddings", None),
+                force_cond_zero_embeddings=options.get(
+                    "force_cond_zero_embeddings", None
+                ),
+                return_latents=False,
+                decoding_t=decoding_t,
+            )
+
+            if isinstance(out, (tuple, list)):
+                samples, samples_z = out
+            else:
+                samples = out
+                samples_z = None
+
+            if save_locally:
+                save_video_as_grid_and_mp4(samples, save_path, T, fps=saving_fps)

MindEye_Imagery/src/generative_models/scripts/sampling/configs/svd.yaml

@@ -0,0 +1,146 @@
+model:
+  target: sgm.models.diffusion.DiffusionEngine
+  params:
+    scale_factor: 0.18215
+    disable_first_stage_autocast: True
+    ckpt_path: checkpoints/svd.safetensors
+
+    denoiser_config:
+      target: sgm.modules.diffusionmodules.denoiser.Denoiser
+      params:
+        scaling_config:
+          target: sgm.modules.diffusionmodules.denoiser_scaling.VScalingWithEDMcNoise
+
+    network_config:
+      target: sgm.modules.diffusionmodules.video_model.VideoUNet
+      params:
+        adm_in_channels: 768
+        num_classes: sequential
+        use_checkpoint: True
+        in_channels: 8
+        out_channels: 4
+        model_channels: 320
+        attention_resolutions: [4, 2, 1]
+        num_res_blocks: 2
+        channel_mult: [1, 2, 4, 4]
+        num_head_channels: 64
+        use_linear_in_transformer: True
+        transformer_depth: 1
+        context_dim: 1024
+        spatial_transformer_attn_type: softmax-xformers
+        extra_ff_mix_layer: True
+        use_spatial_context: True
+        merge_strategy: learned_with_images
+        video_kernel_size: [3, 1, 1]
+
+    conditioner_config:
+      target: sgm.modules.GeneralConditioner
+      params:
+        emb_models:
+        - is_trainable: False
+          input_key: cond_frames_without_noise
+          target: sgm.modules.encoders.modules.FrozenOpenCLIPImagePredictionEmbedder
+          params:
+            n_cond_frames: 1
+            n_copies: 1
+            open_clip_embedding_config:
+              target: sgm.modules.encoders.modules.FrozenOpenCLIPImageEmbedder
+              params:
+                freeze: True
+
+        - input_key: fps_id
+          is_trainable: False
+          target: sgm.modules.encoders.modules.ConcatTimestepEmbedderND
+          params:
+            outdim: 256
+
+        - input_key: motion_bucket_id
+          is_trainable: False
+          target: sgm.modules.encoders.modules.ConcatTimestepEmbedderND
+          params:
+            outdim: 256
+
+        - input_key: cond_frames
+          is_trainable: False
+          target: sgm.modules.encoders.modules.VideoPredictionEmbedderWithEncoder
+          params:
+            disable_encoder_autocast: True
+            n_cond_frames: 1
+            n_copies: 1
+            is_ae: True
+            encoder_config:
+              target: sgm.models.autoencoder.AutoencoderKLModeOnly
+              params:
+                embed_dim: 4
+                monitor: val/rec_loss
+                ddconfig:
+                  attn_type: vanilla-xformers
+                  double_z: True
+                  z_channels: 4
+                  resolution: 256
+                  in_channels: 3
+                  out_ch: 3
+                  ch: 128
+                  ch_mult: [1, 2, 4, 4]
+                  num_res_blocks: 2
+                  attn_resolutions: []
+                  dropout: 0.0
+                lossconfig:
+                  target: torch.nn.Identity
+
+        - input_key: cond_aug
+          is_trainable: False
+          target: sgm.modules.encoders.modules.ConcatTimestepEmbedderND
+          params:
+            outdim: 256
+
+    first_stage_config:
+      target: sgm.models.autoencoder.AutoencodingEngine
+      params:
+        loss_config:
+          target: torch.nn.Identity
+        regularizer_config:
+          target: sgm.modules.autoencoding.regularizers.DiagonalGaussianRegularizer
+        encoder_config: 
+          target: sgm.modules.diffusionmodules.model.Encoder
+          params:
+            attn_type: vanilla
+            double_z: True
+            z_channels: 4
+            resolution: 256
+            in_channels: 3
+            out_ch: 3
+            ch: 128
+            ch_mult: [1, 2, 4, 4]
+            num_res_blocks: 2
+            attn_resolutions: []
+            dropout: 0.0
+        decoder_config:
+          target: sgm.modules.autoencoding.temporal_ae.VideoDecoder
+          params:
+            attn_type: vanilla
+            double_z: True
+            z_channels: 4
+            resolution: 256
+            in_channels: 3
+            out_ch: 3
+            ch: 128
+            ch_mult: [1, 2, 4, 4]
+            num_res_blocks: 2
+            attn_resolutions: []
+            dropout: 0.0
+            video_kernel_size: [3, 1, 1]
+
+    sampler_config:
+      target: sgm.modules.diffusionmodules.sampling.EulerEDMSampler
+      params:
+        discretization_config:
+          target: sgm.modules.diffusionmodules.discretizer.EDMDiscretization
+          params:
+            sigma_max: 700.0
+
+        guider_config:
+          target: sgm.modules.diffusionmodules.guiders.LinearPredictionGuider
+          params:
+            max_scale: 2.5
+            min_scale: 1.0

MindEye_Imagery/src/generative_models/scripts/sampling/configs/svd_image_decoder.yaml

@@ -0,0 +1,129 @@
+model:
+  target: sgm.models.diffusion.DiffusionEngine
+  params:
+    scale_factor: 0.18215
+    disable_first_stage_autocast: True
+    ckpt_path: checkpoints/svd_image_decoder.safetensors
+
+    denoiser_config:
+      target: sgm.modules.diffusionmodules.denoiser.Denoiser
+      params:
+        scaling_config:
+          target: sgm.modules.diffusionmodules.denoiser_scaling.VScalingWithEDMcNoise
+
+    network_config:
+      target: sgm.modules.diffusionmodules.video_model.VideoUNet
+      params:
+        adm_in_channels: 768
+        num_classes: sequential
+        use_checkpoint: True
+        in_channels: 8
+        out_channels: 4
+        model_channels: 320
+        attention_resolutions: [4, 2, 1]
+        num_res_blocks: 2
+        channel_mult: [1, 2, 4, 4]
+        num_head_channels: 64
+        use_linear_in_transformer: True
+        transformer_depth: 1
+        context_dim: 1024
+        spatial_transformer_attn_type: softmax-xformers
+        extra_ff_mix_layer: True
+        use_spatial_context: True
+        merge_strategy: learned_with_images
+        video_kernel_size: [3, 1, 1]
+
+    conditioner_config:
+      target: sgm.modules.GeneralConditioner
+      params:
+        emb_models:
+        - is_trainable: False
+          input_key: cond_frames_without_noise
+          target: sgm.modules.encoders.modules.FrozenOpenCLIPImagePredictionEmbedder
+          params:
+            n_cond_frames: 1
+            n_copies: 1
+            open_clip_embedding_config:
+              target: sgm.modules.encoders.modules.FrozenOpenCLIPImageEmbedder
+              params:
+                freeze: True
+
+        - input_key: fps_id
+          is_trainable: False
+          target: sgm.modules.encoders.modules.ConcatTimestepEmbedderND
+          params:
+            outdim: 256
+
+        - input_key: motion_bucket_id
+          is_trainable: False
+          target: sgm.modules.encoders.modules.ConcatTimestepEmbedderND
+          params:
+            outdim: 256
+
+        - input_key: cond_frames
+          is_trainable: False
+          target: sgm.modules.encoders.modules.VideoPredictionEmbedderWithEncoder
+          params:
+            disable_encoder_autocast: True
+            n_cond_frames: 1
+            n_copies: 1
+            is_ae: True
+            encoder_config:
+              target: sgm.models.autoencoder.AutoencoderKLModeOnly
+              params:
+                embed_dim: 4
+                monitor: val/rec_loss
+                ddconfig:
+                  attn_type: vanilla-xformers
+                  double_z: True
+                  z_channels: 4
+                  resolution: 256
+                  in_channels: 3
+                  out_ch: 3
+                  ch: 128
+                  ch_mult: [1, 2, 4, 4]
+                  num_res_blocks: 2
+                  attn_resolutions: []
+                  dropout: 0.0
+                lossconfig:
+                  target: torch.nn.Identity
+
+        - input_key: cond_aug
+          is_trainable: False
+          target: sgm.modules.encoders.modules.ConcatTimestepEmbedderND
+          params:
+            outdim: 256
+
+    first_stage_config:
+      target: sgm.models.autoencoder.AutoencoderKL
+      params:
+        embed_dim: 4
+        monitor: val/rec_loss
+        ddconfig:
+          attn_type: vanilla-xformers
+          double_z: True
+          z_channels: 4
+          resolution: 256
+          in_channels: 3
+          out_ch: 3
+          ch: 128
+          ch_mult: [1, 2, 4, 4]
+          num_res_blocks: 2
+          attn_resolutions: []
+          dropout: 0.0
+        lossconfig:
+          target: torch.nn.Identity
+
+    sampler_config:
+      target: sgm.modules.diffusionmodules.sampling.EulerEDMSampler
+      params:
+        discretization_config:
+          target: sgm.modules.diffusionmodules.discretizer.EDMDiscretization
+          params:
+            sigma_max: 700.0
+
+        guider_config:
+          target: sgm.modules.diffusionmodules.guiders.LinearPredictionGuider
+          params:
+            max_scale: 2.5
+            min_scale: 1.0

MindEye_Imagery/src/generative_models/scripts/sampling/configs/svd_xt.yaml

@@ -0,0 +1,146 @@
+model:
+  target: sgm.models.diffusion.DiffusionEngine
+  params:
+    scale_factor: 0.18215
+    disable_first_stage_autocast: True
+    ckpt_path: checkpoints/svd_xt.safetensors
+
+    denoiser_config:
+      target: sgm.modules.diffusionmodules.denoiser.Denoiser
+      params:
+        scaling_config:
+          target: sgm.modules.diffusionmodules.denoiser_scaling.VScalingWithEDMcNoise
+
+    network_config:
+      target: sgm.modules.diffusionmodules.video_model.VideoUNet
+      params:
+        adm_in_channels: 768
+        num_classes: sequential
+        use_checkpoint: True
+        in_channels: 8
+        out_channels: 4
+        model_channels: 320
+        attention_resolutions: [4, 2, 1]
+        num_res_blocks: 2
+        channel_mult: [1, 2, 4, 4]
+        num_head_channels: 64
+        use_linear_in_transformer: True
+        transformer_depth: 1
+        context_dim: 1024
+        spatial_transformer_attn_type: softmax-xformers
+        extra_ff_mix_layer: True
+        use_spatial_context: True
+        merge_strategy: learned_with_images
+        video_kernel_size: [3, 1, 1]
+
+    conditioner_config:
+      target: sgm.modules.GeneralConditioner
+      params:
+        emb_models:
+        - is_trainable: False
+          input_key: cond_frames_without_noise
+          target: sgm.modules.encoders.modules.FrozenOpenCLIPImagePredictionEmbedder
+          params:
+            n_cond_frames: 1
+            n_copies: 1
+            open_clip_embedding_config:
+              target: sgm.modules.encoders.modules.FrozenOpenCLIPImageEmbedder
+              params:
+                freeze: True
+
+        - input_key: fps_id
+          is_trainable: False
+          target: sgm.modules.encoders.modules.ConcatTimestepEmbedderND
+          params:
+            outdim: 256
+
+        - input_key: motion_bucket_id
+          is_trainable: False
+          target: sgm.modules.encoders.modules.ConcatTimestepEmbedderND
+          params:
+            outdim: 256
+
+        - input_key: cond_frames
+          is_trainable: False
+          target: sgm.modules.encoders.modules.VideoPredictionEmbedderWithEncoder
+          params:
+            disable_encoder_autocast: True
+            n_cond_frames: 1
+            n_copies: 1
+            is_ae: True
+            encoder_config:
+              target: sgm.models.autoencoder.AutoencoderKLModeOnly
+              params:
+                embed_dim: 4
+                monitor: val/rec_loss
+                ddconfig:
+                  attn_type: vanilla-xformers
+                  double_z: True
+                  z_channels: 4
+                  resolution: 256
+                  in_channels: 3
+                  out_ch: 3
+                  ch: 128
+                  ch_mult: [1, 2, 4, 4]
+                  num_res_blocks: 2
+                  attn_resolutions: []
+                  dropout: 0.0
+                lossconfig:
+                  target: torch.nn.Identity
+
+        - input_key: cond_aug
+          is_trainable: False
+          target: sgm.modules.encoders.modules.ConcatTimestepEmbedderND
+          params:
+            outdim: 256
+
+    first_stage_config:
+      target: sgm.models.autoencoder.AutoencodingEngine
+      params:
+        loss_config:
+          target: torch.nn.Identity
+        regularizer_config:
+          target: sgm.modules.autoencoding.regularizers.DiagonalGaussianRegularizer
+        encoder_config: 
+          target: sgm.modules.diffusionmodules.model.Encoder
+          params:
+            attn_type: vanilla
+            double_z: True
+            z_channels: 4
+            resolution: 256
+            in_channels: 3
+            out_ch: 3
+            ch: 128
+            ch_mult: [1, 2, 4, 4]
+            num_res_blocks: 2
+            attn_resolutions: []
+            dropout: 0.0
+        decoder_config:
+          target: sgm.modules.autoencoding.temporal_ae.VideoDecoder
+          params:
+            attn_type: vanilla
+            double_z: True
+            z_channels: 4
+            resolution: 256
+            in_channels: 3
+            out_ch: 3
+            ch: 128
+            ch_mult: [1, 2, 4, 4]
+            num_res_blocks: 2
+            attn_resolutions: []
+            dropout: 0.0
+            video_kernel_size: [3, 1, 1]
+
+    sampler_config:
+      target: sgm.modules.diffusionmodules.sampling.EulerEDMSampler
+      params:
+        discretization_config:
+          target: sgm.modules.diffusionmodules.discretizer.EDMDiscretization
+          params:
+            sigma_max: 700.0
+
+        guider_config:
+          target: sgm.modules.diffusionmodules.guiders.LinearPredictionGuider
+          params:
+            max_scale: 3.0
+            min_scale: 1.5

MindEye_Imagery/src/generative_models/scripts/sampling/configs/svd_xt_image_decoder.yaml

@@ -0,0 +1,129 @@
+model:
+  target: sgm.models.diffusion.DiffusionEngine
+  params:
+    scale_factor: 0.18215
+    disable_first_stage_autocast: True
+    ckpt_path: checkpoints/svd_xt_image_decoder.safetensors
+
+    denoiser_config:
+      target: sgm.modules.diffusionmodules.denoiser.Denoiser
+      params:
+        scaling_config:
+          target: sgm.modules.diffusionmodules.denoiser_scaling.VScalingWithEDMcNoise
+
+    network_config:
+      target: sgm.modules.diffusionmodules.video_model.VideoUNet
+      params:
+        adm_in_channels: 768
+        num_classes: sequential
+        use_checkpoint: True
+        in_channels: 8
+        out_channels: 4
+        model_channels: 320
+        attention_resolutions: [4, 2, 1]
+        num_res_blocks: 2
+        channel_mult: [1, 2, 4, 4]
+        num_head_channels: 64
+        use_linear_in_transformer: True
+        transformer_depth: 1
+        context_dim: 1024
+        spatial_transformer_attn_type: softmax-xformers
+        extra_ff_mix_layer: True
+        use_spatial_context: True
+        merge_strategy: learned_with_images
+        video_kernel_size: [3, 1, 1]
+
+    conditioner_config:
+      target: sgm.modules.GeneralConditioner
+      params:
+        emb_models:
+        - is_trainable: False
+          input_key: cond_frames_without_noise
+          target: sgm.modules.encoders.modules.FrozenOpenCLIPImagePredictionEmbedder
+          params:
+            n_cond_frames: 1
+            n_copies: 1
+            open_clip_embedding_config:
+              target: sgm.modules.encoders.modules.FrozenOpenCLIPImageEmbedder
+              params:
+                freeze: True
+
+        - input_key: fps_id
+          is_trainable: False
+          target: sgm.modules.encoders.modules.ConcatTimestepEmbedderND
+          params:
+            outdim: 256
+
+        - input_key: motion_bucket_id
+          is_trainable: False
+          target: sgm.modules.encoders.modules.ConcatTimestepEmbedderND
+          params:
+            outdim: 256
+
+        - input_key: cond_frames
+          is_trainable: False
+          target: sgm.modules.encoders.modules.VideoPredictionEmbedderWithEncoder
+          params:
+            disable_encoder_autocast: True
+            n_cond_frames: 1
+            n_copies: 1
+            is_ae: True
+            encoder_config:
+              target: sgm.models.autoencoder.AutoencoderKLModeOnly
+              params:
+                embed_dim: 4
+                monitor: val/rec_loss
+                ddconfig:
+                  attn_type: vanilla-xformers
+                  double_z: True
+                  z_channels: 4
+                  resolution: 256
+                  in_channels: 3
+                  out_ch: 3
+                  ch: 128
+                  ch_mult: [1, 2, 4, 4]
+                  num_res_blocks: 2
+                  attn_resolutions: []
+                  dropout: 0.0
+                lossconfig:
+                  target: torch.nn.Identity
+
+        - input_key: cond_aug
+          is_trainable: False
+          target: sgm.modules.encoders.modules.ConcatTimestepEmbedderND
+          params:
+            outdim: 256
+
+    first_stage_config:
+      target: sgm.models.autoencoder.AutoencoderKL
+      params:
+        embed_dim: 4
+        monitor: val/rec_loss
+        ddconfig:
+          attn_type: vanilla-xformers
+          double_z: True
+          z_channels: 4
+          resolution: 256
+          in_channels: 3
+          out_ch: 3
+          ch: 128
+          ch_mult: [1, 2, 4, 4]
+          num_res_blocks: 2
+          attn_resolutions: []
+          dropout: 0.0
+        lossconfig:
+          target: torch.nn.Identity
+
+    sampler_config:
+      target: sgm.modules.diffusionmodules.sampling.EulerEDMSampler
+      params:
+        discretization_config:
+          target: sgm.modules.diffusionmodules.discretizer.EDMDiscretization
+          params:
+            sigma_max: 700.0
+
+        guider_config:
+          target: sgm.modules.diffusionmodules.guiders.LinearPredictionGuider
+          params:
+            max_scale: 3.0
+            min_scale: 1.5

MindEye_Imagery/src/generative_models/scripts/sampling/simple_video_sample.py

@@ -0,0 +1,278 @@
+import math
+import os
+from glob import glob
+from pathlib import Path
+from typing import Optional
+
+import cv2
+import numpy as np
+import torch
+from einops import rearrange, repeat
+from fire import Fire
+from omegaconf import OmegaConf
+from PIL import Image
+from torchvision.transforms import ToTensor
+
+from scripts.util.detection.nsfw_and_watermark_dectection import \
+    DeepFloydDataFiltering
+from sgm.inference.helpers import embed_watermark
+from sgm.util import default, instantiate_from_config
+
+
+def sample(
+    input_path: str = "assets/test_image.png",  # Can either be image file or folder with image files
+    num_frames: Optional[int] = None,
+    num_steps: Optional[int] = None,
+    version: str = "svd",
+    fps_id: int = 6,
+    motion_bucket_id: int = 127,
+    cond_aug: float = 0.02,
+    seed: int = 23,
+    decoding_t: int = 14,  # Number of frames decoded at a time! This eats most VRAM. Reduce if necessary.
+    device: str = "cuda",
+    output_folder: Optional[str] = None,
+):
+    """
+    Simple script to generate a single sample conditioned on an image `input_path` or multiple images, one for each
+    image file in folder `input_path`. If you run out of VRAM, try decreasing `decoding_t`.
+    """
+
+    if version == "svd":
+        num_frames = default(num_frames, 14)
+        num_steps = default(num_steps, 25)
+        output_folder = default(output_folder, "outputs/simple_video_sample/svd/")
+        model_config = "scripts/sampling/configs/svd.yaml"
+    elif version == "svd_xt":
+        num_frames = default(num_frames, 25)
+        num_steps = default(num_steps, 30)
+        output_folder = default(output_folder, "outputs/simple_video_sample/svd_xt/")
+        model_config = "scripts/sampling/configs/svd_xt.yaml"
+    elif version == "svd_image_decoder":
+        num_frames = default(num_frames, 14)
+        num_steps = default(num_steps, 25)
+        output_folder = default(
+            output_folder, "outputs/simple_video_sample/svd_image_decoder/"
+        )
+        model_config = "scripts/sampling/configs/svd_image_decoder.yaml"
+    elif version == "svd_xt_image_decoder":
+        num_frames = default(num_frames, 25)
+        num_steps = default(num_steps, 30)
+        output_folder = default(
+            output_folder, "outputs/simple_video_sample/svd_xt_image_decoder/"
+        )
+        model_config = "scripts/sampling/configs/svd_xt_image_decoder.yaml"
+    else:
+        raise ValueError(f"Version {version} does not exist.")
+
+    model, filter = load_model(
+        model_config,
+        device,
+        num_frames,
+        num_steps,
+    )
+    torch.manual_seed(seed)
+
+    path = Path(input_path)
+    all_img_paths = []
+    if path.is_file():
+        if any([input_path.endswith(x) for x in ["jpg", "jpeg", "png"]]):
+            all_img_paths = [input_path]
+        else:
+            raise ValueError("Path is not valid image file.")
+    elif path.is_dir():
+        all_img_paths = sorted(
+            [
+                f
+                for f in path.iterdir()
+                if f.is_file() and f.suffix.lower() in [".jpg", ".jpeg", ".png"]
+            ]
+        )
+        if len(all_img_paths) == 0:
+            raise ValueError("Folder does not contain any images.")
+    else:
+        raise ValueError
+
+    for input_img_path in all_img_paths:
+        with Image.open(input_img_path) as image:
+            if image.mode == "RGBA":
+                image = image.convert("RGB")
+            w, h = image.size
+
+            if h % 64 != 0 or w % 64 != 0:
+                width, height = map(lambda x: x - x % 64, (w, h))
+                image = image.resize((width, height))
+                print(
+                    f"WARNING: Your image is of size {h}x{w} which is not divisible by 64. We are resizing to {height}x{width}!"
+                )
+
+            image = ToTensor()(image)
+            image = image * 2.0 - 1.0
+
+        image = image.unsqueeze(0).to(device)
+        H, W = image.shape[2:]
+        assert image.shape[1] == 3
+        F = 8
+        C = 4
+        shape = (num_frames, C, H // F, W // F)
+        if (H, W) != (576, 1024):
+            print(
+                "WARNING: The conditioning frame you provided is not 576x1024. This leads to suboptimal performance as model was only trained on 576x1024. Consider increasing `cond_aug`."
+            )
+        if motion_bucket_id > 255:
+            print(
+                "WARNING: High motion bucket! This may lead to suboptimal performance."
+            )
+
+        if fps_id < 5:
+            print("WARNING: Small fps value! This may lead to suboptimal performance.")
+
+        if fps_id > 30:
+            print("WARNING: Large fps value! This may lead to suboptimal performance.")
+
+        value_dict = {}
+        value_dict["motion_bucket_id"] = motion_bucket_id
+        value_dict["fps_id"] = fps_id
+        value_dict["cond_aug"] = cond_aug
+        value_dict["cond_frames_without_noise"] = image
+        value_dict["cond_frames"] = image + cond_aug * torch.randn_like(image)
+        value_dict["cond_aug"] = cond_aug
+
+        with torch.no_grad():
+            with torch.autocast(device):
+                batch, batch_uc = get_batch(
+                    get_unique_embedder_keys_from_conditioner(model.conditioner),
+                    value_dict,
+                    [1, num_frames],
+                    T=num_frames,
+                    device=device,
+                )
+                c, uc = model.conditioner.get_unconditional_conditioning(
+                    batch,
+                    batch_uc=batch_uc,
+                    force_uc_zero_embeddings=[
+                        "cond_frames",
+                        "cond_frames_without_noise",
+                    ],
+                )
+
+                for k in ["crossattn", "concat"]:
+                    uc[k] = repeat(uc[k], "b ... -> b t ...", t=num_frames)
+                    uc[k] = rearrange(uc[k], "b t ... -> (b t) ...", t=num_frames)
+                    c[k] = repeat(c[k], "b ... -> b t ...", t=num_frames)
+                    c[k] = rearrange(c[k], "b t ... -> (b t) ...", t=num_frames)
+
+                randn = torch.randn(shape, device=device)
+
+                additional_model_inputs = {}
+                additional_model_inputs["image_only_indicator"] = torch.zeros(
+                    2, num_frames
+                ).to(device)
+                additional_model_inputs["num_video_frames"] = batch["num_video_frames"]
+
+                def denoiser(input, sigma, c):
+                    return model.denoiser(
+                        model.model, input, sigma, c, **additional_model_inputs
+                    )
+
+                samples_z = model.sampler(denoiser, randn, cond=c, uc=uc)
+                model.en_and_decode_n_samples_a_time = decoding_t
+                samples_x = model.decode_first_stage(samples_z)
+                samples = torch.clamp((samples_x + 1.0) / 2.0, min=0.0, max=1.0)
+
+                os.makedirs(output_folder, exist_ok=True)
+                base_count = len(glob(os.path.join(output_folder, "*.mp4")))
+                video_path = os.path.join(output_folder, f"{base_count:06d}.mp4")
+                writer = cv2.VideoWriter(
+                    video_path,
+                    cv2.VideoWriter_fourcc(*"MP4V"),
+                    fps_id + 1,
+                    (samples.shape[-1], samples.shape[-2]),
+                )
+
+                samples = embed_watermark(samples)
+                samples = filter(samples)
+                vid = (
+                    (rearrange(samples, "t c h w -> t h w c") * 255)
+                    .cpu()
+                    .numpy()
+                    .astype(np.uint8)
+                )
+                for frame in vid:
+                    frame = cv2.cvtColor(frame, cv2.COLOR_RGB2BGR)
+                    writer.write(frame)
+                writer.release()
+
+
+def get_unique_embedder_keys_from_conditioner(conditioner):
+    return list(set([x.input_key for x in conditioner.embedders]))
+
+
+def get_batch(keys, value_dict, N, T, device):
+    batch = {}
+    batch_uc = {}
+
+    for key in keys:
+        if key == "fps_id":
+            batch[key] = (
+                torch.tensor([value_dict["fps_id"]])
+                .to(device)
+                .repeat(int(math.prod(N)))
+            )
+        elif key == "motion_bucket_id":
+            batch[key] = (
+                torch.tensor([value_dict["motion_bucket_id"]])
+                .to(device)
+                .repeat(int(math.prod(N)))
+            )
+        elif key == "cond_aug":
+            batch[key] = repeat(
+                torch.tensor([value_dict["cond_aug"]]).to(device),
+                "1 -> b",
+                b=math.prod(N),
+            )
+        elif key == "cond_frames":
+            batch[key] = repeat(value_dict["cond_frames"], "1 ... -> b ...", b=N[0])
+        elif key == "cond_frames_without_noise":
+            batch[key] = repeat(
+                value_dict["cond_frames_without_noise"], "1 ... -> b ...", b=N[0]
+            )
+        else:
+            batch[key] = value_dict[key]
+
+    if T is not None:
+        batch["num_video_frames"] = T
+
+    for key in batch.keys():
+        if key not in batch_uc and isinstance(batch[key], torch.Tensor):
+            batch_uc[key] = torch.clone(batch[key])
+    return batch, batch_uc
+
+
+def load_model(
+    config: str,
+    device: str,
+    num_frames: int,
+    num_steps: int,
+):
+    config = OmegaConf.load(config)
+    if device == "cuda":
+        config.model.params.conditioner_config.params.emb_models[
+            0
+        ].params.open_clip_embedding_config.params.init_device = device
+
+    config.model.params.sampler_config.params.num_steps = num_steps
+    config.model.params.sampler_config.params.guider_config.params.num_frames = (
+        num_frames
+    )
+    if device == "cuda":
+        with torch.device(device):
+            model = instantiate_from_config(config.model).to(device).eval()
+    else:
+        model = instantiate_from_config(config.model).to(device).eval()
+
+    filter = DeepFloydDataFiltering(verbose=False, device=device)
+    return model, filter
+
+
+if __name__ == "__main__":
+    Fire(sample)

MindEye_Imagery/src/generative_models/scripts/tests/attention.py

@@ -0,0 +1,319 @@
+import einops
+import torch
+import torch.nn.functional as F
+import torch.utils.benchmark as benchmark
+from torch.backends.cuda import SDPBackend
+
+from sgm.modules.attention import BasicTransformerBlock, SpatialTransformer
+
+
+def benchmark_attn():
+    # Lets define a helpful benchmarking function:
+    # https://pytorch.org/tutorials/intermediate/scaled_dot_product_attention_tutorial.html
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+
+    def benchmark_torch_function_in_microseconds(f, *args, **kwargs):
+        t0 = benchmark.Timer(
+            stmt="f(*args, **kwargs)", globals={"args": args, "kwargs": kwargs, "f": f}
+        )
+        return t0.blocked_autorange().mean * 1e6
+
+    # Lets define the hyper-parameters of our input
+    batch_size = 32
+    max_sequence_len = 1024
+    num_heads = 32
+    embed_dimension = 32
+
+    dtype = torch.float16
+
+    query = torch.rand(
+        batch_size,
+        num_heads,
+        max_sequence_len,
+        embed_dimension,
+        device=device,
+        dtype=dtype,
+    )
+    key = torch.rand(
+        batch_size,
+        num_heads,
+        max_sequence_len,
+        embed_dimension,
+        device=device,
+        dtype=dtype,
+    )
+    value = torch.rand(
+        batch_size,
+        num_heads,
+        max_sequence_len,
+        embed_dimension,
+        device=device,
+        dtype=dtype,
+    )
+
+    print(f"q/k/v shape:", query.shape, key.shape, value.shape)
+
+    # Lets explore the speed of each of the 3 implementations
+    from torch.backends.cuda import SDPBackend, sdp_kernel
+
+    # Helpful arguments mapper
+    backend_map = {
+        SDPBackend.MATH: {
+            "enable_math": True,
+            "enable_flash": False,
+            "enable_mem_efficient": False,
+        },
+        SDPBackend.FLASH_ATTENTION: {
+            "enable_math": False,
+            "enable_flash": True,
+            "enable_mem_efficient": False,
+        },
+        SDPBackend.EFFICIENT_ATTENTION: {
+            "enable_math": False,
+            "enable_flash": False,
+            "enable_mem_efficient": True,
+        },
+    }
+
+    from torch.profiler import ProfilerActivity, profile, record_function
+
+    activities = [ProfilerActivity.CPU, ProfilerActivity.CUDA]
+
+    print(
+        f"The default implementation runs in {benchmark_torch_function_in_microseconds(F.scaled_dot_product_attention, query, key, value):.3f} microseconds"
+    )
+    with profile(
+        activities=activities, record_shapes=False, profile_memory=True
+    ) as prof:
+        with record_function("Default detailed stats"):
+            for _ in range(25):
+                o = F.scaled_dot_product_attention(query, key, value)
+    print(prof.key_averages().table(sort_by="cuda_time_total", row_limit=10))
+
+    print(
+        f"The math implementation runs in {benchmark_torch_function_in_microseconds(F.scaled_dot_product_attention, query, key, value):.3f} microseconds"
+    )
+    with sdp_kernel(**backend_map[SDPBackend.MATH]):
+        with profile(
+            activities=activities, record_shapes=False, profile_memory=True
+        ) as prof:
+            with record_function("Math implmentation stats"):
+                for _ in range(25):
+                    o = F.scaled_dot_product_attention(query, key, value)
+        print(prof.key_averages().table(sort_by="cuda_time_total", row_limit=10))
+
+    with sdp_kernel(**backend_map[SDPBackend.FLASH_ATTENTION]):
+        try:
+            print(
+                f"The flash attention implementation runs in {benchmark_torch_function_in_microseconds(F.scaled_dot_product_attention, query, key, value):.3f} microseconds"
+            )
+        except RuntimeError:
+            print("FlashAttention is not supported. See warnings for reasons.")
+        with profile(
+            activities=activities, record_shapes=False, profile_memory=True
+        ) as prof:
+            with record_function("FlashAttention stats"):
+                for _ in range(25):
+                    o = F.scaled_dot_product_attention(query, key, value)
+        print(prof.key_averages().table(sort_by="cuda_time_total", row_limit=10))
+
+    with sdp_kernel(**backend_map[SDPBackend.EFFICIENT_ATTENTION]):
+        try:
+            print(
+                f"The memory efficient implementation runs in {benchmark_torch_function_in_microseconds(F.scaled_dot_product_attention, query, key, value):.3f} microseconds"
+            )
+        except RuntimeError:
+            print("EfficientAttention is not supported. See warnings for reasons.")
+        with profile(
+            activities=activities, record_shapes=False, profile_memory=True
+        ) as prof:
+            with record_function("EfficientAttention stats"):
+                for _ in range(25):
+                    o = F.scaled_dot_product_attention(query, key, value)
+        print(prof.key_averages().table(sort_by="cuda_time_total", row_limit=10))
+
+
+def run_model(model, x, context):
+    return model(x, context)
+
+
+def benchmark_transformer_blocks():
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    import torch.utils.benchmark as benchmark
+
+    def benchmark_torch_function_in_microseconds(f, *args, **kwargs):
+        t0 = benchmark.Timer(
+            stmt="f(*args, **kwargs)", globals={"args": args, "kwargs": kwargs, "f": f}
+        )
+        return t0.blocked_autorange().mean * 1e6
+
+    checkpoint = True
+    compile = False
+
+    batch_size = 32
+    h, w = 64, 64
+    context_len = 77
+    embed_dimension = 1024
+    context_dim = 1024
+    d_head = 64
+
+    transformer_depth = 4
+
+    n_heads = embed_dimension // d_head
+
+    dtype = torch.float16
+
+    model_native = SpatialTransformer(
+        embed_dimension,
+        n_heads,
+        d_head,
+        context_dim=context_dim,
+        use_linear=True,
+        use_checkpoint=checkpoint,
+        attn_type="softmax",
+        depth=transformer_depth,
+        sdp_backend=SDPBackend.FLASH_ATTENTION,
+    ).to(device)
+    model_efficient_attn = SpatialTransformer(
+        embed_dimension,
+        n_heads,
+        d_head,
+        context_dim=context_dim,
+        use_linear=True,
+        depth=transformer_depth,
+        use_checkpoint=checkpoint,
+        attn_type="softmax-xformers",
+    ).to(device)
+    if not checkpoint and compile:
+        print("compiling models")
+        model_native = torch.compile(model_native)
+        model_efficient_attn = torch.compile(model_efficient_attn)
+
+    x = torch.rand(batch_size, embed_dimension, h, w, device=device, dtype=dtype)
+    c = torch.rand(batch_size, context_len, context_dim, device=device, dtype=dtype)
+
+    from torch.profiler import ProfilerActivity, profile, record_function
+
+    activities = [ProfilerActivity.CPU, ProfilerActivity.CUDA]
+
+    with torch.autocast("cuda"):
+        print(
+            f"The native model runs in {benchmark_torch_function_in_microseconds(model_native.forward, x, c):.3f} microseconds"
+        )
+        print(
+            f"The efficientattn model runs in {benchmark_torch_function_in_microseconds(model_efficient_attn.forward, x, c):.3f} microseconds"
+        )
+
+        print(75 * "+")
+        print("NATIVE")
+        print(75 * "+")
+        torch.cuda.reset_peak_memory_stats()
+        with profile(
+            activities=activities, record_shapes=False, profile_memory=True
+        ) as prof:
+            with record_function("NativeAttention stats"):
+                for _ in range(25):
+                    model_native(x, c)
+        print(prof.key_averages().table(sort_by="cuda_time_total", row_limit=10))
+        print(torch.cuda.max_memory_allocated() * 1e-9, "GB used by native block")
+
+        print(75 * "+")
+        print("Xformers")
+        print(75 * "+")
+        torch.cuda.reset_peak_memory_stats()
+        with profile(
+            activities=activities, record_shapes=False, profile_memory=True
+        ) as prof:
+            with record_function("xformers stats"):
+                for _ in range(25):
+                    model_efficient_attn(x, c)
+        print(prof.key_averages().table(sort_by="cuda_time_total", row_limit=10))
+        print(torch.cuda.max_memory_allocated() * 1e-9, "GB used by xformers block")
+
+
+def test01():
+    # conv1x1 vs linear
+    from sgm.util import count_params
+
+    conv = torch.nn.Conv2d(3, 32, kernel_size=1).cuda()
+    print(count_params(conv))
+    linear = torch.nn.Linear(3, 32).cuda()
+    print(count_params(linear))
+
+    print(conv.weight.shape)
+
+    # use same initialization
+    linear.weight = torch.nn.Parameter(conv.weight.squeeze(-1).squeeze(-1))
+    linear.bias = torch.nn.Parameter(conv.bias)
+
+    print(linear.weight.shape)
+
+    x = torch.randn(11, 3, 64, 64).cuda()
+
+    xr = einops.rearrange(x, "b c h w -> b (h w) c").contiguous()
+    print(xr.shape)
+    out_linear = linear(xr)
+    print(out_linear.mean(), out_linear.shape)
+
+    out_conv = conv(x)
+    print(out_conv.mean(), out_conv.shape)
+    print("done with test01.\n")
+
+
+def test02():
+    # try cosine flash attention
+    import time
+
+    torch.backends.cuda.matmul.allow_tf32 = True
+    torch.backends.cudnn.allow_tf32 = True
+    torch.backends.cudnn.benchmark = True
+    print("testing cosine flash attention...")
+    DIM = 1024
+    SEQLEN = 4096
+    BS = 16
+
+    print(" softmax (vanilla) first...")
+    model = BasicTransformerBlock(
+        dim=DIM,
+        n_heads=16,
+        d_head=64,
+        dropout=0.0,
+        context_dim=None,
+        attn_mode="softmax",
+    ).cuda()
+    try:
+        x = torch.randn(BS, SEQLEN, DIM).cuda()
+        tic = time.time()
+        y = model(x)
+        toc = time.time()
+        print(y.shape, toc - tic)
+    except RuntimeError as e:
+        # likely oom
+        print(str(e))
+
+    print("\n now flash-cosine...")
+    model = BasicTransformerBlock(
+        dim=DIM,
+        n_heads=16,
+        d_head=64,
+        dropout=0.0,
+        context_dim=None,
+        attn_mode="flash-cosine",
+    ).cuda()
+    x = torch.randn(BS, SEQLEN, DIM).cuda()
+    tic = time.time()
+    y = model(x)
+    toc = time.time()
+    print(y.shape, toc - tic)
+    print("done with test02.\n")
+
+
+if __name__ == "__main__":
+    # test01()
+    # test02()
+    # test03()
+
+    # benchmark_attn()
+    benchmark_transformer_blocks()
+
+    print("done.")

MindEye_Imagery/src/generative_models/scripts/util/detection/nsfw_and_watermark_dectection.py

@@ -0,0 +1,110 @@
+import os
+
+import clip
+import numpy as np
+import torch
+import torchvision.transforms as T
+from PIL import Image
+
+RESOURCES_ROOT = "scripts/util/detection/"
+
+
+def predict_proba(X, weights, biases):
+    logits = X @ weights.T + biases
+    proba = np.where(
+        logits >= 0, 1 / (1 + np.exp(-logits)), np.exp(logits) / (1 + np.exp(logits))
+    )
+    return proba.T
+
+
+def load_model_weights(path: str):
+    model_weights = np.load(path)
+    return model_weights["weights"], model_weights["biases"]
+
+
+def clip_process_images(images: torch.Tensor) -> torch.Tensor:
+    min_size = min(images.shape[-2:])
+    return T.Compose(
+        [
+            T.CenterCrop(min_size),  # TODO: this might affect the watermark, check this
+            T.Resize(224, interpolation=T.InterpolationMode.BICUBIC, antialias=True),
+            T.Normalize(
+                (0.48145466, 0.4578275, 0.40821073),
+                (0.26862954, 0.26130258, 0.27577711),
+            ),
+        ]
+    )(images)
+
+
+class DeepFloydDataFiltering(object):
+    def __init__(
+        self, verbose: bool = False, device: torch.device = torch.device("cpu")
+    ):
+        super().__init__()
+        self.verbose = verbose
+        self._device = None
+        self.clip_model, _ = clip.load("ViT-L/14", device=device)
+        self.clip_model.eval()
+
+        self.cpu_w_weights, self.cpu_w_biases = load_model_weights(
+            os.path.join(RESOURCES_ROOT, "w_head_v1.npz")
+        )
+        self.cpu_p_weights, self.cpu_p_biases = load_model_weights(
+            os.path.join(RESOURCES_ROOT, "p_head_v1.npz")
+        )
+        self.w_threshold, self.p_threshold = 0.5, 0.5
+
+    @torch.inference_mode()
+    def __call__(self, images: torch.Tensor) -> torch.Tensor:
+        imgs = clip_process_images(images)
+        if self._device is None:
+            self._device = next(p for p in self.clip_model.parameters()).device
+        image_features = self.clip_model.encode_image(imgs.to(self._device))
+        image_features = image_features.detach().cpu().numpy().astype(np.float16)
+        p_pred = predict_proba(image_features, self.cpu_p_weights, self.cpu_p_biases)
+        w_pred = predict_proba(image_features, self.cpu_w_weights, self.cpu_w_biases)
+        print(f"p_pred = {p_pred}, w_pred = {w_pred}") if self.verbose else None
+        query = p_pred > self.p_threshold
+        if query.sum() > 0:
+            print(f"Hit for p_threshold: {p_pred}") if self.verbose else None
+            images[query] = T.GaussianBlur(99, sigma=(100.0, 100.0))(images[query])
+        query = w_pred > self.w_threshold
+        if query.sum() > 0:
+            print(f"Hit for w_threshold: {w_pred}") if self.verbose else None
+            images[query] = T.GaussianBlur(99, sigma=(100.0, 100.0))(images[query])
+        return images
+
+
+def load_img(path: str) -> torch.Tensor:
+    image = Image.open(path)
+    if not image.mode == "RGB":
+        image = image.convert("RGB")
+    image_transforms = T.Compose(
+        [
+            T.ToTensor(),
+        ]
+    )
+    return image_transforms(image)[None, ...]
+
+
+def test(root):
+    from einops import rearrange
+
+    filter = DeepFloydDataFiltering(verbose=True)
+    for p in os.listdir((root)):
+        print(f"running on {p}...")
+        img = load_img(os.path.join(root, p))
+        filtered_img = filter(img)
+        filtered_img = rearrange(
+            255.0 * (filtered_img.numpy())[0], "c h w -> h w c"
+        ).astype(np.uint8)
+        Image.fromarray(filtered_img).save(
+            os.path.join(root, f"{os.path.splitext(p)[0]}-filtered.jpg")
+        )
+
+
+if __name__ == "__main__":
+    import fire
+
+    fire.Fire(test)
+    print("done.")

MindEye_Imagery/src/vdvae/.ipynb_checkpoints/vae-checkpoint.py

@@ -0,0 +1,252 @@
+import torch
+from torch import nn
+from torch.nn import functional as F
+from vae_helpers import HModule, get_1x1, get_3x3, DmolNet, draw_gaussian_diag_samples, gaussian_analytical_kl
+from collections import defaultdict
+import numpy as np
+import itertools
+
+
+class Block(nn.Module):
+    def __init__(self, in_width, middle_width, out_width, down_rate=None, residual=False, use_3x3=True, zero_last=False):
+        super().__init__()
+        self.down_rate = down_rate
+        self.residual = residual
+        self.c1 = get_1x1(in_width, middle_width)
+        self.c2 = get_3x3(middle_width, middle_width) if use_3x3 else get_1x1(middle_width, middle_width)
+        self.c3 = get_3x3(middle_width, middle_width) if use_3x3 else get_1x1(middle_width, middle_width)
+        self.c4 = get_1x1(middle_width, out_width, zero_weights=zero_last)
+
+    def forward(self, x):
+        xhat = self.c1(F.gelu(x))
+        xhat = self.c2(F.gelu(xhat))
+        xhat = self.c3(F.gelu(xhat))
+        xhat = self.c4(F.gelu(xhat))
+        out = x + xhat if self.residual else xhat
+        if self.down_rate is not None:
+            out = F.avg_pool2d(out, kernel_size=self.down_rate, stride=self.down_rate)
+        return out
+
+
+def parse_layer_string(s):
+    layers = []
+    for ss in s.split(','):
+        if 'x' in ss:
+            res, num = ss.split('x')
+            count = int(num)
+            layers += [(int(res), None) for _ in range(count)]
+        elif 'm' in ss:
+            res, mixin = [int(a) for a in ss.split('m')]
+            layers.append((res, mixin))
+        elif 'd' in ss:
+            res, down_rate = [int(a) for a in ss.split('d')]
+            layers.append((res, down_rate))
+        else:
+            res = int(ss)
+            layers.append((res, None))
+    return layers
+
+
+def pad_channels(t, width):
+    d1, d2, d3, d4 = t.shape
+    empty = torch.zeros(d1, width, d3, d4, device=t.device)
+    empty[:, :d2, :, :] = t
+    return empty
+
+
+def get_width_settings(width, s):
+    mapping = defaultdict(lambda: width)
+    if s:
+        s = s.split(',')
+        for ss in s:
+            k, v = ss.split(':')
+            mapping[int(k)] = int(v)
+    return mapping
+
+
+class Encoder(HModule):
+    def build(self):
+        H = self.H
+        self.in_conv = get_3x3(H.image_channels, H.width)
+        self.widths = get_width_settings(H.width, H.custom_width_str)
+        enc_blocks = []
+        blockstr = parse_layer_string(H.enc_blocks)
+        for res, down_rate in blockstr:
+            use_3x3 = res > 2  # Don't use 3x3s for 1x1, 2x2 patches
+            enc_blocks.append(Block(self.widths[res], int(self.widths[res] * H.bottleneck_multiple), self.widths[res], down_rate=down_rate, residual=True, use_3x3=use_3x3))
+        n_blocks = len(blockstr)
+        for b in enc_blocks:
+            b.c4.weight.data *= np.sqrt(1 / n_blocks)
+        self.enc_blocks = nn.ModuleList(enc_blocks)
+
+    def forward(self, x):
+        x = x.permute(0, 3, 1, 2).contiguous()
+        x = self.in_conv(x)
+        activations = {}
+        activations[x.shape[2]] = x
+        for block in self.enc_blocks:
+            x = block(x)
+            res = x.shape[2]
+            x = x if x.shape[1] == self.widths[res] else pad_channels(x, self.widths[res])
+            activations[res] = x
+        return activations
+
+
+class DecBlock(nn.Module):
+    def __init__(self, H, res, mixin, n_blocks):
+        super().__init__()
+        self.base = res
+        self.mixin = mixin
+        self.H = H
+        self.widths = get_width_settings(H.width, H.custom_width_str)
+        width = self.widths[res]
+        use_3x3 = res > 2
+        cond_width = int(width * H.bottleneck_multiple)
+        self.zdim = H.zdim
+        self.enc = Block(width * 2, cond_width, H.zdim * 2, residual=False, use_3x3=use_3x3)
+        self.prior = Block(width, cond_width, H.zdim * 2 + width, residual=False, use_3x3=use_3x3, zero_last=True)
+        self.z_proj = get_1x1(H.zdim, width)
+        self.z_proj.weight.data *= np.sqrt(1 / n_blocks)
+        self.resnet = Block(width, cond_width, width, residual=True, use_3x3=use_3x3)
+        self.resnet.c4.weight.data *= np.sqrt(1 / n_blocks)
+        self.z_fn = lambda x: self.z_proj(x)
+
+    def sample(self, x, acts):
+        qm, qv = self.enc(torch.cat([x, acts], dim=1)).chunk(2, dim=1)
+        feats = self.prior(x)
+        pm, pv, xpp = feats[:, :self.zdim, ...], feats[:, self.zdim:self.zdim * 2, ...], feats[:, self.zdim * 2:, ...]
+        x = x + xpp
+        z = draw_gaussian_diag_samples(qm, qv)
+        kl = gaussian_analytical_kl(qm, pm, qv, pv)
+        return z, x, kl
+
+    def sample_uncond(self, x, t=None, lvs=None):
+        n, c, h, w = x.shape
+        feats = self.prior(x)
+        pm, pv, xpp = feats[:, :self.zdim, ...], feats[:, self.zdim:self.zdim * 2, ...], feats[:, self.zdim * 2:, ...]
+        x = x + xpp
+        if lvs is not None:
+            z = lvs
+        else:
+            if t is not None:
+                pv = pv + torch.ones_like(pv) * np.log(t)
+            z = draw_gaussian_diag_samples(pm, pv)
+        return z, x
+
+    def get_inputs(self, xs, activations):
+        acts = activations[self.base]
+        try:
+            x = xs[self.base]
+        except KeyError:
+            x = torch.zeros_like(acts)
+        if acts.shape[0] != x.shape[0]:
+            x = x.repeat(acts.shape[0], 1, 1, 1)
+        return x, acts
+
+    def forward(self, xs, activations, get_latents=False):
+        x, acts = self.get_inputs(xs, activations)
+        if self.mixin is not None:
+            x = x + F.interpolate(xs[self.mixin][:, :x.shape[1], ...], scale_factor=self.base // self.mixin)
+        z, x, kl = self.sample(x, acts)
+        x = x + self.z_fn(z)
+        x = self.resnet(x)
+        xs[self.base] = x
+        if get_latents:
+            return xs, dict(z=z.detach(), kl=kl)
+        return xs, dict(kl=kl)
+
+    def forward_uncond(self, xs, t=None, lvs=None):
+        try:
+            x = xs[self.base]
+        except KeyError:
+            ref = xs[list(xs.keys())[0]]
+            x = torch.zeros(dtype=ref.dtype, size=(ref.shape[0], self.widths[self.base], self.base, self.base), device=ref.device)
+        if self.mixin is not None:
+            x = x + F.interpolate(xs[self.mixin][:, :x.shape[1], ...], scale_factor=self.base // self.mixin)
+        z, x = self.sample_uncond(x, t, lvs=lvs)
+        x = x + self.z_fn(z)
+        x = self.resnet(x)
+        xs[self.base] = x
+        return xs
+
+
+class Decoder(HModule):
+
+    def build(self):
+        H = self.H
+        resos = set()
+        dec_blocks = []
+        self.widths = get_width_settings(H.width, H.custom_width_str)
+        blocks = parse_layer_string(H.dec_blocks)
+        for idx, (res, mixin) in enumerate(blocks):
+            dec_blocks.append(DecBlock(H, res, mixin, n_blocks=len(blocks)))
+            resos.add(res)
+        self.resolutions = sorted(resos)
+        self.dec_blocks = nn.ModuleList(dec_blocks)
+        self.bias_xs = nn.ParameterList([nn.Parameter(torch.zeros(1, self.widths[res], res, res)) for res in self.resolutions if res <= H.no_bias_above])
+        self.out_net = DmolNet(H)
+        self.gain = nn.Parameter(torch.ones(1, H.width, 1, 1))
+        self.bias = nn.Parameter(torch.zeros(1, H.width, 1, 1))
+        self.final_fn = lambda x: x * self.gain + self.bias
+
+    def forward(self, activations, get_latents=False):
+        stats = []
+        xs = {a.shape[2]: a for a in self.bias_xs}
+        for block in self.dec_blocks:
+            xs, block_stats = block(xs, activations, get_latents=get_latents)
+            stats.append(block_stats)
+        xs[self.H.image_size] = self.final_fn(xs[self.H.image_size])
+        return xs[self.H.image_size], stats
+
+    def forward_uncond(self, n, t=None, y=None):
+        xs = {}
+        for bias in self.bias_xs:
+            xs[bias.shape[2]] = bias.repeat(n, 1, 1, 1)
+        for idx, block in enumerate(self.dec_blocks):
+            try:
+                temp = t[idx]
+            except TypeError:
+                temp = t
+            xs = block.forward_uncond(xs, temp)
+        xs[self.H.image_size] = self.final_fn(xs[self.H.image_size])
+        return xs[self.H.image_size]
+
+    def forward_manual_latents(self, n, latents, t=None):
+        xs = {}
+        for bias in self.bias_xs:
+            xs[bias.shape[2]] = bias.repeat(n, 1, 1, 1)
+        for block, lvs in itertools.zip_longest(self.dec_blocks, latents):
+            xs = block.forward_uncond(xs, t, lvs=lvs)
+        xs[self.H.image_size] = self.final_fn(xs[self.H.image_size])
+        return xs[self.H.image_size]
+
+
+class VAE(HModule):
+    def build(self):
+        self.encoder = Encoder(self.H)
+        self.decoder = Decoder(self.H)
+
+    def forward(self, x, x_target):
+        activations = self.encoder.forward(x)
+        px_z, stats = self.decoder.forward(activations)
+        distortion_per_pixel = self.decoder.out_net.nll(px_z, x_target)
+        rate_per_pixel = torch.zeros_like(distortion_per_pixel)
+        ndims = np.prod(x.shape[1:])
+        for statdict in stats:
+            rate_per_pixel += statdict['kl'].sum(dim=(1, 2, 3))
+        rate_per_pixel /= ndims
+        elbo = (distortion_per_pixel + rate_per_pixel).mean()
+        return dict(elbo=elbo, distortion=distortion_per_pixel.mean(), rate=rate_per_pixel.mean())
+
+    def forward_get_latents(self, x):
+        activations = self.encoder.forward(x)
+        _, stats = self.decoder.forward(activations, get_latents=True)
+        return stats
+
+    def forward_uncond_samples(self, n_batch, t=None):
+        px_z = self.decoder.forward_uncond(n_batch, t=t)
+        return self.decoder.out_net.sample(px_z)
+
+    def forward_samples_set_latents(self, n_batch, latents, t=None):
+        px_z = self.decoder.forward_manual_latents(n_batch, latents, t=t)
+        return self.decoder.out_net.sample(px_z)

MindEye_Imagery/src/vdvae/LICENSE.md

@@ -0,0 +1,7 @@
+Copyright © 2020 OpenAI
+
+Permission is hereby granted, free of charge, to any person obtaining a copy of this software and associated documentation files (the “Software”), to deal in the Software without restriction, including without limitation the rights to use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of the Software, and to permit persons to whom the Software is furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED “AS IS”, WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.

MindEye_Imagery/src/vdvae/README.md

@@ -0,0 +1,103 @@
+# Very Deep VAEs
+
+Repository for the paper "Very Deep VAEs Generalize Autoregressive Models and Can Outperform Them on Images" (https://arxiv.org/abs/2011.10650)
+
+Some model samples and a visualization of how it generates them:
+![image](header-image.png)
+
+This repository is tested with PyTorch 1.6, CUDA 10.1, Numpy 1.16, Ubuntu 18.04, and V100 GPUs.
+
+# Setup
+Several additional packages are required, including NVIDIA Apex:
+```
+pip install imageio
+pip install mpi4py
+pip install sklearn
+git clone https://github.com/NVIDIA/apex
+cd apex
+pip install -v --no-cache-dir --global-option="--cpp_ext" --global-option="--cuda_ext" ./
+cd ..
+```
+
+Also, you'll have to download the data, depending on which one you want to run:
+```
+./setup_cifar10.sh
+./setup_imagenet.sh imagenet32
+./setup_imagenet.sh imagenet64
+./setup_ffhq256.sh
+./setup_ffhq1024.sh  /path/to/images1024x1024  # this one depends on you first downloading the subfolder `images_1024x1024` from https://github.com/NVlabs/ffhq-dataset on your own
+```
+
+# Training models
+Hyperparameters all reside in `hps.py`. We use 2 gpus for our CIFAR-10 runs, and 32 for the rest of the models. (Using a lower batch size is also possible and results in slower learning, and may also require a lower learning rate).
+
+The `mpiexec` arguments you use for runs with more than 1 node depend on the configuration of your system, so please adapt accordingly.
+
+```bash
+mpiexec -n 2 python train.py --hps cifar10
+mpiexec -n 32 python train.py --hps imagenet32
+mpiexec -n 32 python train.py --hps imagenet64
+mpiexec -n 32 python train.py --hps ffhq256
+mpiexec -n 32 python train.py --hps ffhq1024
+```
+
+# Restoring saved models
+For convenience, we have included training checkpoints which can be restored in order to confirm performance, continue training, or generate samples.
+
+### ImageNet 32
+```bash
+# 119M parameter model, trained for 1.7M iters (about 2.5 weeks on 32 V100)
+wget https://openaipublic.blob.core.windows.net/very-deep-vaes-assets/vdvae-assets/imagenet32-iter-1700000-log.jsonl
+wget https://openaipublic.blob.core.windows.net/very-deep-vaes-assets/vdvae-assets/imagenet32-iter-1700000-model.th
+wget https://openaipublic.blob.core.windows.net/very-deep-vaes-assets/vdvae-assets/imagenet32-iter-1700000-model-ema.th
+wget https://openaipublic.blob.core.windows.net/very-deep-vaes-assets/vdvae-assets/imagenet32-iter-1700000-opt.th
+python train.py --hps imagenet32 --restore_path imagenet32-iter-1700000-model.th --restore_ema_path imagenet32-iter-1700000-model-ema.th --restore_log_path imagenet32-iter-1700000-log.jsonl --restore_optimizer_path imagenet32-iter-1700000-opt.th --test_eval
+# should give 2.6364 nats per dim, which is 3.80 bpd
+```
+
+### ImageNet 64
+```bash
+# 125M parameter model, trained for 1.6M iters (about 2.5 weeks on 32 V100)
+wget https://openaipublic.blob.core.windows.net/very-deep-vaes-assets/vdvae-assets-2/imagenet64-iter-1600000-log.jsonl
+wget https://openaipublic.blob.core.windows.net/very-deep-vaes-assets/vdvae-assets-2/imagenet64-iter-1600000-model.th
+wget https://openaipublic.blob.core.windows.net/very-deep-vaes-assets/vdvae-assets-2/imagenet64-iter-1600000-model-ema.th
+wget https://openaipublic.blob.core.windows.net/very-deep-vaes-assets/vdvae-assets-2/imagenet64-iter-1600000-opt.th
+python train.py --hps imagenet64 --restore_path imagenet64-iter-1600000-model.th --restore_ema_path imagenet64-iter-1600000-model-ema.th --restore_log_path imagenet64-iter-1600000-log.jsonl --restore_optimizer_path imagenet64-iter-1600000-opt.th --test_eval
+# should be 2.44 nats, or 3.52 bits per dim
+```
+
+### FFHQ-256
+```bash
+# 115M parameters, trained for 1.7M iterations (or about 2.5 weeks) on 32 V100
+wget https://openaipublic.blob.core.windows.net/very-deep-vaes-assets/vdvae-assets/ffhq256-iter-1700000-log.jsonl
+wget https://openaipublic.blob.core.windows.net/very-deep-vaes-assets/vdvae-assets/ffhq256-iter-1700000-model.th
+wget https://openaipublic.blob.core.windows.net/very-deep-vaes-assets/vdvae-assets/ffhq256-iter-1700000-model-ema.th
+wget https://openaipublic.blob.core.windows.net/very-deep-vaes-assets/vdvae-assets/ffhq256-iter-1700000-opt.th
+python train.py --hps ffhq256 --restore_path ffhq256-iter-1700000-model.th --restore_ema_path ffhq256-iter-1700000-model-ema.th --restore_log_path ffhq256-iter-1700000-log.jsonl --restore_optimizer_path ffhq256-iter-1700000-opt.th --test_eval
+# should be 0.4232 nats, or 0.61 bits per dim
+```
+
+### FFHQ-1024
+```bash
+# 115M parameters, trained for 1.7M iterations (or about 2.5 weeks) on 32 V100
+wget https://openaipublic.blob.core.windows.net/very-deep-vaes-assets/vdvae-assets/ffhq1024-iter-1700000-log.jsonl
+wget https://openaipublic.blob.core.windows.net/very-deep-vaes-assets/vdvae-assets/ffhq1024-iter-1700000-model.th
+wget https://openaipublic.blob.core.windows.net/very-deep-vaes-assets/vdvae-assets/ffhq1024-iter-1700000-model-ema.th
+wget https://openaipublic.blob.core.windows.net/very-deep-vaes-assets/vdvae-assets/ffhq1024-iter-1700000-opt.th
+python train.py --hps ffhq1024 --restore_path ffhq1024-iter-1700000-model.th --restore_ema_path ffhq1024-iter-1700000-model-ema.th --restore_log_path ffhq1024-iter-1700000-log.jsonl --restore_optimizer_path ffhq1024-iter-1700000-opt.th --test_eval
+# should be 1.678 nats, or 2.42 bits per dim
+```
+
+### CIFAR-10
+```bash
+# 39M parameters, trained for ~1M iterations with early stopping (a little less than a week on 2 GPUs)
+wget https://openaipublic.blob.core.windows.net/very-deep-vaes-assets/vdvae-assets-2/cifar10-seed0-iter-900000-model-ema.th
+wget https://openaipublic.blob.core.windows.net/very-deep-vaes-assets/vdvae-assets-2/cifar10-seed1-iter-1050000-model-ema.th
+wget https://openaipublic.blob.core.windows.net/very-deep-vaes-assets/vdvae-assets-2/cifar10-seed2-iter-650000-model-ema.th
+wget https://openaipublic.blob.core.windows.net/very-deep-vaes-assets/vdvae-assets-2/cifar10-seed3-iter-1050000-model-ema.th
+python train.py --hps cifar10 --restore_ema_path cifar10-seed0-iter-900000-model-ema.th --test_eval
+python train.py --hps cifar10 --restore_ema_path cifar10-seed1-iter-1050000-model-ema.th --test_eval
+python train.py --hps cifar10 --restore_ema_path cifar10-seed2-iter-650000-model-ema.th --test_eval
+python train.py --hps cifar10 --restore_ema_path cifar10-seed3-iter-1050000-model-ema.th --test_eval
+# seeds 0, 1, 2, 3 should give 2.879, 2.842, 2.898, 2.864 bits per dim, for an average of 2.87 bits per dim.
+```

MindEye_Imagery/src/vdvae/data.py

@@ -0,0 +1,163 @@
+import numpy as np
+import pickle
+import os
+import torch
+from torch.utils.data import TensorDataset
+from torchvision.datasets import ImageFolder
+import torchvision.transforms as transforms
+from sklearn.model_selection import train_test_split
+
+
+def set_up_data(H):
+    shift_loss = -127.5
+    scale_loss = 1. / 127.5
+    if H.dataset == 'imagenet32':
+        trX, vaX, teX = imagenet32(H.data_root)
+        H.image_size = 32
+        H.image_channels = 3
+        shift = -116.2373
+        scale = 1. / 69.37404
+    elif H.dataset == 'imagenet64':
+        trX, vaX, teX = imagenet64(H.data_root)
+        H.image_size = 64
+        H.image_channels = 3
+        shift = -115.92961967
+        scale = 1. / 69.37404
+    elif H.dataset == 'ffhq_256':
+        trX, vaX, teX = ffhq256(H.data_root)
+        H.image_size = 256
+        H.image_channels = 3
+        shift = -112.8666757481
+        scale = 1. / 69.84780273
+    elif H.dataset == 'ffhq_1024':
+        trX, vaX, teX = ffhq1024(H.data_root)
+        H.image_size = 1024
+        H.image_channels = 3
+        shift = -0.4387
+        scale = 1.0 / 0.2743
+        shift_loss = -0.5
+        scale_loss = 2.0
+    elif H.dataset == 'cifar10':
+        (trX, _), (vaX, _), (teX, _) = cifar10(H.data_root, one_hot=False)
+        H.image_size = 32
+        H.image_channels = 3
+        shift = -120.63838
+        scale = 1. / 64.16736
+    else:
+        raise ValueError('unknown dataset: ', H.dataset)
+
+    do_low_bit = H.dataset in ['ffhq_256']
+
+    if H.test_eval:
+        print('DOING TEST')
+        eval_dataset = teX
+    else:
+        eval_dataset = vaX
+
+    shift = torch.tensor([shift]).cuda().view(1, 1, 1, 1)
+    scale = torch.tensor([scale]).cuda().view(1, 1, 1, 1)
+    shift_loss = torch.tensor([shift_loss]).cuda().view(1, 1, 1, 1)
+    scale_loss = torch.tensor([scale_loss]).cuda().view(1, 1, 1, 1)
+
+    if H.dataset == 'ffhq_1024':
+        train_data = ImageFolder(trX, transforms.ToTensor())
+        valid_data = ImageFolder(eval_dataset, transforms.ToTensor())
+        untranspose = True
+    else:
+        train_data = TensorDataset(torch.as_tensor(trX))
+        valid_data = TensorDataset(torch.as_tensor(eval_dataset))
+        untranspose = False
+
+    def preprocess_func(x):
+        nonlocal shift
+        nonlocal scale
+        nonlocal shift_loss
+        nonlocal scale_loss
+        nonlocal do_low_bit
+        nonlocal untranspose
+        'takes in a data example and returns the preprocessed input'
+        'as well as the input processed for the loss'
+        if untranspose:
+            x[0] = x[0].permute(0, 2, 3, 1)
+        inp = x[0].cuda(non_blocking=True).float()
+        out = inp.clone()
+        inp.add_(shift).mul_(scale)
+        if do_low_bit:
+            # 5 bits of precision
+            out.mul_(1. / 8.).floor_().mul_(8.)
+        out.add_(shift_loss).mul_(scale_loss)
+        return inp, out
+
+    return H, train_data, valid_data, preprocess_func
+
+
+def mkdir_p(path):
+    os.makedirs(path, exist_ok=True)
+
+
+def flatten(outer):
+    return [el for inner in outer for el in inner]
+
+
+def unpickle_cifar10(file):
+    fo = open(file, 'rb')
+    data = pickle.load(fo, encoding='bytes')
+    fo.close()
+    data = dict(zip([k.decode() for k in data.keys()], data.values()))
+    return data
+
+
+def imagenet32(data_root):
+    trX = np.load(os.path.join(data_root, 'imagenet32-train.npy'), mmap_mode='r')
+    np.random.seed(42)
+    tr_va_split_indices = np.random.permutation(trX.shape[0])
+    train = trX[tr_va_split_indices[:-5000]]
+    valid = trX[tr_va_split_indices[-5000:]]
+    test = np.load(os.path.join(data_root, 'imagenet32-valid.npy'), mmap_mode='r')
+    return train, valid, test
+
+
+def imagenet64(data_root):
+    trX = np.load(os.path.join(data_root, 'imagenet64-train.npy'), mmap_mode='r')
+    np.random.seed(42)
+    tr_va_split_indices = np.random.permutation(trX.shape[0])
+    train = trX[tr_va_split_indices[:-5000]]
+    valid = trX[tr_va_split_indices[-5000:]]
+    test = np.load(os.path.join(data_root, 'imagenet64-valid.npy'), mmap_mode='r')  # this is test.
+    return train, valid, test
+
+
+def ffhq1024(data_root):
+    # we did not significantly tune hyperparameters on ffhq-1024, and so simply evaluate on the test set
+    return os.path.join(data_root, 'ffhq1024/train'), os.path.join(data_root, 'ffhq1024/valid'), os.path.join(data_root, 'ffhq1024/valid')
+
+
+def ffhq256(data_root):
+    trX = np.load(os.path.join(data_root, 'ffhq-256.npy'), mmap_mode='r')
+    np.random.seed(5)
+    tr_va_split_indices = np.random.permutation(trX.shape[0])
+    train = trX[tr_va_split_indices[:-7000]]
+    valid = trX[tr_va_split_indices[-7000:]]
+    # we did not significantly tune hyperparameters on ffhq-256, and so simply evaluate on the test set
+    return train, valid, valid
+
+
+def cifar10(data_root, one_hot=True):
+    tr_data = [unpickle_cifar10(os.path.join(data_root, 'cifar-10-batches-py/', 'data_batch_%d' % i)) for i in range(1, 6)]
+    trX = np.vstack(data['data'] for data in tr_data)
+    trY = np.asarray(flatten([data['labels'] for data in tr_data]))
+    te_data = unpickle_cifar10(os.path.join(data_root, 'cifar-10-batches-py/', 'test_batch'))
+    teX = np.asarray(te_data['data'])
+    teY = np.asarray(te_data['labels'])
+    trX = trX.reshape(-1, 3, 32, 32).transpose(0, 2, 3, 1)
+    teX = teX.reshape(-1, 3, 32, 32).transpose(0, 2, 3, 1)
+    trX, vaX, trY, vaY = train_test_split(trX, trY, test_size=5000, random_state=11172018)
+    if one_hot:
+        trY = np.eye(10, dtype=np.float32)[trY]
+        vaY = np.eye(10, dtype=np.float32)[vaY]
+        teY = np.eye(10, dtype=np.float32)[teY]
+    else:
+        trY = np.reshape(trY, [-1, 1])
+        vaY = np.reshape(vaY, [-1, 1])
+        teY = np.reshape(teY, [-1, 1])
+    return (trX, trY), (vaX, vaY), (teX, teY)

MindEye_Imagery/src/vdvae/files_to_npy.py

@@ -0,0 +1,14 @@
+import sys
+import numpy as np
+import imageio
+import glob
+import os
+
+if __name__ == "__main__":
+    print("moving images in", sys.argv[1], "to", sys.argv[2])
+    files = glob.glob(os.path.join(sys.argv[1], "*.png"))
+    shape = imageio.imread(files[0]).shape
+    data = np.zeros(shape=(len(files), *shape), dtype=np.uint8)
+    for idx, f in enumerate(files):
+        data[idx] = imageio.imread(f)
+    np.save(sys.argv[2], data)

MindEye_Imagery/src/vdvae/hps.py

@@ -0,0 +1,159 @@
+HPARAMS_REGISTRY = {}
+
+
+class Hyperparams(dict):
+    def __getattr__(self, attr):
+        try:
+            return self[attr]
+        except KeyError:
+            return None
+
+    def __setattr__(self, attr, value):
+        self[attr] = value
+
+
+cifar10 = Hyperparams()
+cifar10.width = 384
+cifar10.lr = 0.0002
+cifar10.zdim = 16
+cifar10.wd = 0.01
+cifar10.dec_blocks = "1x1,4m1,4x2,8m4,8x5,16m8,16x10,32m16,32x21"
+cifar10.enc_blocks = "32x11,32d2,16x6,16d2,8x6,8d2,4x3,4d4,1x3"
+cifar10.warmup_iters = 100
+cifar10.dataset = 'cifar10'
+cifar10.n_batch = 16
+cifar10.ema_rate = 0.9999
+HPARAMS_REGISTRY['cifar10'] = cifar10
+
+
+i32 = Hyperparams()
+i32.update(cifar10)
+i32.dataset = 'imagenet32'
+i32.ema_rate = 0.999
+i32.dec_blocks = "1x2,4m1,4x4,8m4,8x9,16m8,16x19,32m16,32x40"
+i32.enc_blocks = "32x15,32d2,16x9,16d2,8x8,8d2,4x6,4d4,1x6"
+i32.width = 512
+i32.n_batch = 8
+i32.lr = 0.00015
+i32.grad_clip = 200.
+i32.skip_threshold = 300.
+i32.epochs_per_eval = 1
+i32.epochs_per_eval_save = 1
+HPARAMS_REGISTRY['imagenet32'] = i32
+
+i64 = Hyperparams()
+i64.update(i32)
+i64.n_batch = 4
+i64.grad_clip = 220.0
+i64.skip_threshold = 380.0
+i64.dataset = 'imagenet64'
+i64.dec_blocks = "1x2,4m1,4x3,8m4,8x7,16m8,16x15,32m16,32x31,64m32,64x12"
+i64.enc_blocks = "64x11,64d2,32x20,32d2,16x9,16d2,8x8,8d2,4x7,4d4,1x5"
+HPARAMS_REGISTRY['imagenet64'] = i64
+
+ffhq_256 = Hyperparams()
+ffhq_256.update(i64)
+ffhq_256.n_batch = 1
+ffhq_256.lr = 0.00015
+ffhq_256.dataset = 'ffhq_256'
+ffhq_256.epochs_per_eval = 1
+ffhq_256.epochs_per_eval_save = 1
+ffhq_256.num_images_visualize = 2
+ffhq_256.num_variables_visualize = 3
+ffhq_256.num_temperatures_visualize = 1
+ffhq_256.dec_blocks = "1x2,4m1,4x3,8m4,8x4,16m8,16x9,32m16,32x21,64m32,64x13,128m64,128x7,256m128"
+ffhq_256.enc_blocks = "256x3,256d2,128x8,128d2,64x12,64d2,32x17,32d2,16x7,16d2,8x5,8d2,4x5,4d4,1x4"
+ffhq_256.no_bias_above = 64
+ffhq_256.grad_clip = 130.
+ffhq_256.skip_threshold = 180.
+HPARAMS_REGISTRY['ffhq256'] = ffhq_256
+
+
+ffhq1024 = Hyperparams()
+ffhq1024.update(ffhq_256)
+ffhq1024.dataset = 'ffhq_1024'
+ffhq1024.data_root = './ffhq_images1024x1024'
+ffhq1024.epochs_per_eval = 1
+ffhq1024.epochs_per_eval_save = 1
+ffhq1024.num_images_visualize = 1
+ffhq1024.iters_per_images = 25000
+ffhq1024.num_variables_visualize = 0
+ffhq1024.num_temperatures_visualize = 4
+ffhq1024.grad_clip = 360.
+ffhq1024.skip_threshold = 500.
+ffhq1024.num_mixtures = 2
+ffhq1024.width = 16
+ffhq1024.lr = 0.00007
+ffhq1024.dec_blocks = "1x2,4m1,4x3,8m4,8x4,16m8,16x9,32m16,32x20,64m32,64x14,128m64,128x7,256m128,256x2,512m256,1024m512"
+ffhq1024.enc_blocks = "1024x1,1024d2,512x3,512d2,256x5,256d2,128x7,128d2,64x10,64d2,32x14,32d2,16x7,16d2,8x5,8d2,4x5,4d4,1x4"
+ffhq1024.custom_width_str = "512:32,256:64,128:512,64:512,32:512,16:512,8:512,4:512,1:512"
+HPARAMS_REGISTRY['ffhq1024'] = ffhq1024
+
+
+def parse_args_and_update_hparams(H, parser, s=None):
+    args = parser.parse_args(s)
+    valid_args = set(args.__dict__.keys())
+    hparam_sets = [x for x in args.hparam_sets.split(',') if x]
+    for hp_set in hparam_sets:
+        hps = HPARAMS_REGISTRY[hp_set]
+        for k in hps:
+            if k not in valid_args:
+                raise ValueError(f"{k} not in default args")
+        parser.set_defaults(**hps)
+    H.update(parser.parse_args(s).__dict__)
+
+
+def add_vae_arguments(parser):
+    parser.add_argument('--seed', type=int, default=0)
+    parser.add_argument('--port', type=int, default=29500)
+    parser.add_argument('--save_dir', type=str, default='./saved_models')
+    parser.add_argument('--data_root', type=str, default='./')
+
+    parser.add_argument('--desc', type=str, default='test')
+    parser.add_argument('--hparam_sets', '--hps', type=str)
+    parser.add_argument('--restore_path', type=str, default=None)
+    parser.add_argument('--restore_ema_path', type=str, default=None)
+    parser.add_argument('--restore_log_path', type=str, default=None)
+    parser.add_argument('--restore_optimizer_path', type=str, default=None)
+    parser.add_argument('--dataset', type=str, default='cifar10')
+
+    parser.add_argument('--ema_rate', type=float, default=0.999)
+
+    parser.add_argument('--enc_blocks', type=str, default=None)
+    parser.add_argument('--dec_blocks', type=str, default=None)
+    parser.add_argument('--zdim', type=int, default=16)
+    parser.add_argument('--width', type=int, default=512)
+    parser.add_argument('--custom_width_str', type=str, default='')
+    parser.add_argument('--bottleneck_multiple', type=float, default=0.25)
+
+    parser.add_argument('--no_bias_above', type=int, default=64)
+    parser.add_argument('--scale_encblock', action="store_true")
+
+    parser.add_argument('--test_eval', action="store_true")
+    parser.add_argument('--warmup_iters', type=float, default=0)
+
+    parser.add_argument('--num_mixtures', type=int, default=10)
+    parser.add_argument('--grad_clip', type=float, default=200.0)
+    parser.add_argument('--skip_threshold', type=float, default=400.0)
+    parser.add_argument('--lr', type=float, default=0.00015)
+    parser.add_argument('--lr_prior', type=float, default=0.00015)
+    parser.add_argument('--wd', type=float, default=0.0)
+    parser.add_argument('--wd_prior', type=float, default=0.0)
+    parser.add_argument('--num_epochs', type=int, default=10000)
+    parser.add_argument('--n_batch', type=int, default=32)
+    parser.add_argument('--adam_beta1', type=float, default=0.9)
+    parser.add_argument('--adam_beta2', type=float, default=0.9)
+
+    parser.add_argument('--temperature', type=float, default=1.0)
+
+    parser.add_argument('--iters_per_ckpt', type=int, default=25000)
+    parser.add_argument('--iters_per_print', type=int, default=1000)
+    parser.add_argument('--iters_per_save', type=int, default=10000)
+    parser.add_argument('--iters_per_images', type=int, default=10000)
+    parser.add_argument('--epochs_per_eval', type=int, default=10)
+    parser.add_argument('--epochs_per_probe', type=int, default=None)
+    parser.add_argument('--epochs_per_eval_save', type=int, default=20)
+    parser.add_argument('--num_images_visualize', type=int, default=8)
+    parser.add_argument('--num_variables_visualize', type=int, default=6)
+    parser.add_argument('--num_temperatures_visualize', type=int, default=3)
+    return parser

MindEye_Imagery/src/vdvae/image_utils.py

@@ -0,0 +1,68 @@
+import io
+import IPython.display
+import PIL.Image
+import os
+from pprint import pformat
+import numpy as np
+
+def imgrid(imarray, cols=4, pad=1, padval=255, row_major=True):
+  """Lays out a [N, H, W, C] image array as a single image grid."""
+  pad = int(pad)
+  if pad < 0:
+    raise ValueError('pad must be non-negative')
+  cols = int(cols)
+  assert cols >= 1
+  N, H, W, C = imarray.shape
+  rows = N // cols + int(N % cols != 0)
+  batch_pad = rows * cols - N
+  assert batch_pad >= 0
+  post_pad = [batch_pad, pad, pad, 0]
+  pad_arg = [[0, p] for p in post_pad]
+  imarray = np.pad(imarray, pad_arg, 'constant', constant_values=padval)
+  H += pad
+  W += pad
+  grid = (imarray
+          .reshape(rows, cols, H, W, C)
+          .transpose(0, 2, 1, 3, 4)
+          .reshape(rows*H, cols*W, C))
+  if pad:
+    grid = grid[:-pad, :-pad]
+  return grid
+
+def interleave(*args):
+  """Interleaves input arrays of the same shape along the batch axis."""
+  if not args:
+    raise ValueError('At least one argument is required.')
+  a0 = args[0]
+  if any(a.shape != a0.shape for a in args):
+    raise ValueError('All inputs must have the same shape.')
+  if not a0.shape:
+    raise ValueError('Inputs must have at least one axis.')
+  out = np.transpose(args, [1, 0] + list(range(2, len(a0.shape) + 1)))
+  out = out.reshape(-1, *a0.shape[1:])
+  return out
+
+def imshow(a, format='png', jpeg_fallback=True):
+  """Displays an image in the given format."""
+  a = a.astype(np.uint8)
+  data = io.BytesIO()
+  PIL.Image.fromarray(a).save(data, format)
+  im_data = data.getvalue()
+  try:
+    disp = IPython.display.display(IPython.display.Image(im_data))
+  except IOError:
+    if jpeg_fallback and format != 'jpeg':
+      print ('Warning: image was too large to display in format "{}"; '
+             'trying jpeg instead.').format(format)
+      return imshow(a, format='jpeg')
+    else:
+      raise
+  return disp
+
+def image_to_uint8(x):
+  """Converts [-1, 1] float array to [0, 255] uint8."""
+  x = np.asarray(x)
+  x = (256. / 2.) * (x + 1.)
+  x = np.clip(x, 0, 255)
+  x = x.astype(np.uint8)
+  return x

MindEye_Imagery/src/vdvae/model_utils.py

@@ -0,0 +1,102 @@
+import torch
+import numpy as np
+#from mpi4py import MPI
+import socket
+import argparse
+import os
+import json
+import subprocess
+from hps import Hyperparams, parse_args_and_update_hparams, add_vae_arguments
+from vdvae_utils import (logger,
+                   local_mpi_rank,
+                   mpi_size,
+                   maybe_download,
+                   mpi_rank)
+from data import mkdir_p
+from contextlib import contextmanager
+import torch.distributed as dist
+#from apex.optimizers import FusedAdam as AdamW
+from vae import VAE
+from torch.nn.parallel.distributed import DistributedDataParallel
+from train_helpers import restore_params
+
+def set_up_hyperparams(s=None):
+    H = Hyperparams()
+    parser = argparse.ArgumentParser()
+    parser = add_vae_arguments(parser)
+    parse_args_and_update_hparams(H, parser, s=s)
+    setup_mpi(H)
+    setup_save_dirs(H)
+    logprint = logger(H.logdir)
+    for i, k in enumerate(sorted(H)):
+        logprint(type='hparam', key=k, value=H[k])
+    np.random.seed(H.seed)
+    torch.manual_seed(H.seed)
+    torch.cuda.manual_seed(H.seed)
+    logprint('training model', H.desc, 'on', H.dataset)
+    return H, logprint
+
+def set_up_data(H, device="cuda"):
+    shift_loss = -127.5
+    scale_loss = 1. / 127.5
+    
+    #trX, vaX, teX = imagenet64(H.data_root)
+    H.image_size = 64
+    H.image_channels = 3
+    shift = -115.92961967
+    scale = 1. / 69.37404
+    
+    #if H.test_eval:
+    #    print('DOING TEST')
+    #    eval_dataset = teX
+    #else:
+    #    eval_dataset = vaX
+
+    shift = torch.tensor([shift]).to(device).view(1, 1, 1, 1)
+    scale = torch.tensor([scale]).to(device).view(1, 1, 1, 1)
+    shift_loss = torch.tensor([shift_loss]).to(device).view(1, 1, 1, 1)
+    scale_loss = torch.tensor([scale_loss]).to(device).view(1, 1, 1, 1)
+    
+    #train_data = TensorDataset(torch.as_tensor(trX))
+    #valid_data = TensorDataset(torch.as_tensor(eval_dataset))
+    #untranspose = False
+
+    def preprocess_func(x):
+        nonlocal shift
+        nonlocal scale
+        nonlocal shift_loss
+        nonlocal scale_loss
+        'takes in a data example and returns the preprocessed input'
+        'as well as the input processed for the loss'
+        #untranspose = False
+        #if untranspose:
+        #    x[0] = x[0].permute(0, 2, 3, 1)
+        inp = x.to(device, non_blocking=True).float()
+        out = inp.clone()
+        inp.add_(shift).mul_(scale)
+        out.add_(shift_loss).mul_(scale_loss)
+        return inp, out
+    
+    return H, preprocess_func
+
+def load_vaes(H, logprint=None, device="cuda"):
+
+    ema_vae = VAE(H)
+    if H.restore_ema_path:
+        print(f'Restoring ema vae from {H.restore_ema_path}')
+        restore_params(ema_vae, H.restore_ema_path, map_cpu=True, local_rank=H.local_rank, mpi_size=H.mpi_size)
+    else:
+        ema_vae.load_state_dict(vae.state_dict())
+    ema_vae.requires_grad_(False)
+    ema_vae = ema_vae.to(device)
+
+    #vae = DistributedDataParallel(vae, device_ids=[H.local_rank], output_device=H.local_rank)
+
+    #if len(list(vae.named_parameters())) != len(list(vae.parameters())):
+    #    raise ValueError('Some params are not named. Please name all params.')
+    #total_params = 0
+    #for name, p in vae.named_parameters():
+    #    total_params += np.prod(p.shape)
+
+    #print(total_params=total_params, readable=f'{total_params:,}')
+    return ema_vae

MindEye_Imagery/src/vdvae/setup_cifar10.sh

@@ -0,0 +1,2 @@
+wget https://www.cs.toronto.edu/~kriz/cifar-10-python.tar.gz
+tar -xf cifar-10-python.tar.gz

MindEye_Imagery/src/vdvae/setup_ffhq1024.sh

@@ -0,0 +1,13 @@
+# the first argument to this script should be the path to the ffhq_images1024x1024 folder
+# the same path should be provided as the `data_root` argument to train.py
+cd $1
+mkdir train
+mkdir train/0
+mkdir valid
+mkdir valid/0
+for i in $(seq -f "%05g" 0 64999); do
+    mv $i.png train/0
+done
+for i in $(seq -f "%05g" 65000 69999); do
+    mv $i.png valid/0
+done

MindEye_Imagery/src/vdvae/setup_ffhq256.sh

@@ -0,0 +1,12 @@
+# we provide a copy of ffhq-256 for convenience, downsampled using the same function as NVAE (https://github.com/NVlabs/NVAE) (personal communication with author)
+
+# Resizing function is this one, with the default second argument and size=256
+# https://pytorch.org/docs/stable/torchvision/transforms.html#torchvision.transforms.Resize
+
+# 5-bit precision is calculated using the following lines, with num_bits=5 and for an x in [0, 1]
+# x = torch.floor(x * 255 / 2 ** (8 - num_bits))
+# x /= (2 ** num_bits - 1)
+
+# the DMOL loss should also be adjusted to have 32 buckets instead of 256 (this code, or NVAE, can be used as reference)
+
+wget https://openaipublic.blob.core.windows.net/very-deep-vaes-assets/vdvae-assets/ffhq-256.npy

MindEye_Imagery/src/vdvae/setup_imagenet.sh

@@ -0,0 +1,28 @@
+if [ "$1" == "imagenet32" ]; then
+
+echo "downloading imagenet32"
+wget http://www.image-net.org/small/train_32x32.tar
+wget http://www.image-net.org/small/valid_32x32.tar
+tar -xvf train_32x32.tar
+tar -xvf valid_32x32.tar
+python files_to_npy.py train_32x32/ imagenet32-train.npy
+python files_to_npy.py valid_32x32/ imagenet32-valid.npy
+
+elif [ "$1" == "imagenet64" ]; then
+
+echo "downloading imagenet64"
+wget http://www.image-net.org/small/train_64x64.tar
+wget http://www.image-net.org/small/valid_64x64.tar
+tar -xvf train_64x64.tar
+tar -xvf valid_64x64.tar
+python files_to_npy.py train_64x64/ imagenet64-train.npy
+python files_to_npy.py valid_64x64/ imagenet64-valid.npy
+
+else
+
+echo "please pass the string imagenet32 or imagenet64 as an argument"
+
+fi
+
+
+

MindEye_Imagery/src/vdvae/train.py

@@ -0,0 +1,136 @@
+import numpy as np
+import imageio
+import os
+import time
+import torch
+from torch.utils.data import DataLoader
+from torch.utils.data.distributed import DistributedSampler
+from data import set_up_data
+from vdvae_utils import get_cpu_stats_over_ranks
+from train_helpers import set_up_hyperparams, load_vaes, load_opt, accumulate_stats, save_model, update_ema
+
+
+def training_step(H, data_input, target, vae, ema_vae, optimizer, iterate):
+    t0 = time.time()
+    vae.zero_grad()
+    stats = vae.forward(data_input, target)
+    stats['elbo'].backward()
+    grad_norm = torch.nn.utils.clip_grad_norm_(vae.parameters(), H.grad_clip).item()
+    distortion_nans = torch.isnan(stats['distortion']).sum()
+    rate_nans = torch.isnan(stats['rate']).sum()
+    stats.update(dict(rate_nans=0 if rate_nans == 0 else 1, distortion_nans=0 if distortion_nans == 0 else 1))
+    stats = get_cpu_stats_over_ranks(stats)
+
+    skipped_updates = 1
+    # only update if no rank has a nan and if the grad norm is below a specific threshold
+    if stats['distortion_nans'] == 0 and stats['rate_nans'] == 0 and (H.skip_threshold == -1 or grad_norm < H.skip_threshold):
+        optimizer.step()
+        skipped_updates = 0
+        update_ema(vae, ema_vae, H.ema_rate)
+
+    t1 = time.time()
+    stats.update(skipped_updates=skipped_updates, iter_time=t1 - t0, grad_norm=grad_norm)
+    return stats
+
+
+def eval_step(data_input, target, ema_vae):
+    with torch.no_grad():
+        stats = ema_vae.forward(data_input, target)
+    stats = get_cpu_stats_over_ranks(stats)
+    return stats
+
+
+def get_sample_for_visualization(data, preprocess_fn, num, dataset):
+    for x in DataLoader(data, batch_size=num):
+        break
+    orig_image = (x[0] * 255.0).to(torch.uint8).permute(0, 2, 3, 1) if dataset == 'ffhq_1024' else x[0]
+    preprocessed = preprocess_fn(x)[0]
+    return orig_image, preprocessed
+
+
+def train_loop(H, data_train, data_valid, preprocess_fn, vae, ema_vae, logprint):
+    optimizer, scheduler, cur_eval_loss, iterate, starting_epoch = load_opt(H, vae, logprint)
+    train_sampler = DistributedSampler(data_train, num_replicas=H.mpi_size, rank=H.rank)
+    viz_batch_original, viz_batch_processed = get_sample_for_visualization(data_valid, preprocess_fn, H.num_images_visualize, H.dataset)
+    early_evals = set([1] + [2 ** exp for exp in range(3, 14)])
+    stats = []
+    iters_since_starting = 0
+    H.ema_rate = torch.as_tensor(H.ema_rate).cuda()
+    for epoch in range(starting_epoch, H.num_epochs):
+        train_sampler.set_epoch(epoch)
+        for x in DataLoader(data_train, batch_size=H.n_batch, drop_last=True, pin_memory=True, sampler=train_sampler):
+            data_input, target = preprocess_fn(x)
+            training_stats = training_step(H, data_input, target, vae, ema_vae, optimizer, iterate)
+            stats.append(training_stats)
+            scheduler.step()
+            if iterate % H.iters_per_print == 0 or iters_since_starting in early_evals:
+                logprint(model=H.desc, type='train_loss', lr=scheduler.get_last_lr()[0], epoch=epoch, step=iterate, **accumulate_stats(stats, H.iters_per_print))
+
+            if iterate % H.iters_per_images == 0 or (iters_since_starting in early_evals and H.dataset != 'ffhq_1024') and H.rank == 0:
+                write_images(H, ema_vae, viz_batch_original, viz_batch_processed, f'{H.save_dir}/samples-{iterate}.png', logprint)
+
+            iterate += 1
+            iters_since_starting += 1
+            if iterate % H.iters_per_save == 0 and H.rank == 0:
+                if np.isfinite(stats[-1]['elbo']):
+                    logprint(model=H.desc, type='train_loss', epoch=epoch, step=iterate, **accumulate_stats(stats, H.iters_per_print))
+                    fp = os.path.join(H.save_dir, 'latest')
+                    logprint(f'Saving model@ {iterate} to {fp}')
+                    save_model(fp, vae, ema_vae, optimizer, H)
+
+            if iterate % H.iters_per_ckpt == 0 and H.rank == 0:
+                save_model(os.path.join(H.save_dir, f'iter-{iterate}'), vae, ema_vae, optimizer, H)
+
+        if epoch % H.epochs_per_eval == 0:
+            valid_stats = evaluate(H, ema_vae, data_valid, preprocess_fn)
+            logprint(model=H.desc, type='eval_loss', epoch=epoch, step=iterate, **valid_stats)
+
+
+def evaluate(H, ema_vae, data_valid, preprocess_fn):
+    stats_valid = []
+    valid_sampler = DistributedSampler(data_valid, num_replicas=H.mpi_size, rank=H.rank)
+    for x in DataLoader(data_valid, batch_size=H.n_batch, drop_last=True, pin_memory=True, sampler=valid_sampler):
+        data_input, target = preprocess_fn(x)
+        stats_valid.append(eval_step(data_input, target, ema_vae))
+    vals = [a['elbo'] for a in stats_valid]
+    finites = np.array(vals)[np.isfinite(vals)]
+    stats = dict(n_batches=len(vals), filtered_elbo=np.mean(finites), **{k: np.mean([a[k] for a in stats_valid]) for k in stats_valid[-1]})
+    return stats
+
+
+def write_images(H, ema_vae, viz_batch_original, viz_batch_processed, fname, logprint):
+    zs = [s['z'].cuda() for s in ema_vae.forward_get_latents(viz_batch_processed)]
+    batches = [viz_batch_original.numpy()]
+    mb = viz_batch_processed.shape[0]
+    lv_points = np.floor(np.linspace(0, 1, H.num_variables_visualize + 2) * len(zs)).astype(int)[1:-1]
+    for i in lv_points:
+        batches.append(ema_vae.forward_samples_set_latents(mb, zs[:i], t=0.1))
+    for t in [1.0, 0.9, 0.8, 0.7][:H.num_temperatures_visualize]:
+        batches.append(ema_vae.forward_uncond_samples(mb, t=t))
+    n_rows = len(batches)
+    im = np.concatenate(batches, axis=0).reshape((n_rows, mb, *viz_batch_processed.shape[1:])).transpose([0, 2, 1, 3, 4]).reshape([n_rows * viz_batch_processed.shape[1], mb * viz_batch_processed.shape[2], 3])
+    logprint(f'printing samples to {fname}')
+    imageio.imwrite(fname, im)
+
+
+def run_test_eval(H, ema_vae, data_test, preprocess_fn, logprint):
+    print('evaluating')
+    stats = evaluate(H, ema_vae, data_test, preprocess_fn)
+    print('test results')
+    for k in stats:
+        print(k, stats[k])
+    logprint(type='test_loss', **stats)
+
+
+def main():
+    H, logprint = set_up_hyperparams()
+    H, data_train, data_valid_or_test, preprocess_fn = set_up_data(H)
+    vae, ema_vae = load_vaes(H, logprint)
+    if H.test_eval:
+        run_test_eval(H, ema_vae, data_valid_or_test, preprocess_fn, logprint)
+    else:
+        train_loop(H, data_train, data_valid_or_test, preprocess_fn, vae, ema_vae, logprint)
+
+
+if __name__ == "__main__":
+    main()

MindEye_Imagery/src/vdvae/train_helpers.py

@@ -0,0 +1,190 @@
+import torch
+import numpy as np
+#from mpi4py import MPI
+import socket
+import argparse
+import os
+import json
+import subprocess
+from hps import Hyperparams, parse_args_and_update_hparams, add_vae_arguments
+from vdvae_utils import (logger,
+                   local_mpi_rank,
+                   mpi_size,
+                   maybe_download,
+                   mpi_rank)
+from data import mkdir_p
+from contextlib import contextmanager
+import torch.distributed as dist
+#from apex.optimizers import FusedAdam as AdamW
+from vae import VAE
+from torch.nn.parallel.distributed import DistributedDataParallel
+
+
+def update_ema(vae, ema_vae, ema_rate):
+    for p1, p2 in zip(vae.parameters(), ema_vae.parameters()):
+        p2.data.mul_(ema_rate)
+        p2.data.add_(p1.data * (1 - ema_rate))
+
+
+def save_model(path, vae, ema_vae, optimizer, H):
+    torch.save(vae.state_dict(), f'{path}-model.th')
+    torch.save(ema_vae.state_dict(), f'{path}-model-ema.th')
+    torch.save(optimizer.state_dict(), f'{path}-opt.th')
+    from_log = os.path.join(H.save_dir, 'log.jsonl')
+    to_log = f'{os.path.dirname(path)}/{os.path.basename(path)}-log.jsonl'
+    subprocess.check_output(['cp', from_log, to_log])
+
+
+def accumulate_stats(stats, frequency):
+    z = {}
+    for k in stats[-1]:
+        if k in ['distortion_nans', 'rate_nans', 'skipped_updates', 'gcskip']:
+            z[k] = np.sum([a[k] for a in stats[-frequency:]])
+        elif k == 'grad_norm':
+            vals = [a[k] for a in stats[-frequency:]]
+            finites = np.array(vals)[np.isfinite(vals)]
+            if len(finites) == 0:
+                z[k] = 0.0
+            else:
+                z[k] = np.max(finites)
+        elif k == 'elbo':
+            vals = [a[k] for a in stats[-frequency:]]
+            finites = np.array(vals)[np.isfinite(vals)]
+            z['elbo'] = np.mean(vals)
+            z['elbo_filtered'] = np.mean(finites)
+        elif k == 'iter_time':
+            z[k] = stats[-1][k] if len(stats) < frequency else np.mean([a[k] for a in stats[-frequency:]])
+        else:
+            z[k] = np.mean([a[k] for a in stats[-frequency:]])
+    return z
+
+
+def linear_warmup(warmup_iters):
+    def f(iteration):
+        return 1.0 if iteration > warmup_iters else iteration / warmup_iters
+    return f
+
+
+def setup_mpi(H):
+    H.mpi_size = mpi_size()
+    H.local_rank = local_mpi_rank()
+    H.rank = mpi_rank()
+    os.environ["RANK"] = str(H.rank)
+    os.environ["WORLD_SIZE"] = str(H.mpi_size)
+    os.environ["MASTER_PORT"] = str(H.port)
+    # os.environ["NCCL_LL_THRESHOLD"] = "0"
+    os.environ["MASTER_ADDR"] = MPI.COMM_WORLD.bcast(socket.gethostname(), root=0)
+    torch.cuda.set_device(H.local_rank)
+    dist.init_process_group(backend='nccl', init_method=f"env://")
+
+
+def distributed_maybe_download(path, local_rank, mpi_size):
+    if not path.startswith('gs://'):
+        return path
+    filename = path[5:].replace('/', '-')
+    with first_rank_first(local_rank, mpi_size):
+        fp = maybe_download(path, filename)
+    return fp
+
+
+@contextmanager
+def first_rank_first(local_rank, mpi_size):
+    if mpi_size > 1 and local_rank > 0:
+        dist.barrier()
+
+    try:
+        yield
+    finally:
+        if mpi_size > 1 and local_rank == 0:
+            dist.barrier()
+
+
+def setup_save_dirs(H):
+    H.save_dir = os.path.join(H.save_dir, H.desc)
+    mkdir_p(H.save_dir)
+    H.logdir = os.path.join(H.save_dir, 'log')
+
+
+def set_up_hyperparams(s=None):
+    H = Hyperparams()
+    parser = argparse.ArgumentParser()
+    parser = add_vae_arguments(parser)
+    parse_args_and_update_hparams(H, parser, s=s)
+    setup_mpi(H)
+    setup_save_dirs(H)
+    logprint = logger(H.logdir)
+    for i, k in enumerate(sorted(H)):
+        logprint(type='hparam', key=k, value=H[k])
+    np.random.seed(H.seed)
+    torch.manual_seed(H.seed)
+    torch.cuda.manual_seed(H.seed)
+    logprint('training model', H.desc, 'on', H.dataset)
+    return H, logprint
+
+
+def restore_params(model, path, local_rank, mpi_size, map_ddp=True, map_cpu=False):
+    state_dict = torch.load(distributed_maybe_download(path, local_rank, mpi_size), map_location='cpu' if map_cpu else None)
+    if map_ddp:
+        new_state_dict = {}
+        l = len('module.')
+        for k in state_dict:
+            if k.startswith('module.'):
+                new_state_dict[k[l:]] = state_dict[k]
+            else:
+                new_state_dict[k] = state_dict[k]
+        state_dict = new_state_dict
+    model.load_state_dict(state_dict)
+
+
+def restore_log(path, local_rank, mpi_size):
+    loaded = [json.loads(l) for l in open(distributed_maybe_download(path, local_rank, mpi_size))]
+    try:
+        cur_eval_loss = min([z['elbo'] for z in loaded if 'type' in z and z['type'] == 'eval_loss'])
+    except ValueError:
+        cur_eval_loss = float('inf')
+    starting_epoch = max([z['epoch'] for z in loaded if 'type' in z and z['type'] == 'train_loss'])
+    iterate = max([z['step'] for z in loaded if 'type' in z and z['type'] == 'train_loss'])
+    return cur_eval_loss, iterate, starting_epoch
+
+
+def load_vaes(H, logprint):
+    vae = VAE(H)
+    if H.restore_path:
+        logprint(f'Restoring vae from {H.restore_path}')
+        restore_params(vae, H.restore_path, map_cpu=True, local_rank=H.local_rank, mpi_size=H.mpi_size)
+
+    ema_vae = VAE(H)
+    if H.restore_ema_path:
+        logprint(f'Restoring ema vae from {H.restore_ema_path}')
+        restore_params(ema_vae, H.restore_ema_path, map_cpu=True, local_rank=H.local_rank, mpi_size=H.mpi_size)
+    else:
+        ema_vae.load_state_dict(vae.state_dict())
+    ema_vae.requires_grad_(False)
+
+    vae = vae.cuda(H.local_rank)
+    ema_vae = ema_vae.cuda(H.local_rank)
+
+    vae = DistributedDataParallel(vae, device_ids=[H.local_rank], output_device=H.local_rank)
+
+    if len(list(vae.named_parameters())) != len(list(vae.parameters())):
+        raise ValueError('Some params are not named. Please name all params.')
+    total_params = 0
+    for name, p in vae.named_parameters():
+        total_params += np.prod(p.shape)
+    logprint(total_params=total_params, readable=f'{total_params:,}')
+    return vae, ema_vae
+
+
+def load_opt(H, vae, logprint):
+    optimizer = AdamW(vae.parameters(), weight_decay=H.wd, lr=H.lr, betas=(H.adam_beta1, H.adam_beta2))
+    scheduler = torch.optim.lr_scheduler.LambdaLR(optimizer, lr_lambda=linear_warmup(H.warmup_iters))
+
+    if H.restore_optimizer_path:
+        optimizer.load_state_dict(
+            torch.load(distributed_maybe_download(H.restore_optimizer_path, H.local_rank, H.mpi_size), map_location='cpu'))
+    if H.restore_log_path:
+        cur_eval_loss, iterate, starting_epoch = restore_log(H.restore_log_path, H.local_rank, H.mpi_size)
+    else:
+        cur_eval_loss, iterate, starting_epoch = float('inf'), 0, 0
+    logprint('starting at epoch', starting_epoch, 'iterate', iterate, 'eval loss', cur_eval_loss)
+    return optimizer, scheduler, cur_eval_loss, iterate, starting_epoch

MindEye_Imagery/src/vdvae/vae.py

@@ -0,0 +1,252 @@
+import torch
+from torch import nn
+from torch.nn import functional as F
+from vae_helpers import HModule, get_1x1, get_3x3, DmolNet, draw_gaussian_diag_samples, gaussian_analytical_kl
+from collections import defaultdict
+import numpy as np
+import itertools
+
+
+class Block(nn.Module):
+    def __init__(self, in_width, middle_width, out_width, down_rate=None, residual=False, use_3x3=True, zero_last=False):
+        super().__init__()
+        self.down_rate = down_rate
+        self.residual = residual
+        self.c1 = get_1x1(in_width, middle_width)
+        self.c2 = get_3x3(middle_width, middle_width) if use_3x3 else get_1x1(middle_width, middle_width)
+        self.c3 = get_3x3(middle_width, middle_width) if use_3x3 else get_1x1(middle_width, middle_width)
+        self.c4 = get_1x1(middle_width, out_width, zero_weights=zero_last)
+
+    def forward(self, x):
+        xhat = self.c1(F.gelu(x))
+        xhat = self.c2(F.gelu(xhat))
+        xhat = self.c3(F.gelu(xhat))
+        xhat = self.c4(F.gelu(xhat))
+        out = x + xhat if self.residual else xhat
+        if self.down_rate is not None:
+            out = F.avg_pool2d(out, kernel_size=self.down_rate, stride=self.down_rate)
+        return out
+
+
+def parse_layer_string(s):
+    layers = []
+    for ss in s.split(','):
+        if 'x' in ss:
+            res, num = ss.split('x')
+            count = int(num)
+            layers += [(int(res), None) for _ in range(count)]
+        elif 'm' in ss:
+            res, mixin = [int(a) for a in ss.split('m')]
+            layers.append((res, mixin))
+        elif 'd' in ss:
+            res, down_rate = [int(a) for a in ss.split('d')]
+            layers.append((res, down_rate))
+        else:
+            res = int(ss)
+            layers.append((res, None))
+    return layers
+
+
+def pad_channels(t, width):
+    d1, d2, d3, d4 = t.shape
+    empty = torch.zeros(d1, width, d3, d4, device=t.device)
+    empty[:, :d2, :, :] = t
+    return empty
+
+
+def get_width_settings(width, s):
+    mapping = defaultdict(lambda: width)
+    if s:
+        s = s.split(',')
+        for ss in s:
+            k, v = ss.split(':')
+            mapping[int(k)] = int(v)
+    return mapping
+
+
+class Encoder(HModule):
+    def build(self):
+        H = self.H
+        self.in_conv = get_3x3(H.image_channels, H.width)
+        self.widths = get_width_settings(H.width, H.custom_width_str)
+        enc_blocks = []
+        blockstr = parse_layer_string(H.enc_blocks)
+        for res, down_rate in blockstr:
+            use_3x3 = res > 2  # Don't use 3x3s for 1x1, 2x2 patches
+            enc_blocks.append(Block(self.widths[res], int(self.widths[res] * H.bottleneck_multiple), self.widths[res], down_rate=down_rate, residual=True, use_3x3=use_3x3))
+        n_blocks = len(blockstr)
+        for b in enc_blocks:
+            b.c4.weight.data *= np.sqrt(1 / n_blocks)
+        self.enc_blocks = nn.ModuleList(enc_blocks)
+
+    def forward(self, x):
+        x = x.permute(0, 3, 1, 2).contiguous()
+        x = self.in_conv(x)
+        activations = {}
+        activations[x.shape[2]] = x
+        for block in self.enc_blocks:
+            x = block(x)
+            res = x.shape[2]
+            x = x if x.shape[1] == self.widths[res] else pad_channels(x, self.widths[res])
+            activations[res] = x
+        return activations
+
+
+class DecBlock(nn.Module):
+    def __init__(self, H, res, mixin, n_blocks):
+        super().__init__()
+        self.base = res
+        self.mixin = mixin
+        self.H = H
+        self.widths = get_width_settings(H.width, H.custom_width_str)
+        width = self.widths[res]
+        use_3x3 = res > 2
+        cond_width = int(width * H.bottleneck_multiple)
+        self.zdim = H.zdim
+        self.enc = Block(width * 2, cond_width, H.zdim * 2, residual=False, use_3x3=use_3x3)
+        self.prior = Block(width, cond_width, H.zdim * 2 + width, residual=False, use_3x3=use_3x3, zero_last=True)
+        self.z_proj = get_1x1(H.zdim, width)
+        self.z_proj.weight.data *= np.sqrt(1 / n_blocks)
+        self.resnet = Block(width, cond_width, width, residual=True, use_3x3=use_3x3)
+        self.resnet.c4.weight.data *= np.sqrt(1 / n_blocks)
+        self.z_fn = lambda x: self.z_proj(x)
+
+    def sample(self, x, acts):
+        qm, qv = self.enc(torch.cat([x, acts], dim=1)).chunk(2, dim=1)
+        feats = self.prior(x)
+        pm, pv, xpp = feats[:, :self.zdim, ...], feats[:, self.zdim:self.zdim * 2, ...], feats[:, self.zdim * 2:, ...]
+        x = x + xpp
+        z = draw_gaussian_diag_samples(qm, qv)
+        kl = gaussian_analytical_kl(qm, pm, qv, pv)
+        return z, x, kl
+
+    def sample_uncond(self, x, t=None, lvs=None):
+        n, c, h, w = x.shape
+        feats = self.prior(x)
+        pm, pv, xpp = feats[:, :self.zdim, ...], feats[:, self.zdim:self.zdim * 2, ...], feats[:, self.zdim * 2:, ...]
+        x = x + xpp
+        if lvs is not None:
+            z = lvs
+        else:
+            if t is not None:
+                pv = pv + torch.ones_like(pv) * np.log(t)
+            z = draw_gaussian_diag_samples(pm, pv)
+        return z, x
+
+    def get_inputs(self, xs, activations):
+        acts = activations[self.base]
+        try:
+            x = xs[self.base]
+        except KeyError:
+            x = torch.zeros_like(acts)
+        if acts.shape[0] != x.shape[0]:
+            x = x.repeat(acts.shape[0], 1, 1, 1)
+        return x, acts
+
+    def forward(self, xs, activations, get_latents=False):
+        x, acts = self.get_inputs(xs, activations)
+        if self.mixin is not None:
+            x = x + F.interpolate(xs[self.mixin][:, :x.shape[1], ...], scale_factor=self.base // self.mixin)
+        z, x, kl = self.sample(x, acts)
+        x = x + self.z_fn(z)
+        x = self.resnet(x)
+        xs[self.base] = x
+        if get_latents:
+            return xs, dict(z=z.detach(), kl=kl)
+        return xs, dict(kl=kl)
+
+    def forward_uncond(self, xs, t=None, lvs=None):
+        try:
+            x = xs[self.base]
+        except KeyError:
+            ref = xs[list(xs.keys())[0]]
+            x = torch.zeros(dtype=ref.dtype, size=(ref.shape[0], self.widths[self.base], self.base, self.base), device=ref.device)
+        if self.mixin is not None:
+            x = x + F.interpolate(xs[self.mixin][:, :x.shape[1], ...], scale_factor=self.base // self.mixin)
+        z, x = self.sample_uncond(x, t, lvs=lvs)
+        x = x + self.z_fn(z)
+        x = self.resnet(x)
+        xs[self.base] = x
+        return xs
+
+
+class Decoder(HModule):
+
+    def build(self):
+        H = self.H
+        resos = set()
+        dec_blocks = []
+        self.widths = get_width_settings(H.width, H.custom_width_str)
+        blocks = parse_layer_string(H.dec_blocks)
+        for idx, (res, mixin) in enumerate(blocks):
+            dec_blocks.append(DecBlock(H, res, mixin, n_blocks=len(blocks)))
+            resos.add(res)
+        self.resolutions = sorted(resos)
+        self.dec_blocks = nn.ModuleList(dec_blocks)
+        self.bias_xs = nn.ParameterList([nn.Parameter(torch.zeros(1, self.widths[res], res, res)) for res in self.resolutions if res <= H.no_bias_above])
+        self.out_net = DmolNet(H)
+        self.gain = nn.Parameter(torch.ones(1, H.width, 1, 1))
+        self.bias = nn.Parameter(torch.zeros(1, H.width, 1, 1))
+        self.final_fn = lambda x: x * self.gain + self.bias
+
+    def forward(self, activations, get_latents=False):
+        stats = []
+        xs = {a.shape[2]: a for a in self.bias_xs}
+        for block in self.dec_blocks:
+            xs, block_stats = block(xs, activations, get_latents=get_latents)
+            stats.append(block_stats)
+        xs[self.H.image_size] = self.final_fn(xs[self.H.image_size])
+        return xs[self.H.image_size], stats
+
+    def forward_uncond(self, n, t=None, y=None):
+        xs = {}
+        for bias in self.bias_xs:
+            xs[bias.shape[2]] = bias.repeat(n, 1, 1, 1)
+        for idx, block in enumerate(self.dec_blocks):
+            try:
+                temp = t[idx]
+            except TypeError:
+                temp = t
+            xs = block.forward_uncond(xs, temp)
+        xs[self.H.image_size] = self.final_fn(xs[self.H.image_size])
+        return xs[self.H.image_size]
+
+    def forward_manual_latents(self, n, latents, t=None):
+        xs = {}
+        for bias in self.bias_xs:
+            xs[bias.shape[2]] = bias.repeat(n, 1, 1, 1)
+        for block, lvs in itertools.zip_longest(self.dec_blocks, latents):
+            xs = block.forward_uncond(xs, t, lvs=lvs)
+        xs[self.H.image_size] = self.final_fn(xs[self.H.image_size])
+        return xs[self.H.image_size]
+
+
+class VAE(HModule):
+    def build(self):
+        self.encoder = Encoder(self.H)
+        self.decoder = Decoder(self.H)
+
+    def forward(self, x, x_target):
+        activations = self.encoder.forward(x)
+        px_z, stats = self.decoder.forward(activations)
+        distortion_per_pixel = self.decoder.out_net.nll(px_z, x_target)
+        rate_per_pixel = torch.zeros_like(distortion_per_pixel)
+        ndims = np.prod(x.shape[1:])
+        for statdict in stats:
+            rate_per_pixel += statdict['kl'].sum(dim=(1, 2, 3))
+        rate_per_pixel /= ndims
+        elbo = (distortion_per_pixel + rate_per_pixel).mean()
+        return dict(elbo=elbo, distortion=distortion_per_pixel.mean(), rate=rate_per_pixel.mean())
+
+    def forward_get_latents(self, x):
+        activations = self.encoder.forward(x)
+        _, stats = self.decoder.forward(activations, get_latents=True)
+        return stats
+
+    def forward_uncond_samples(self, n_batch, t=None):
+        px_z = self.decoder.forward_uncond(n_batch, t=t)
+        return self.decoder.out_net.sample(px_z)
+
+    def forward_samples_set_latents(self, n_batch, latents, t=None):
+        px_z = self.decoder.forward_manual_latents(n_batch, latents, t=t)
+        return self.decoder.out_net.sample(px_z)

MindEye_Imagery/src/vdvae/vae_helpers.py

@@ -0,0 +1,165 @@
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+@torch.jit.script
+def gaussian_analytical_kl(mu1, mu2, logsigma1, logsigma2):
+    return -0.5 + logsigma2 - logsigma1 + 0.5 * (logsigma1.exp() ** 2 + (mu1 - mu2) ** 2) / (logsigma2.exp() ** 2)
+
+
+@torch.jit.script
+def draw_gaussian_diag_samples(mu, logsigma):
+    eps = torch.empty_like(mu).normal_(0., 1.)
+    return torch.exp(logsigma) * eps + mu
+
+
+def get_conv(in_dim, out_dim, kernel_size, stride, padding, zero_bias=True, zero_weights=False, groups=1, scaled=False):
+    c = nn.Conv2d(in_dim, out_dim, kernel_size, stride, padding, groups=groups)
+    if zero_bias:
+        c.bias.data *= 0.0
+    if zero_weights:
+        c.weight.data *= 0.0
+    return c
+
+
+def get_3x3(in_dim, out_dim, zero_bias=True, zero_weights=False, groups=1, scaled=False):
+    return get_conv(in_dim, out_dim, 3, 1, 1, zero_bias, zero_weights, groups=groups, scaled=scaled)
+
+
+def get_1x1(in_dim, out_dim, zero_bias=True, zero_weights=False, groups=1, scaled=False):
+    return get_conv(in_dim, out_dim, 1, 1, 0, zero_bias, zero_weights, groups=groups, scaled=scaled)
+
+
+def log_prob_from_logits(x):
+    """ numerically stable log_softmax implementation that prevents overflow """
+    axis = len(x.shape) - 1
+    m = x.max(dim=axis, keepdim=True)[0]
+    return x - m - torch.log(torch.exp(x - m).sum(dim=axis, keepdim=True))
+
+
+def const_max(t, constant):
+    other = torch.ones_like(t) * constant
+    return torch.max(t, other)
+
+
+def const_min(t, constant):
+    other = torch.ones_like(t) * constant
+    return torch.min(t, other)
+
+
+def discretized_mix_logistic_loss(x, l, low_bit=False):
+    """ log-likelihood for mixture of discretized logistics, assumes the data has been rescaled to [-1,1] interval """
+    # Adapted from https://github.com/openai/pixel-cnn/blob/master/pixel_cnn_pp/nn.py
+    xs = [s for s in x.shape]  # true image (i.e. labels) to regress to, e.g. (B,32,32,3)
+    ls = [s for s in l.shape]  # predicted distribution, e.g. (B,32,32,100)
+    nr_mix = int(ls[-1] / 10)  # here and below: unpacking the params of the mixture of logistics
+    logit_probs = l[:, :, :, :nr_mix]
+    l = torch.reshape(l[:, :, :, nr_mix:], xs + [nr_mix * 3])
+    means = l[:, :, :, :, :nr_mix]
+    log_scales = const_max(l[:, :, :, :, nr_mix:2 * nr_mix], -7.)
+    coeffs = torch.tanh(l[:, :, :, :, 2 * nr_mix:3 * nr_mix])
+    x = torch.reshape(x, xs + [1]) + torch.zeros(xs + [nr_mix]).to(x.device)  # here and below: getting the means and adjusting them based on preceding sub-pixels
+    m2 = torch.reshape(means[:, :, :, 1, :] + coeffs[:, :, :, 0, :] * x[:, :, :, 0, :], [xs[0], xs[1], xs[2], 1, nr_mix])
+    m3 = torch.reshape(means[:, :, :, 2, :] + coeffs[:, :, :, 1, :] * x[:, :, :, 0, :] + coeffs[:, :, :, 2, :] * x[:, :, :, 1, :], [xs[0], xs[1], xs[2], 1, nr_mix])
+    means = torch.cat([torch.reshape(means[:, :, :, 0, :], [xs[0], xs[1], xs[2], 1, nr_mix]), m2, m3], dim=3)
+    centered_x = x - means
+    inv_stdv = torch.exp(-log_scales)
+    if low_bit:
+        plus_in = inv_stdv * (centered_x + 1. / 31.)
+        cdf_plus = torch.sigmoid(plus_in)
+        min_in = inv_stdv * (centered_x - 1. / 31.)
+    else:
+        plus_in = inv_stdv * (centered_x + 1. / 255.)
+        cdf_plus = torch.sigmoid(plus_in)
+        min_in = inv_stdv * (centered_x - 1. / 255.)
+    cdf_min = torch.sigmoid(min_in)
+    log_cdf_plus = plus_in - F.softplus(plus_in)  # log probability for edge case of 0 (before scaling)
+    log_one_minus_cdf_min = -F.softplus(min_in)  # log probability for edge case of 255 (before scaling)
+    cdf_delta = cdf_plus - cdf_min  # probability for all other cases
+    mid_in = inv_stdv * centered_x
+    log_pdf_mid = mid_in - log_scales - 2. * F.softplus(mid_in)  # log probability in the center of the bin, to be used in extreme cases (not actually used in our code)
+
+    # now select the right output: left edge case, right edge case, normal case, extremely low prob case (doesn't actually happen for us)
+
+    # this is what we are really doing, but using the robust version below for extreme cases in other applications and to avoid NaN issue with tf.select()
+    # log_probs = tf.select(x < -0.999, log_cdf_plus, tf.select(x > 0.999, log_one_minus_cdf_min, tf.log(cdf_delta)))
+
+    # robust version, that still works if probabilities are below 1e-5 (which never happens in our code)
+    # tensorflow backpropagates through tf.select() by multiplying with zero instead of selecting: this requires use to use some ugly tricks to avoid potential NaNs
+    # the 1e-12 in tf.maximum(cdf_delta, 1e-12) is never actually used as output, it's purely there to get around the tf.select() gradient issue
+    # if the probability on a sub-pixel is below 1e-5, we use an approximation based on the assumption that the log-density is constant in the bin of the observed sub-pixel value
+    if low_bit:
+        log_probs = torch.where(x < -0.999,
+                                log_cdf_plus,
+                                torch.where(x > 0.999,
+                                            log_one_minus_cdf_min,
+                                            torch.where(cdf_delta > 1e-5,
+                                                        torch.log(const_max(cdf_delta, 1e-12)),
+                                                        log_pdf_mid - np.log(15.5))))
+    else:
+        log_probs = torch.where(x < -0.999,
+                                log_cdf_plus,
+                                torch.where(x > 0.999,
+                                            log_one_minus_cdf_min,
+                                            torch.where(cdf_delta > 1e-5,
+                                                        torch.log(const_max(cdf_delta, 1e-12)),
+                                                        log_pdf_mid - np.log(127.5))))
+    log_probs = log_probs.sum(dim=3) + log_prob_from_logits(logit_probs)
+    mixture_probs = torch.logsumexp(log_probs, -1)
+    return -1. * mixture_probs.sum(dim=[1, 2]) / np.prod(xs[1:])
+
+
+def sample_from_discretized_mix_logistic(l, nr_mix):
+    ls = [s for s in l.shape]
+    xs = ls[:-1] + [3]
+    # unpack parameters
+    logit_probs = l[:, :, :, :nr_mix]
+    l = torch.reshape(l[:, :, :, nr_mix:], xs + [nr_mix * 3])
+    # sample mixture indicator from softmax
+    eps = torch.empty(logit_probs.shape, device=l.device).uniform_(1e-5, 1. - 1e-5)
+    amax = torch.argmax(logit_probs - torch.log(-torch.log(eps)), dim=3)
+    sel = F.one_hot(amax, num_classes=nr_mix).float()
+    sel = torch.reshape(sel, xs[:-1] + [1, nr_mix])
+    # select logistic parameters
+    means = (l[:, :, :, :, :nr_mix] * sel).sum(dim=4)
+    log_scales = const_max((l[:, :, :, :, nr_mix:nr_mix * 2] * sel).sum(dim=4), -7.)
+    coeffs = (torch.tanh(l[:, :, :, :, nr_mix * 2:nr_mix * 3]) * sel).sum(dim=4)
+    # sample from logistic & clip to interval
+    # we don't actually round to the nearest 8bit value when sampling
+    u = torch.empty(means.shape, device=means.device).uniform_(1e-5, 1. - 1e-5)
+    x = means + torch.exp(log_scales) * (torch.log(u) - torch.log(1. - u))
+    x0 = const_min(const_max(x[:, :, :, 0], -1.), 1.)
+    x1 = const_min(const_max(x[:, :, :, 1] + coeffs[:, :, :, 0] * x0, -1.), 1.)
+    x2 = const_min(const_max(x[:, :, :, 2] + coeffs[:, :, :, 1] * x0 + coeffs[:, :, :, 2] * x1, -1.), 1.)
+    return torch.cat([torch.reshape(x0, xs[:-1] + [1]), torch.reshape(x1, xs[:-1] + [1]), torch.reshape(x2, xs[:-1] + [1])], dim=3)
+
+
+class HModule(nn.Module):
+    def __init__(self, H):
+        super().__init__()
+        self.H = H
+        self.build()
+
+
+class DmolNet(nn.Module):
+    def __init__(self, H):
+        super().__init__()
+        self.H = H
+        self.width = H.width
+        self.out_conv = get_conv(H.width, H.num_mixtures * 10, kernel_size=1, stride=1, padding=0)
+
+    def nll(self, px_z, x):
+        return discretized_mix_logistic_loss(x=x, l=self.forward(px_z), low_bit=self.H.dataset in ['ffhq_256'])
+
+    def forward(self, px_z):
+        xhat = self.out_conv(px_z)
+        return xhat.permute(0, 2, 3, 1)
+
+    def sample(self, px_z):
+        im = sample_from_discretized_mix_logistic(self.forward(px_z), self.H.num_mixtures)
+        xhat = (im + 1.0) * 127.5
+        xhat = xhat.detach().cpu().numpy()
+        xhat = np.minimum(np.maximum(0.0, xhat), 255.0).astype(np.uint8)
+        return xhat

MindEye_Imagery/src/vdvae/vdvae_utils.py

@@ -0,0 +1,140 @@
+#from mpi4py import MPI
+import os
+import json
+import tempfile
+import numpy as np
+import torch
+import time
+import subprocess
+import torch.distributed as dist
+
+
+def allreduce(x, average):
+    if mpi_size() > 1:
+        dist.all_reduce(x, dist.ReduceOp.SUM)
+    return x / mpi_size() if average else x
+
+
+def get_cpu_stats_over_ranks(stat_dict):
+    keys = sorted(stat_dict.keys())
+    allreduced = allreduce(torch.stack([torch.as_tensor(stat_dict[k]).detach().cuda().float() for k in keys]), average=True).cpu()
+    return {k: allreduced[i].item() for (i, k) in enumerate(keys)}
+
+
+class Hyperparams(dict):
+    def __getattr__(self, attr):
+        try:
+            return self[attr]
+        except KeyError:
+            return None
+
+    def __setattr__(self, attr, value):
+        self[attr] = value
+
+
+def logger(log_prefix):
+    'Prints the arguments out to stdout, .txt, and .jsonl files'
+
+    jsonl_path = f'{log_prefix}.jsonl'
+    txt_path = f'{log_prefix}.txt'
+
+    def log(*args, pprint=False, **kwargs):
+        if mpi_rank() != 0:
+            return
+        t = time.ctime()
+        argdict = {'time': t}
+        if len(args) > 0:
+            argdict['message'] = ' '.join([str(x) for x in args])
+        argdict.update(kwargs)
+
+        txt_str = []
+        args_iter = sorted(argdict) if pprint else argdict
+        for k in args_iter:
+            val = argdict[k]
+            if isinstance(val, np.ndarray):
+                val = val.tolist()
+            elif isinstance(val, np.integer):
+                val = int(val)
+            elif isinstance(val, np.floating):
+                val = float(val)
+            argdict[k] = val
+            if isinstance(val, float):
+                val = f'{val:.5f}'
+            txt_str.append(f'{k}: {val}')
+        txt_str = ', '.join(txt_str)
+
+        if pprint:
+            json_str = json.dumps(argdict, sort_keys=True)
+            txt_str = json.dumps(argdict, sort_keys=True, indent=4)
+        else:
+            json_str = json.dumps(argdict)
+
+        print(txt_str, flush=True)
+
+        with open(txt_path, "a+") as f:
+            print(txt_str, file=f, flush=True)
+        with open(jsonl_path, "a+") as f:
+            print(json_str, file=f, flush=True)
+
+    return log
+
+
+def maybe_download(path, filename=None):
+    '''If a path is a gsutil path, download it and return the local link,
+    otherwise return link'''
+    if not path.startswith('gs://'):
+        return path
+    if filename:
+        local_dest = f'/tmp/'
+        out_path = f'/tmp/{filename}'
+        if os.path.isfile(out_path):
+            return out_path
+        subprocess.check_output(['gsutil', '-m', 'cp', '-R', path, out_path])
+        return out_path
+    else:
+        local_dest = tempfile.mkstemp()[1]
+        subprocess.check_output(['gsutil', '-m', 'cp', path, local_dest])
+    return local_dest
+
+
+def tile_images(images, d1=4, d2=4, border=1):
+    id1, id2, c = images[0].shape
+    out = np.ones([d1 * id1 + border * (d1 + 1),
+                   d2 * id2 + border * (d2 + 1),
+                   c], dtype=np.uint8)
+    out *= 255
+    if len(images) != d1 * d2:
+        raise ValueError('Wrong num of images')
+    for imgnum, im in enumerate(images):
+        num_d1 = imgnum // d2
+        num_d2 = imgnum % d2
+        start_d1 = num_d1 * id1 + border * (num_d1 + 1)
+        start_d2 = num_d2 * id2 + border * (num_d2 + 1)
+        out[start_d1:start_d1 + id1, start_d2:start_d2 + id2, :] = im
+    return out
+
+
+def mpi_size():
+    return MPI.COMM_WORLD.Get_size()
+
+
+def mpi_rank():
+    return MPI.COMM_WORLD.Get_rank()
+
+
+def num_nodes():
+    nn = mpi_size()
+    if nn % 8 == 0:
+        return nn // 8
+    return nn // 8 + 1
+
+
+def gpus_per_node():
+    size = mpi_size()
+    if size > 1:
+        return max(size // num_nodes(), 1)
+    return 1
+
+
+def local_mpi_rank():
+    return mpi_rank() % gpus_per_node()

spurious_reconstruction/.gitignore

@@ -0,0 +1,13 @@
+# python generated files
+__pycache__/
+*.py[oc]
+build/
+dist/
+wheels/
+*.egg-info
+
+# venv
+.venv
+./results
+./analysis/my_playground
+./data

spurious_reconstruction/README.md

@@ -0,0 +1,58 @@
+# Spurious-reconstruction
+
+This repository contains demo code for the paper:  
+Ken Shirakawa, Yoshihiro Nagano, Misato Tanaka, Shuntaro C. Aoki, Yusuke Muraki, Kei Majima, and Yukiyasu Kamitani, "Spurious reconstruction from brain activity" ([arXiv](https://arxiv.org/abs/2405.10078)).
+
+## Getting Started
+
+### Installation
+To clone this repository on your local machine, use the following `git clone` command with the project URL:
+
+
+``` 
+git clone https://github.com/KamitaniLab/spurious_reconstruction.git
+```
+### Environment Setup
+The environment for this project was created using [Rye](https://rye.astral.sh/). To set it up, follow the official instructions to install Rye. After installation, you can synchronize the environment for this repository using:
+
+```
+cd spurious_reconstruction 
+rye sync
+```
+
+## Usage
+
+---
+
+We are currently preparing to share the preprocessed data and files. These files will be made available on Figshare at [this URL](https://figshare.com/articles/dataset/Spurious_reconstruction_from_brain_activity/27013342).
+
+Reproducing the full analysis requires a large amount of data (approximately 3TB) and some analyses take considerable time to complete (up to ~2 weeks). Therefore, we recommend downloading only the minimum necessary files to reproduce specific analysis, such as image reconstruction, UMAP results or simulation analysis.
+
+For example, if you want to reproduce the UMAP results (related to Figure 4), you can download the required data by running the following command:
+
+```
+python ./download.py "UMAP visualization analysis"
+```
+
+Alternatively, you can directly download the results of UMAP analysis with:
+```
+python ./download.py "UMAP visualization results"
+```
+Note: At the moment, only the download of **result data** is supported. Please wait a little longer for the availability of **analysis data**.
+
+Each analysis directory contains additional information on how to reproduce the corresponding results. Please check the README files within each directory for specific instructions.
+
+
+### Original Dataset
+We used publicly available datasets for this project. The key datasets are:
+
+- **NSD (Natural Scenes Dataset)**  
+    - Raw fMRI data (including visual images and text annotations): Available upon request via [Spurious Reconstruction@figshare](https://forms.gle/eT4jHxaWwYUDEf2i9).
+  
+- **Deep Image Reconstruction**  
+    - Raw fMRI data: Available from [OpenNeuro](https://openneuro.org/datasets/ds001506).
+    - Preprocessed fMRI data, DNN features extracted from images, and decoded DNN features: Available via [Deep Image Reconstruction@figshare](https://github.com/KamitaniLab/DeepImageReconstruction?tab=readme-ov-file#:~:text=Preprocessed%20fMRI%20data,Image%20Reconstruction%40figshare).
+    - Visual images: Available upon request via [this form](https://forms.gle/ujvA34948Xg49jdn9).
+    - Text annotations: Available from [GOD Stimuli Annotations](https://github.com/KamitaniLab/GOD_stimuli_annotations).
+
+If you want to fully reproduce the analysis, you can use these brain data, image stimuli and corresponding text annotations.

spurious_reconstruction/analysis/0_preprocessing/.ipynb_checkpoints/extract_features-checkpoint.ipynb

@@ -0,0 +1,437 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": 1,
+   "id": "5082cdc1-b052-433d-b018-29ba6fa3d313",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# %% [Block 1]\n",
+    "from glob import glob\n",
+    "from tqdm.auto import tqdm\n",
+    "import os\n",
+    "from PIL import Image\n",
+    "import numpy as np\n",
+    "import torch\n",
+    "from scipy.io import savemat\n",
+    "from bdpy.dl.torch import FeatureExtractor\n",
+    "from torchvision import models, transforms\n",
+    "import yaml\n",
+    "import h5py"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 2,
+   "id": "a31be173-61ff-4adf-a558-63a3efe74cff",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# %% [Block 2]\n",
+    "# Custom transform class to convert RGB to BGR\n",
+    "class ConvertRGBtoBGR:\n",
+    "    def __call__(self, image):\n",
+    "        return image[[2, 1, 0], :, :]  \n",
+    "    \n",
+    "# Custom transform class to convert Image format to tensor while keeping pixel ranges\n",
+    "class ToTensorWithoutScaling:\n",
+    "    def __call__(self, image):\n",
+    "        image = np.array(image).astype(np.float32)  # Convert PIL Image to NumPy array\n",
+    "        tensor = torch.from_numpy(image).permute(2, 0, 1).float()  # Convert NumPy array to tensor and rearrange dimensions\n",
+    "        return tensor\n",
+    "    \n",
+    "class Convert32to16:\n",
+    "    def __call__(self, image):\n",
+    "        return image.type(torch.float16)\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 34,
+   "id": "969abfe4-e4b6-4fa2-bd09-17f77ddf2292",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# %% [Block 3]\n",
+    "def load_model(network):\n",
+    "    \n",
+    "    if network == \"VGG19_ILSVRC_19_layers\":\n",
+    "        from bdpy.dl.torch.models import layer_map, model_factory\n",
+    "        model = model_factory('vgg19')\n",
+    "        encoder_param_file = '/weka/proj-fmri/ckadirt/spurious_reconstruction/analysis/VGG_ILSVRC_19_layers/VGG_ILSVRC_19_layers.pt'\n",
+    "        model.load_state_dict(torch.load(encoder_param_file))\n",
+    "        model.eval()\n",
+    "        layer_mapping = layer_map('vgg19')\n",
+    "        mean_image = [104., 117., 123.] # BGR\n",
+    "        preprocess = transforms.Compose([\n",
+    "                    transforms.Resize((224, 224), interpolation=Image.BICUBIC),\n",
+    "                    # ToTensorWithoutScaling(),  # Custom transform added\n",
+    "                    # ConvertRGBtoBGR(),  # Custom transform added\n",
+    "                    transforms.Normalize(mean=mean_image, std=[1.,1.,1.]),\n",
+    "                ])\n",
+    "        def preprocess_and_check(img):\n",
+    "            print(\"Original shape:\", img.shape)\n",
+    "            img = transforms.Resize((224, 224), interpolation=Image.BICUBIC)(img)\n",
+    "            print(\"After ToTensor:\", img.shape)\n",
+    "            print(\"After Resize:\", img.shape)\n",
+    "            img = transforms.Normalize(mean=mean_image, std=[1.,1.,1.])(img)\n",
+    "            return img\n",
+    "        \n",
+    "    elif network == \"vgg19_torchvision\":\n",
+    "        model = models.vgg19(pretrained=True)\n",
+    "        model.eval()\n",
+    "        \n",
+    "        preprocess = transforms.Compose([\n",
+    "            transforms.Resize((224, 224)), \n",
+    "            transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),\n",
+    "        ]) \n",
+    "        def preprocess_and_check(img):\n",
+    "            print(\"Original shape:\", img.shape)\n",
+    "            img = transforms.ToTensor()(img)\n",
+    "            print(\"After ToTensor:\", img.shape)\n",
+    "            img = transforms.Resize(224)(img)\n",
+    "            print(\"After Resize:\", img.shape)\n",
+    "            img = transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])(img)\n",
+    "            return img\n",
+    "    else:\n",
+    "        raise ValueError(\"Network not supported. Please choose from: VGG19_ILSVRC_19_layers\"\n",
+    "                        +\"or define by yourself. The output should be model and preprocessing functions from RGB image to model input.\")\n",
+    "    \n",
+    "    return model, preprocess\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 35,
+   "id": "8020d232-72cd-42ed-bb8e-183c8c0ee9b6",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# %% [Block 4]\n",
+    "def load_config(config_path):\n",
+    "    '''Load configuration file.'''\n",
+    "    with open(config_path, 'r') as file:\n",
+    "        config = yaml.safe_load(file)\n",
+    "    return config\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 36,
+   "id": "a766c128-7e43-419c-9d7e-6f317e44753e",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "/weka/proj-medarc/shared/mindeyev2_dataset//wds/subj01/train/{0..39}.tar\n",
+      "/weka/proj-medarc/shared/mindeyev2_dataset//wds/subj01/new_test/0.tar\n",
+      "Loaded test dl for subj1!\n",
+      "\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "65cb44139d1e4cf788e5a8956d9c1eb1",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "0it [00:00, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "981c8b8963d84fbf9994421e1e8ce02b",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "0it [00:00, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    }
+   ],
+   "source": [
+    "import webdataset as wds\n",
+    "import random\n",
+    "\n",
+    "def my_split_by_node(urls): return urls\n",
+    "data_path = '/weka/proj-medarc/shared/mindeyev2_dataset/'\n",
+    "subj_list = [1]\n",
+    "num_sessions = 40\n",
+    "multi_subject = False\n",
+    "batch_size = 1\n",
+    "data_type = torch.float16\n",
+    "subj = 1\n",
+    "\n",
+    "train_data = {}\n",
+    "train_dl = {}\n",
+    "num_voxels = {}\n",
+    "voxels = {}\n",
+    "\n",
+    "train_url = f\"{data_path}/wds/subj0{subj}/train/\" + \"{0..\" + f\"{num_sessions-1}\" + \"}.tar\"\n",
+    "print(train_url)\n",
+    "\n",
+    "train_data = wds.WebDataset(train_url,resampled=False,nodesplitter=my_split_by_node)\\\n",
+    "                    .decode(\"torch\")\\\n",
+    "                    .rename(behav=\"behav.npy\", past_behav=\"past_behav.npy\", future_behav=\"future_behav.npy\", olds_behav=\"olds_behav.npy\")\\\n",
+    "                    .to_tuple(*[\"behav\", \"past_behav\", \"future_behav\", \"olds_behav\"])\n",
+    "train_dl = torch.utils.data.DataLoader(train_data, batch_size=batch_size, shuffle=False, drop_last=False, pin_memory=True)\n",
+    "\n",
+    "# f = h5py.File(f'{data_path}/betas_all_subj0{s}_fp32_renorm.hdf5', 'r')\n",
+    "# betas = f['betas'][:]\n",
+    "# betas = torch.Tensor(betas).to(\"cpu\").to(data_type)\n",
+    "# num_voxels_list.append(betas[0].shape[-1])\n",
+    "# num_voxels[f'subj0{s}'] = betas[0].shape[-1]\n",
+    "# voxels[f'subj0{s}'] = betas\n",
+    "# print(f\"num_voxels for subj0{s}: {num_voxels[f'subj0{s}']}\")\n",
+    "\n",
+    "\n",
+    "num_test=3000\n",
+    "test_url = f\"{data_path}/wds/subj0{subj}/new_test/\" + \"0.tar\"\n",
+    "print(test_url)\n",
+    "test_data = wds.WebDataset(test_url,resampled=False,nodesplitter=my_split_by_node)\\\n",
+    "                    .decode(\"torch\")\\\n",
+    "                    .rename(behav=\"behav.npy\", past_behav=\"past_behav.npy\", future_behav=\"future_behav.npy\", olds_behav=\"olds_behav.npy\")\\\n",
+    "                    .to_tuple(*[\"behav\", \"past_behav\", \"future_behav\", \"olds_behav\"])\n",
+    "test_dl = torch.utils.data.DataLoader(test_data, batch_size=1, shuffle=False, drop_last=True, pin_memory=True)\n",
+    "print(f\"Loaded test dl for subj{subj}!\\n\")\n",
+    "\n",
+    "all_indexes_train = []\n",
+    "\n",
+    "for behav0, past_behav0, future_behav0, old_behav0 in tqdm(train_dl):\n",
+    "    img_idx = behav0[:,0,0].cpu().long().numpy()[0]\n",
+    "    all_indexes_train.append(img_idx)\n",
+    "\n",
+    "all_indexes_test = []\n",
+    "\n",
+    "for behav0, past_behav0, future_behav0, old_behav0 in tqdm(test_dl):\n",
+    "    img_idx = behav0[:,0,0].cpu().long().numpy()[0]\n",
+    "    all_indexes_test.append(img_idx)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 37,
+   "id": "17a4377a-67dc-49fc-b89c-7d641ef218ce",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "def extract_features(config, device='cuda'):\n",
+    "    '''Extract features based on the configuration.'''\n",
+    "    print(\"Extracting features using network:\", config['network'])\n",
+    "    \n",
+    "    f = h5py.File(f'{config[\"image path\"]}/coco_images_224_float16.hdf5', 'r')\n",
+    "    images = f['images']\n",
+    "    print(images.shape)\n",
+    "    model, preprocess = load_model(config['network'])\n",
+    "    layers = config[\"features\"]\n",
+    "    feature_extractor = FeatureExtractor(model, layers, device=device, detach=True)\n",
+    "    \n",
+    "    output_dir = os.path.join(config[\"output base dir\"], \"pytorch\", config['network'])\n",
+    "\n",
+    "    all_indexes_to_compute = set(all_indexes_train + list(set(all_indexes_test)))\n",
+    "    \n",
+    "    # First, ensure all layer directories exist\n",
+    "    # for layer in layers:\n",
+    "    #     layer_output_dir = os.path.join(output_dir, layer)\n",
+    "    #     os.makedirs(layer_output_dir, exist_ok=True)\n",
+    "    #     img = torch.Tensor(images[0])\n",
+    "\n",
+    "    #     x = preprocess(img).unsqueeze(0).to(device)\n",
+    "    #     # Extract features\n",
+    "    #     features = feature_extractor.run(x)\n",
+    "        \n",
+    "    #     # Get the feature for the current layer\n",
+    "    #     f = features.get(layer)\n",
+    "    #     print(f.shape)\n",
+    "        \n",
+    "\n",
+    "    \n",
+    "    for layer in tqdm(layers, desc=\"Processing layers\"):  # Iterate over each layer first\n",
+    "        layer_output_dir = os.path.join(output_dir, layer)\n",
+    "\n",
+    "        with h5py.File(f'{layer_output_dir}/{subj}.h5', 'a') as hdf5_file:\n",
+    "\n",
+    "                # initial_shape = (0, 100)  # Start with 0 rows and 100 columns (or your array shape)\n",
+    "                # max_shape = (None, 100)   # Allow unlimited rows\n",
+    "                \n",
+    "                # if 'dataset' not in hdf5_file:\n",
+    "                #     dataset = hdf5_file.create_dataset(\n",
+    "                #         'dataset', \n",
+    "                #         shape=initial_shape, \n",
+    "                #         maxshape=max_shape, \n",
+    "                #         dtype='float64'\n",
+    "                #     )\n",
+    "                # else:\n",
+    "                #     dataset = hdf5_file['dataset']\n",
+    "            for i, image_index in tqdm(enumerate(all_indexes_to_compute), total = len(all_indexes_to_compute)):\n",
+    "                output_file = os.path.join(layer_output_dir, f\"{image_index}.npy\")\n",
+    "                \n",
+    "                if os.path.exists(output_file):\n",
+    "                    continue  # Skip if the feature for this image and layer already exists\n",
+    "                \n",
+    "                # Load and preprocess the image\n",
+    "                img = torch.Tensor(images[image_index])\n",
+    "                x = preprocess(img).unsqueeze(0).to(device)\n",
+    "                \n",
+    "                # Extract features\n",
+    "                features = feature_extractor.run(x)\n",
+    "                \n",
+    "                # Get the feature for the current layer\n",
+    "                f = features.get(layer)\n",
+    "                if f is None:\n",
+    "                    print(f\"Warning: Layer '{layer}' not found in the extracted features.\")\n",
+    "                    continue\n",
+    "\n",
+    "                if i == 0:\n",
+    "                    initial_shape = f.shape\n",
+    "                    max_shape = (None,) + initial_shape[1:]\n",
+    "\n",
+    "                    if 'dataset' not in hdf5_file:\n",
+    "                        dataset = hdf5_file.create_dataset(\n",
+    "                            'dataset', \n",
+    "                            shape=initial_shape, \n",
+    "                            maxshape=max_shape, \n",
+    "                            dtype='float16'\n",
+    "                        )\n",
+    "                    else:\n",
+    "                        dataset = hdf5_file['dataset']\n",
+    "\n",
+    "                dataset.resize(dataset.shape[0] + f.shape[0], axis=0)\n",
+    "        \n",
+    "                # Append the data\n",
+    "                dataset[-f.shape[0]:] = f.astype(np.float16)\n",
+    "                \n",
+    "                # # Save the feature\n",
+    "                # np.save(output_file, f.astype(np.float16))\n",
+    "        \n",
+    "        # Optional: Break early for testing\n",
+    "        # if layer == some_condition:\n",
+    "        #     break\n",
+    "\n",
+    "    print('All done')"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 38,
+   "id": "5dc66f30-9406-4e7a-adca-81161df3611a",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Extracting features using network: VGG19_ILSVRC_19_layers\n",
+      "(73000, 3, 224, 224)\n",
+      "(1, 64, 224, 224)\n",
+      "(1, 64, 224, 224)\n",
+      "(1, 128, 112, 112)\n",
+      "(1, 128, 112, 112)\n",
+      "(1, 256, 56, 56)\n",
+      "(1, 256, 56, 56)\n",
+      "(1, 256, 56, 56)\n",
+      "(1, 256, 56, 56)\n",
+      "(1, 512, 28, 28)\n",
+      "(1, 512, 28, 28)\n",
+      "(1, 512, 28, 28)\n",
+      "(1, 512, 28, 28)\n",
+      "(1, 512, 14, 14)\n",
+      "(1, 512, 14, 14)\n",
+      "(1, 512, 14, 14)\n",
+      "(1, 512, 14, 14)\n",
+      "(1, 4096)\n",
+      "(1, 4096)\n",
+      "(1, 1000)\n"
+     ]
+    },
+    {
+     "ename": "NameError",
+     "evalue": "name 'ththth' is not defined",
+     "output_type": "error",
+     "traceback": [
+      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
+      "\u001b[0;31mNameError\u001b[0m                                 Traceback (most recent call last)",
+      "Cell \u001b[0;32mIn[38], line 8\u001b[0m\n\u001b[1;32m      5\u001b[0m \u001b[38;5;66;03m# %%\u001b[39;00m\n\u001b[1;32m      6\u001b[0m \u001b[38;5;66;03m# Load config and extract features\u001b[39;00m\n\u001b[1;32m      7\u001b[0m config \u001b[38;5;241m=\u001b[39m load_config(config_path)\n\u001b[0;32m----> 8\u001b[0m \u001b[43mextract_features\u001b[49m\u001b[43m(\u001b[49m\u001b[43mconfig\u001b[49m\u001b[43m)\u001b[49m\n",
+      "Cell \u001b[0;32mIn[37], line 31\u001b[0m, in \u001b[0;36mextract_features\u001b[0;34m(config, device)\u001b[0m\n\u001b[1;32m     27\u001b[0m     f \u001b[38;5;241m=\u001b[39m features\u001b[38;5;241m.\u001b[39mget(layer)\n\u001b[1;32m     28\u001b[0m     \u001b[38;5;28mprint\u001b[39m(f\u001b[38;5;241m.\u001b[39mshape)\n\u001b[0;32m---> 31\u001b[0m \u001b[43mththth\u001b[49m\n\u001b[1;32m     32\u001b[0m \u001b[38;5;28;01mfor\u001b[39;00m layer \u001b[38;5;129;01min\u001b[39;00m tqdm(layers, desc\u001b[38;5;241m=\u001b[39m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mProcessing layers\u001b[39m\u001b[38;5;124m\"\u001b[39m):  \u001b[38;5;66;03m# Iterate over each layer first\u001b[39;00m\n\u001b[1;32m     33\u001b[0m     layer_output_dir \u001b[38;5;241m=\u001b[39m os\u001b[38;5;241m.\u001b[39mpath\u001b[38;5;241m.\u001b[39mjoin(output_dir, layer)\n",
+      "\u001b[0;31mNameError\u001b[0m: name 'ththth' is not defined"
+     ]
+    }
+   ],
+   "source": [
+    "# %% [Block 6]\n",
+    "# Instead of using argparse, manually provide the config file path for a notebook\n",
+    "config_path = '/weka/proj-fmri/ckadirt/spurious_reconstruction/analysis/0_preprocessing/config/NSD/NSD.yaml'  # Replace with your config file path\n",
+    "\n",
+    "# %%\n",
+    "# Load config and extract features\n",
+    "config = load_config(config_path)\n",
+    "extract_features(config)\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "35e4f307-1652-4af7-ad82-3ee4c3f142fd",
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 8,
+   "id": "2d6c4809-0aab-4257-8994-4e892e23de9f",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# with h5py.File(f'/weka/proj-fmri/ckadirt/spurious_reconstruction/analysis/0_preprocessing/features/pytorch/vgg19_torchvision/features[0]/1.h5', 'a') as hdf5_file:\n",
+    "#     dataset = hdf5_file['dataset']\n",
+    "#     print(len(dataset))\n",
+    "#     print(dataset[10].shape, dataset[10].dtype)\n",
+    "#     print(dataset.shape)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 9,
+   "id": "fcf6d112-e24c-4fd5-ac59-f21e3611ebbc",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# dataset"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "mindeye",
+   "language": "python",
+   "name": "mindeye"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.11.10"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}

spurious_reconstruction/analysis/0_preprocessing/.ipynb_checkpoints/iCNN_image_vgg19_feature_extraction-checkpoint.py

@@ -0,0 +1,130 @@
+'''DNN Feature extraction script'''
+
+# %%%
+from glob import glob
+from tqdm import tqdm
+import os
+import argparse
+from PIL import Image
+import numpy as np
+import torch
+from scipy.io import savemat
+from bdpy.dl.torch import FeatureExtractor
+from torchvision import models, transforms
+import yaml
+# %%%
+# Custum transform class to convert RGB to BGR 
+class ConvertRGBtoBGR:
+    def __call__(self, image):
+        return image[[2, 1, 0], :, :]  
+    
+# Custum transform class to convert Image format to tensor while keeping pixel ranges
+class ToTensorWithoutScaling:
+    def __call__(self, image):
+        image = np.array(image).astype(np.float32)  # PIL ImageをNumPy配列に変換
+        tensor = torch.from_numpy(image).permute(2, 0, 1).float()  # NumPy配列をテンソルに変換し、次元を並べ替える
+        return tensor
+    
+class Convert32to16:
+    def __call__(self, image):
+        return image.type(torch.float16)
+
+# %%%
+def load_model(network):
+    
+    if network == "VGG19_ILSVRC_19_layers":
+        from bdpy.dl.torch.models import layer_map, model_factory
+        model = model_factory('vgg19')
+        encoder_param_file = './data/models_shared/pytorch/VGG_ILSVRC_19_layers/VGG_ILSVRC_19_layers.pt'
+        model.load_state_dict(torch.load(encoder_param_file))
+        model.eval()
+        layer_mapping = layer_map('vgg19')
+        # change RGB-> BGR and subtract mean via lambda function
+        mean_image = [104., 117., 123.] #BGR!
+        #mean_image = [103.939, 116.779, 123.68] #BGR!
+        preprocess = transforms.Compose([
+                    transforms.Resize(224, interpolation=Image.BICUBIC),
+                    #transforms.CenterCrop(224),
+                    #transforms.ToTensor(),
+                    ToTensorWithoutScaling(), # カスタム変換を追加
+                    ConvertRGBtoBGR(),  # カスタム変換を追加
+                    transforms.Normalize(mean=mean_image, std=[1.,1.,1.]),
+                ])
+        
+    elif network == "vgg19_torchvision":
+        model = models.vgg19(pretrained=True)
+        model.eval()
+        
+        # 画像の前処理
+        preprocess = transforms.Compose([
+            transforms.Resize(224),#transforms.Resize(256),
+            #transforms.CenterCrop(224),
+            transforms.ToTensor(),
+            transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
+        ]) 
+    else:
+        raise ValueError("Network not supported. Please choose from: VGG19_ILSVRC_19_layers"
+                        +"or define by yourself. The output should be model and preprocessing functions from RGB image to model input."
+                         )
+    
+    return model, preprocess
+
+def load_config(config_path):
+    '''Load configuration file.'''
+    with open(config_path, 'r') as file:
+        config = yaml.safe_load(file)
+    return config
+
+def extract_features(config, device='cuda'):
+    '''Extract features based on the configuration.'''
+    print("Extracting features using network:", config['network'])
+    
+    image_dir = config['image path']
+    image_ext = config["image_ext"]
+    image_files = glob(os.path.join(image_dir, '*.' + image_ext))
+    
+    model, preprocess = load_model(config['network'])
+    layers = config["features"]
+    feature_extractor = FeatureExtractor(model, layers, device=device, detach=True)
+    
+    output_dir = os.path.join(config["output base dir"], "pytorch",config['network'])
+    
+    for image_file in tqdm(image_files):
+        img = Image.open(image_file).convert('RGB')
+        
+        x = preprocess(img).unsqueeze(0).to(device)
+        
+        # Extract features
+        features = feature_extractor.run(x)
+
+        # Save features
+        for layer in features.keys():
+            f = features[layer]
+
+            output_file = os.path.join(
+                output_dir,
+                layer,
+                os.path.splitext(os.path.basename(image_file))[0] + '.mat'
+            )
+
+            if os.path.exists(output_file):
+                continue
+
+            os.makedirs(os.path.join(output_dir, layer), exist_ok=True)
+
+            savemat(output_file, {'feat': f})
+        #print('Saved {}'.format(output_file))
+
+    print('All done')
+
+        
+    
+# %%%
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser(description='DNN Feature extraction')
+    parser.add_argument('config', help='Configuration file')
+    args = parser.parse_args()
+
+    config = load_config(args.config)
+    extract_features(config)
+# %%

spurious_reconstruction/analysis/0_preprocessing/BD_extract_CLIP_vision_features.py

@@ -0,0 +1,115 @@
+## Code is based on https://github.com/ozcelikfu/brain-diffuser/blob/main/scripts/clipvision_extract_features.py
+
+# %%
+import sys
+# %%
+sys.path.append('brain-diffuser/versatile_diffusion')
+import os
+import PIL
+from PIL import Image
+import numpy as np
+
+import torch
+from lib.cfg_helper import model_cfg_bank
+from lib.model_zoo import get_model
+#from lib.experiments.sd_default import color_adjust, auto_merge_imlist
+from torch.utils.data import DataLoader, Dataset
+
+from lib.model_zoo.vd import VD
+from lib.cfg_holder import cfg_unique_holder as cfguh
+from lib.cfg_helper import get_command_line_args, cfg_initiates, load_cfg_yaml
+import torchvision.transforms as T
+import yaml
+from easydict import EasyDict as edict
+import argparse
+import os.path as osp
+
+from scipy.io import savemat
+# %%
+# %%
+class batch_generator_external_image_files(Dataset):
+
+    def __init__(self, data_path):
+        self.data_path = data_path
+        self.im = sorted(os.listdir(self.data_path))
+
+
+    def __getitem__(self,idx):
+        img = Image.open(os.path.join(self.data_path,self.im[idx])).convert('RGB')
+        img = T.functional.resize(img,(512,512))
+        img = torch.tensor(np.array(img)).float()
+        #img = img/255
+        img = img*2 - 1
+        return img.permute(2,0, 1), self.im[idx]
+
+    def __len__(self):
+        return  len(self.im)
+
+def load_config(config_path):
+    '''Load configuration file.'''
+    with open(config_path, 'r') as file:
+        config = yaml.safe_load(file)
+    return config
+
+def extract_CLIP_vision_features(config):
+    batch_size=1
+    # %%
+    cfgm_name = 'vd_noema'
+    cfgm = model_cfg_bank()(cfgm_name)
+    net = get_model()(cfgm)
+    # %%
+    pth = 'versatile_diffusion/pretrained/vd-four-flow-v1-0-fp16-deprecated.pth'
+    sd = torch.load(pth, map_location='cpu')
+    net.load_state_dict(sd, strict=False)    
+    # %%
+    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+    net.clip = net.clip.to(device)
+    # %%
+
+    data_path = config['image path']
+    # %%
+    image_data = batch_generator_external_image_files(data_path)
+    data_loader = DataLoader(image_data,batch_size,shuffle=False)
+    # %%
+    num_embed, num_features, num_image = 257, 768, image_data.__len__()
+    
+    clip_feat = np.zeros((num_image,num_embed,num_features))
+    # %%
+    output_dir = os.path.join(config["output base dir"], "pytorch", 'brain_diffuser_versatile_diffusion')
+    # %%
+    layer = 'vision_encoder'
+    # %%
+    with torch.no_grad():
+        # %%
+        for i,(imgs_torch_tensor, image_files) in enumerate(data_loader):
+            
+            # %%
+            c_v_feature = net.clip_encode_vision(imgs_torch_tensor)
+            
+            for j, image_file in enumerate(image_files):
+                f = c_v_feature[j][np.newaxis].cpu().numpy()
+                
+                output_file = os.path.join(
+                    output_dir,
+                    layer,
+                    os.path.splitext(os.path.basename(image_file))[0] + '.mat'
+                )
+                
+                if os.path.exists(output_file):
+                    continue
+                
+                os.makedirs(os.path.join(output_dir, layer), exist_ok=True)
+            
+                savemat(output_file, {'feat': f})
+                
+        
+# %%
+
+# %%%
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser(description='DNN Feature extraction')
+    parser.add_argument('config', help='Configuration file')
+    args = parser.parse_args()
+
+    config = load_config(args.config)
+    extract_CLIP_vision_features(config)

spurious_reconstruction/analysis/0_preprocessing/BD_extract_Deeprecon_CLIP_text_features.py

@@ -0,0 +1,101 @@
+# %%
+import sys
+sys.path.append('versatile_diffusion')
+import os
+import numpy as np
+
+import torch
+from lib.cfg_helper import model_cfg_bank
+from lib.model_zoo import get_model
+from torch.utils.data import DataLoader, Dataset
+
+from lib.model_zoo.vd import VD
+from lib.cfg_holder import cfg_unique_holder as cfguh
+from lib.cfg_helper import get_command_line_args, cfg_initiates, load_cfg_yaml
+import matplotlib.pyplot as plt
+import torchvision.transforms as T
+from scipy.io import savemat
+from glob import glob
+import pandas as pd
+# %%
+
+#def extract_NSD_CLIP_text_features():
+def extract_Deeprecon_CLIP_text_features():
+    cfgm_name = 'vd_noema'
+    pth = 'versatile_diffusion/pretrained/vd-four-flow-v1-0-fp16-deprecated.pth'
+    cfgm = model_cfg_bank()(cfgm_name)
+    net = get_model()(cfgm)
+    sd = torch.load(pth, map_location='cpu')
+    net.load_state_dict(sd, strict=False)   
+    # %%
+    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+    net.clip = net.clip.to(device)
+
+    # %%
+    num_embed, num_features = 77, 768
+    layer = 'text_encoder'
+    # Deeprecon test setting
+    image_list, df, save_base_path = Deeprecon_test_setting()
+    
+    output_dir = os.path.join(save_base_path, "pytorch", 'brain_diffuser_versatile_diffusion')
+    print('Extract text features from Deeprecon test images')
+    extract_CLIP_text_features(image_list, df, output_dir, layer, net)
+    print('Done')
+    # %%
+    
+    # Deeprecon traininig setting
+    training_image_list, df, save_base_path_training = Deeprecon_training_setting()
+    output_dir_training = os.path.join(save_base_path_training, "pytorch", 'brain_diffuser_versatile_diffusion')
+    print('Extract text features from Deeprecon training images')
+    extract_CLIP_text_features(training_image_list, df, output_dir_training, layer, net)
+    print('Done')
+    
+        
+def Deeprecon_test_setting():
+    original_DeepreconTest_dir = './data/ImageNetTest'
+    image_list = glob(os.path.join(original_DeepreconTest_dir, 'source', '*.JPEG'))
+    # we didn't share the test annotation file due to the possibility to traing the future AI models.
+    df = pd.read_csv(os.path.join(original_DeepreconTest_dir, 'derivatives/captions/amt_20181204/amt_20181204.csv'))
+    save_base_path = "./data/ImageNetTest/derivatives/features"
+    return image_list, df, save_base_path
+    
+def Deeprecon_training_setting():
+    original_DeepreconTraining_dir = './data/contents_shared/ImageNetTraining'
+    image_list = glob(os.path.join(original_DeepreconTraining_dir, 'source', '*.JPEG'))
+    df = pd.read_csv(os.path.join(original_DeepreconTraining_dir, 'derivatives/captions/amt_20181204/amt_20181204.csv'))
+    save_base_path = "./data/ImageNetTraining/derivatives/features"
+    return image_list, df, save_base_path
+    
+def extract_CLIP_text_features(image_list, df, output_dir, layer, net):
+    for image in image_list:
+        file_name = os.path.basename(image).split('.')[0]
+        caption_texts = df[df['content_id'] == file_name]['caption'].values
+        #print(img)
+        out_text_list  =[]
+        # %%
+        for cap in caption_texts:
+            c_t_feature = net.clip_encode_text(cap)
+            out_text_list.append(c_t_feature.cpu().detach().numpy())
+        out_text_list = np.array(out_text_list).squeeze() #[5, 77, 768]
+        
+        # %%
+        output_file = os.path.join(
+            output_dir,
+            layer,
+            file_name + '.mat'
+        )
+        if os.path.exists(output_file):
+            continue
+        
+        os.makedirs(os.path.join(output_dir, layer), exist_ok=True)
+
+        savemat(output_file, dict([('feat', np.mean(out_text_list, 0, keepdims=True)),
+            ('raw_feat', out_text_list)]) )
+    
+            
+# %%
+
+# %%%
+if __name__ == '__main__':
+
+    extract_Deeprecon_CLIP_text_features()

spurious_reconstruction/analysis/0_preprocessing/BD_extract_NSD_CLIP_text_features.py

@@ -0,0 +1,78 @@
+# %%
+import sys
+sys.path.append('versatile_diffusion')
+import os
+import numpy as np
+
+import torch
+from lib.cfg_helper import model_cfg_bank
+from lib.model_zoo import get_model
+from torch.utils.data import DataLoader, Dataset
+
+from lib.model_zoo.vd import VD
+from lib.cfg_holder import cfg_unique_holder as cfguh
+from lib.cfg_helper import get_command_line_args, cfg_initiates, load_cfg_yaml
+import matplotlib.pyplot as plt
+import torchvision.transforms as T
+from nsd_access.nsd_access.nsda import NSDAccess
+from scipy.io import savemat
+# %%
+
+def extract_NSD_CLIP_text_features():
+
+    cfgm_name = 'vd_noema'
+    pth = 'versatile_diffusion/pretrained/vd-four-flow-v1-0-fp16-deprecated.pth'
+    cfgm = model_cfg_bank()(cfgm_name)
+    net = get_model()(cfgm)
+    sd = torch.load(pth, map_location='cpu')
+    net.load_state_dict(sd, strict=False)   
+    # %%
+    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+    net.clip = net.clip.to(device)
+
+    # %%
+    num_embed, num_features = 77, 768
+
+    origina_nsd_dir = '[PATH_TO_NSD_DATASET_DIR]' #'/home/nu/data/fmri_shared/public/NSD/latest/'
+    nsda = NSDAccess(origina_nsd_dir)
+
+    output_dir = os.path.join("./data/NSD-stimuli/derivatives/features/", "pytorch", 'brain_diffuser_versatile_diffusion')
+    layer = 'text_encoder'
+    # %%
+    total_num = 73000
+    for img in range(total_num):
+        file_name = f'nsd{img+1:05}'+ '.mat'
+        ci = nsda.read_image_coco_info([img], info_type='captions', show_annot=True)
+        #print(img)
+        out_text_list  =[]
+        # %%
+
+        for annotator in ci:
+            cap = annotator['caption']
+            
+            c_v_feature = net.clip_encode_text(cap)
+            out_text_list.append(c_v_feature.cpu().detach().numpy())
+        out_text_list = np.array(out_text_list).squeeze() #[5, 77, 768]
+        
+        # %%
+        output_file = os.path.join(
+            output_dir,
+            layer,
+            file_name
+        )
+        if os.path.exists(output_file):
+            continue
+        
+            
+        os.makedirs(os.path.join(output_dir, layer), exist_ok=True)
+
+        savemat(output_file, dict([('feat', np.mean(out_text_list, 0, keepdims=True)),
+            ('raw_feat', out_text_list)]) )
+        
+            
+# %%
+
+# %%%
+if __name__ == '__main__':
+
+    extract_NSD_CLIP_text_features()

spurious_reconstruction/analysis/0_preprocessing/README.md

@@ -0,0 +1,15 @@
+Scripts for extracting deep neural networks (DNNs) features from images (or text annotation), located at `./data/[dataset]/source` folders. You should put the stimuli files appropriately beforehand or run the download script.
+
+# Usage (using rye environment)
+If you want to extract VGG19 (for iCNN) in the Deeprecon dataset, running this command:
+```rye run python ./analysis/0_preprocessing/iCNN_image_vgg19_feature_extraction.py  ./analysis/0_preprocessing/config/Deeprecon/ImageNetTest.yaml``` 
+
+If you want to extract CLIP text (for Brain-Diffuser) in the Deeprecon dataset, running this command:
+```rye run python ./analysis/0_preprocessing/BD_extract_NSD_CLIP_tex_feature.py  ./analysis/0_preprocessing/config/NSD/NSD.yaml``` 
+
+If you want to extract CLIP vision (for Brain-Diffuser) in the NSD dataset, running this command:
+```rye run python ./analysis/0_preprocessing/BD_extract_CLIP_vision_feature.py  ./analysis/0_preprocessing/config/NSD/NSD.yaml``` 
+
+Note:
+If you want to extract CLIP features used in the Brain-Diffuser models, you need additonaly prepare[versatile_diffusion](https://github.com/ozcelikfu/brain-diffuser/tree/main/versatile_diffusion) at the top of this directory. 
+

spurious_reconstruction/analysis/0_preprocessing/config/Deeprecon/ArtificialShapes.yaml

@@ -0,0 +1,34 @@
+image path:
+  ./data/ArtificialShapes/source
+
+image_ext:
+  tiff
+
+output base dir:
+  ./data/ArtificialShapes/derivatives/features/
+
+
+# This is only related to the script "iCNN_image_vgg19_feature_extraction.py"
+network:
+  VGG19_ILSVRC_19_layers
+
+features:
+  - features[0]   # conv1_1
+  - features[2]   # conv1_2
+  - features[5]   # conv2_1
+  - features[7]   # conv2_2
+  - features[10]  # conv3_1
+  - features[12]  # conv3_2
+  - features[14]  # conv3_3
+  - features[16]  # conv3_4
+  - features[19]  # conv4_1
+  - features[21]  # conv4_2
+  - features[23]  # conv4_3
+  - features[25]  # conv4_4
+  - features[28]  # conv5_1
+  - features[30]  # conv5_2
+  - features[32]  # conv5_3
+  - features[34]  # conv5_4
+  - classifier[0] # fc6
+  - classifier[3] # fc7
+  - classifier[6] # fc8