Official Release of NeurIPS 2023 Spotlight paper "LSD: Object-Centric Slot Diffusion"
import os
import argparse
import warnings

def parse_args(input_args=None):
    parser = argparse.ArgumentParser(description="Simple example of a training script.")
    parser.add_argument(
        "--pretrained_model_name",
        type=str,
        default="stabilityai/stable-diffusion-2-1",
        help="Path to pretrained model or model identifier from huggingface.co/models.",
    )
    parser.add_argument(
        "--revision",
        type=str,
        default=None,
        required=False,
        help=(
            "Revision of pretrained model identifier from huggingface.co/models. Trainable model components should be"
            " float32 precision."
        ),
    )

    parser.add_argument(
        "--output_dir",
        type=str,
        default="log",
        help="The output directory where the model predictions and checkpoints will be written.",
    )
    parser.add_argument("--seed", type=int, default=None, help="A seed for reproducible training.")
    parser.add_argument(
        "--resolution",
        type=int,
        default=512,
        help=(
            "The resolution for input images, all the images in the train/validation dataset will be resized to this"
            " resolution"
        ),
    )
    
    parser.add_argument(
        "--train_batch_size", type=int, default=4, 
        help="Batch size (per device) for the training dataloader."
    )
    parser.add_argument(
        "--val_batch_size", type=int, default=4, 
        help="Batch size (per device) for the validation dataloader."
    )
    parser.add_argument("--num_train_epochs", type=int, default=1)
    parser.add_argument(
        "--max_train_steps",
        type=int,
        default=None,
        help="Total number of training steps to perform.  If provided, overrides num_train_epochs.",
    )
    parser.add_argument(
        "--checkpointing_steps",
        type=int,
        default=500,
        help=(
            "Save a checkpoint of the training state every X updates. Checkpoints can be used for resuming training via `--resume_from_checkpoint`. "
        ),
    )
    parser.add_argument(
        "--checkpoints_total_limit",
        type=int,
        default=None,
        help=(
            "Max number of checkpoints to store."
        ),
    ) # for reference: for movi-e experiment, each pack of ckpt takes roughly 3.9GB. model weights are only < 1.4g though, just the optimizer state is huge 2.6g
    parser.add_argument(
        "--resume_from_checkpoint",
        type=str,
        default=None,
        help=(
            "Whether training should be resumed from a previous checkpoint. Use a path saved by"
            ' `--checkpointing_steps`, or `"latest"` to automatically select the last available checkpoint.'
        ),
    )
    parser.add_argument(
        "--gradient_accumulation_steps",
        type=int,
        default=1,
        help="Number of updates steps to accumulate before performing a backward/update pass.",
    )
    parser.add_argument(
        "--gradient_checkpointing",
        action="store_true",
        help="Whether or not to use gradient checkpointing to save memory at the expense of slower backward pass.",
    )
    parser.add_argument(
        "--learning_rate",
        type=float,
        default=5e-6,
        help="Initial learning rate (after the potential warmup period) to use.",
    )
    parser.add_argument(
        "--scale_lr",
        action="store_true",
        default=False,
        help="Scale the learning rate by the number of GPUs, gradient accumulation steps, and batch size.",
    )
    
    parser.add_argument(
        "--use_8bit_adam", action="store_true", help="Whether or not to use 8-bit Adam from bitsandbytes."
    )
    parser.add_argument(
        "--dataloader_num_workers",
        type=int,
        default=0,
        help=(
            "Number of subprocesses to use for data loading. 0 means that the data will be loaded in the main process."
        ),
    )
    parser.add_argument("--adam_beta1", type=float, default=0.9, help="The beta1 parameter for the Adam optimizer.")
    parser.add_argument("--adam_beta2", type=float, default=0.999, help="The beta2 parameter for the Adam optimizer.")
    parser.add_argument("--adam_weight_decay", type=float, default=1e-2, help="Weight decay to use.")
    parser.add_argument("--adam_epsilon", type=float, default=1e-08, help="Epsilon value for the Adam optimizer")
    parser.add_argument("--max_grad_norm", default=1.0, type=float, help="Max gradient norm.")
    parser.add_argument("--push_to_hub", action="store_true", help="Whether or not to push the model to the Hub.")
    parser.add_argument("--hub_token", type=str, default=None, help="The token to use to push to the Model Hub.")
    # parser.add_argument(
    #     "--hub_model_id",
    #     type=str,
    #     default=None,
    #     help="The name of the repository to keep in sync with the local `output_dir`.",
    # ) # not supported yet
    parser.add_argument(
        "--logging_dir",
        type=str,
        default="logs",
        help=(
            "[TensorBoard](https://www.tensorflow.org/tensorboard) log directory. Will default to"
            " *output_dir/runs/**CURRENT_DATETIME_HOSTNAME***."
        ),
    )
    parser.add_argument(
        "--allow_tf32",
        action="store_true",
        help=(
            "Whether or not to allow TF32 on Ampere GPUs. Can be used to speed up training. For more information, see"
            " https://pytorch.org/docs/stable/notes/cuda.html#tensorfloat-32-tf32-on-ampere-devices"
        ),
    )
    parser.add_argument(
        "--report_to",
        type=str,
        default="tensorboard",
        help=(
            'The integration to report the results and logs to. Supported platforms are `"tensorboard"`'
            ' (default), `"wandb"` and `"comet_ml"`. Use `"all"` to report to all integrations.'
        ),
    )
    
    parser.add_argument(
        "--num_validation_images",
        type=int,
        default=4,
        help="Number of images that should be generated during validation.",
    )
    parser.add_argument(
        "--validation_steps",
        type=int,
        default=100,
        help=(
            "Run validation every X steps. Validation consists of running the prompt"
            " `args.validation_prompt` multiple times: `args.num_validation_images`"
            " and logging the images."
        ),
    )
    parser.add_argument(
        "--mixed_precision",
        type=str,
        default=None,
        choices=["no", "fp16", "bf16"],
        help=(
            "Whether to use mixed precision. Choose between fp16 and bf16 (bfloat16). Bf16 requires PyTorch >="
            " 1.10.and an Nvidia Ampere GPU.  Default to the value of accelerate config of the current system or the"
            " flag passed with the `accelerate.launch` command. Use this argument to override the accelerate config."
        ),
    ) # recommended 
    
    parser.add_argument("--local_rank", type=int, default=-1, help="For distributed training: local_rank")
    parser.add_argument(
        "--enable_xformers_memory_efficient_attention", action="store_true", help="Whether or not to use xformers."
    )
    parser.add_argument(
        "--set_grads_to_none",
        action="store_true",
        help=(
            "Save more memory by using setting grads to None instead of zero. Be aware, that this changes certain"
            " behaviors, so disable this argument if it causes any problems. More info:"
            " https://pytorch.org/docs/stable/generated/torch.optim.Optimizer.zero_grad.html"
        ),
    )

    parser.add_argument(
        "--offset_noise",
        action="store_true",
        default=False,
        help=(
            "Fine-tuning against a modified noise"
            " See: https://www.crosslabs.org//blog/diffusion-with-offset-noise for more information."
        ),
    )
    parser.add_argument(
        "--snr_gamma",
        type=float,
        default=None,
        help="SNR weighting gamma to be used if rebalancing the loss. Recommended value is 5.0. "
        "More details here: https://arxiv.org/abs/2303.09556.",
    )
    
    parser.add_argument(
        "--validation_scheduler",
        type=str,
        default="DPMSolverMultistepScheduler",
        choices=["DPMSolverMultistepScheduler", "DDPMScheduler"],
        help="Select which scheduler to use for validation. DDPMScheduler is recommended for DeepFloyd IF.",
    )

    # ------------------------------------ Latent slot diffusion ------------------------------------
    parser.add_argument(
        "--tracker_project_name",
        type=str,
        default="latent_slot_diffusion",
        help="The name of the tracker project to use for logging.",
    )
    parser.add_argument(
        "--backbone_config",
        type=str,
        default=None,
        help="Path to a config file for the backbone.",
        required=True,
    )
    parser.add_argument(
        "--slot_attn_config",
        type=str,
        default=None,
        help="Path to a config file for the slot attention.",
        required=True,
    )
    parser.add_argument(
        "--unet_config",
        type=str,
        default=None,
        help="Path to a config file for the unet.",
        required=True,
    )
    parser.add_argument(
        "--scheduler_config",
        type=str,
        default=None,
        help="Path to a config file for the scheduler.",
        required=True,
    )
    parser.add_argument(
        "--dataset_root",
        type=str,
        default=None,
        help="Path to the dataset root.",
        required=True,
    )
    parser.add_argument(
        "--dataset_glob",
        type=str,
        default="**/*.png",
        help="Glob pattern for the dataset.",
    )
    parser.add_argument(
        "--encoder_lr_scale",
        type=float,
        default=1.0,
        help="Scale the learning rate of the encoder by this factor. 1.0 means same learning rate as the ldm model.",
    )

    parser.add_argument(
        "--flip_images",
        action="store_true",
        help="Whether to flip the image horizontally in training.",
    )
    parser.add_argument(
        "--train_split_portion",
        type=float,
        default=0.9,
        help="Portion of the dataset to use for training.",
    )

    if input_args is not None:
        args = parser.parse_args(input_args)
    else:
        args = parser.parse_args()

    env_local_rank = int(os.environ.get("LOCAL_RANK", -1))
    if env_local_rank != -1 and env_local_rank != args.local_rank:
        args.local_rank = env_local_rank

    args.output_dir = os.path.join(args.output_dir, args.tracker_project_name)

    return args

if __name__ == "__main__":
    # define and save scheduler config from here
    import os
    import sys
    if __name__ == "__main__":
        sys.path.append(os.path.join(os.path.dirname(__file__), '../../'))
    from diffusers.schedulers import DDPMScheduler
    scheduler = DDPMScheduler(
        num_train_timesteps=1000,
        beta_start= 0.00085,
        beta_end=0.0120,
        beta_schedule='linear',
    )
    scheduler.save_config("./configs/movi-e/scheduler")
configs/parser.py
JindongJiang/latent-slot-diffusion
# conda create -n tfds && conda activate tfds && pip install tensorflow-datasets gcfs tqdm pillow
import os
from tqdm.auto import tqdm
import json
import argparse
from PIL import Image
import tensorflow as tf
import numpy as np
import tensorflow_datasets as tfds 

parser = argparse.ArgumentParser()
parser.add_argument(
    '--dataset_split', 
    type=str, 
    default='movi-a',
    choices=['movi-a', 'movi-b', 'movi-c', 'movi-d', 'movi-e', 'movi-f'],
    help='Dataset split to use (movi-a, movi-b, movi-c, movi-d, movi-e, movi-f)'
)
parser.add_argument(
    '--data_dir',
    type=str,
    default='/path_to_your_movi/',
    help='Directory to save the dataset'
)

args = parser.parse_args()

ds, ds_info = tfds.load(f"{args.dataset_split.replace('-', '_')}/256x256:1.0.0", 
                        data_dir="gs://kubric-public/tfds", with_info=True)

for section in ["train", "validation", "test"]:
    out_path_images = os.path.join(args.data_dir, f'{args.dataset_split}/{args.dataset_split}-{section}-with-label/images')
    out_path_labels = os.path.join(args.data_dir, f'{args.dataset_split}/{args.dataset_split}-{section}-with-label/labels')

    try:
        dataset = tfds.as_numpy(ds[section])
        data_iter = iter(dataset)
    except:
        continue

    # to_tensor = transforms.ToTensor()

    class JsonEncoder(json.JSONEncoder):
        """ Special json encoder for numpy types """
        def default(self, obj):
            if isinstance(obj, np.integer):
                return int(obj)
            elif isinstance(obj, np.floating):
                return float(obj)
            elif isinstance(obj, np.ndarray):
                return obj.tolist()
            elif isinstance(obj, tf.RaggedTensor):
                return obj.to_tensor().numpy().tolist()
            elif isinstance(obj, bytes):
                return obj.decode()

            return json.JSONEncoder.default(self, obj)

    b = 0
    progress_bar = tqdm(
        range(0, len(dataset)),
        initial=0,
        desc=f"{section} Steps",
        position=0, leave=True
    )
    for record in data_iter:
        video = record['video']
        segments = record['segmentations']
        instances = {}
        for k, v in record['instances'].items():
            if isinstance(v, tf.RaggedTensor):
                v = v.to_tensor().numpy()
            instances[k] = v
        T, *_ = video.shape

        # setup dirs
        path_vid_images = os.path.join(out_path_images, f"{b:08}")
        os.makedirs(path_vid_images, exist_ok=True)
        path_vid_labels = os.path.join(out_path_labels, f"{b:08}")
        os.makedirs(path_vid_labels, exist_ok=True)

        for t in range(T):
            img = video[t]
            seg = segments[t]
            Image.fromarray(img).save(os.path.join(path_vid_images, f"{t:08}_image.png"), optimize=False)
            Image.fromarray(seg[..., 0]).save(os.path.join(path_vid_labels, f"{t:08}_segment.png"), optimize=False)

            instances_t = {}
            for k, v in instances.items():
                if len(v.shape) > 1 and v.shape[1] == T:
                    instances_t[k] = v[:, t]
                else:
                    instances_t[k] = v
            with open(os.path.join(path_vid_labels, f"{t:08}_instances.json"), 'w') as f:
                json.dump(instances_t, f, cls=JsonEncoder)

        b += 1
        progress_bar.update(1)
scripts/data_preprocess/movi_kubric_dump_with_labels.py
JindongJiang/latent-slot-diffusion
import os
import glob
import random
import torchvision
import numpy as np
from PIL import Image
from torchvision import transforms
from torch.utils.data import Dataset

class GlobDataset(Dataset):
    def __init__(self, root, img_size, img_glob='*.png', 
                 data_portion=(),  random_data_on_portion=True,
                 random_flip=False):
        super().__init__()
        if isinstance(root, str) or not hasattr(root, '__iter__'):
            root = [root]
            img_glob = [img_glob]
        if not all(hasattr(sublist, '__iter__') for sublist in data_portion) or data_portion == (): # if not iterable or empty
            data_portion = [data_portion]
        self.root = root
        self.img_size = img_size
        self.episodes = []
        self.random_flip = random_flip

        for n, (r, g) in enumerate(zip(root, img_glob)):
            episodes = glob.glob(os.path.join(r, g), recursive=True)

            episodes = sorted(episodes)

            data_p = data_portion[n]

            assert (len(data_p) == 0 or len(data_p) == 2)
            if len(data_p) == 2:
                assert max(data_p) <= 1.0 and min(data_p) >= 0.0

            if data_p and data_p != (0., 1.):
                if random_data_on_portion:
                    random.Random(42).shuffle(episodes) # fix results
                episodes = \
                    episodes[int(len(episodes)*data_p[0]):int(len(episodes)*data_p[1])]

            self.episodes += episodes
        
        # resize the shortest side to img_size and center crop
        self.transform = transforms.Compose([
            transforms.Resize(img_size, interpolation=torchvision.transforms.InterpolationMode.BILINEAR),
            transforms.CenterCrop(img_size),
            transforms.ToTensor(),
            transforms.Normalize(mean=[0.5], std=[0.5])
        ])
    
    def __len__(self):
        return len(self.episodes)

    def __getitem__(self, i):
        example = {}
        image = Image.open(self.episodes[i])
        if not image.mode == "RGB":
            image = image.convert("RGB")
        if self.random_flip:
            if random.random() > 0.5:
                image = image.transpose(Image.FLIP_LEFT_RIGHT)
        pixel_values = self.transform(image)
        example["pixel_values"] = pixel_values
        return example

if __name__ == "__main__":
    dataset = GlobDataset(
        root="/research/projects/object_centric/shared_datasets/movi/movi-e/movi-e-train-with-label/images/",
        img_size=256,
        img_glob="**/*.png",
        data_portion=(0.0, 0.9)
    )
    pass
src/data/dataset.py
JindongJiang/latent-slot-diffusion
from pathlib import Path
import sys
path_root = Path(__file__).parents[2]
sys.path.append(str(path_root))
import random
import torch
import argparse
import numpy as np
from torch import nn
from tqdm.auto import tqdm
from accelerate import Accelerator
from einops import rearrange, reduce
from sklearn.metrics import r2_score

from packaging import version
import torch.nn.functional as F
from torch.utils.data import DataLoader
from accelerate.logging import get_logger
from diffusers.utils.import_utils import is_xformers_available
from src.eval.eval_utils import GlobVideoDatasetWithLabel, ari, get_mask_cosine_distance, \
    hungarian_algorithm, clevrtex_label_reading, movi_label_reading
from accelerate.utils import set_seed

from src.models.backbone import UNetEncoder
from src.models.slot_attn import MultiHeadSTEVESA

parser = argparse.ArgumentParser()


parser.add_argument(
    "--ckpt_path",
    type=str,
    default=None,
    help="Path to a checkpoint folder for the model.",
)

parser.add_argument(
    "--seed",
    type=int,
    default=666,
)

parser.add_argument(
    "--dataset_root",
    type=str,
    default=None,
    help="Path to the dataset root.", # note that we should use val split for some of kubric dataset such as the movi-e, text split is for ood eval
    required=True,
)
parser.add_argument(
    "--dataset_glob",
    type=str,
    default="**/*.png",
    help="Glob pattern for the dataset.",
)
parser.add_argument(
    "--batch_size",
    type=int,
    default=50,
)

parser.add_argument(
    "--num_workers",
    type=int,
    default=8,
)

parser.add_argument(
    "--data_portion",
    nargs=2, type=float, metavar=('x', 'y'),
    default=()
)

parser.add_argument(
    "--mixed_precision",
    type=str,
    default=None,
    choices=["no", "fp16", "bf16"],
    help=(
        "Whether to use mixed precision. Choose between fp16 and bf16 (bfloat16). Bf16 requires PyTorch >="
        " 1.10.and an Nvidia Ampere GPU.  Default to the value of accelerate config of the current system or the"
        " flag passed with the `accelerate.launch` command. Use this argument to override the accelerate config."
    ),
)

parser.add_argument(
    "--resolution",
    type=int,
    default=256,
    help=(
        "The resolution for input images, all the images in the train/validation dataset will be resized to this"
        " resolution"
    ),
)

parser.add_argument(
    "--enable_xformers_memory_efficient_attention",
    action="store_true",
    help=(
        "Whether to use the memory efficient attention implementation of xFormers. This is an experimental feature"
        " and is only available for PyTorch >= 1.10.0 and xFormers >= 0.0.17."
    ),
)

parser.add_argument(
    "--linear_prob_train_portion",
    type=float,
    default=0.83, 
    help=(
        "The portion of the training data (in the testing slots) to use for attribute prediction tasks."
    ),
)

logger = get_logger(__name__)

args = parser.parse_args()

set_seed(args.seed)
torch.backends.cudnn.deterministic = True

accelerator = Accelerator(mixed_precision=args.mixed_precision)
weight_dtype = torch.float32
if accelerator.mixed_precision == "fp16":
    weight_dtype = torch.float16
elif accelerator.mixed_precision == "bf16":
    weight_dtype = torch.bfloat16

max_num_obj = None
max_num_obj_dict = {'movi-a': 10, 'movi-e': 23, 'movi-c': 10,
                    'clevr_with_masks': 10, 'clevrtex': 10}
if "movi-c-test" in args.dataset_root:
    max_num_obj = 24
else:
    for k, v in max_num_obj_dict.items():
        if k in args.dataset_root:
            max_num_obj = v
            break
assert max_num_obj is not None, 'max_num_obj is not set'
if 'movi-a' in args.dataset_root:
    classification_label_names = ['color_label',
                                  'material_label', 'shape_label', 'size_label']
    regresion_label_names = ['image_positions']
elif 'movi-e' in args.dataset_root or 'movi-c' in args.dataset_root:
    classification_label_names = ['category']
    regresion_label_names = ['image_positions', 'bboxes_3d', 'scale']
elif 'clevr_with_masks' in args.dataset_root:
    classification_label_names = ['color', 'material', 'shape', 'size']
    regresion_label_names = ['x', 'y', 'z', 'pixel_coords', 'rotation']
elif 'clevrtex' in args.dataset_root:
    classification_label_names = ['color', 'material', 'shape', 'size']
    regresion_label_names = ['pixel_coords', 'rotation']
else:
    raise NotImplementedError

label_reading = clevrtex_label_reading if 'clevrtex' in args.dataset_root else movi_label_reading

if args.data_portion:
    data_portion = args.data_portion
else:
    data_portion = ()

test_dataset = GlobVideoDatasetWithLabel(root=args.dataset_root, 
                                         max_num_obj=max_num_obj,
                                         img_size=args.resolution,
                                         img_glob=args.dataset_glob,
                                         data_portion=data_portion,
                                         keys_to_log=classification_label_names + regresion_label_names,
                                         label_reading=label_reading,
                                         )

test_sampler = None

loader_kwargs = {
    'batch_size': args.batch_size,
    'shuffle': False,
    'num_workers': args.num_workers,
    'pin_memory': True,
    'drop_last': False,
}

test_loader = DataLoader(test_dataset, sampler=test_sampler, **loader_kwargs)

backbone = UNetEncoder.from_pretrained(
    args.ckpt_path, subfolder="UNetEncoder".lower())
backbone = backbone.to(device=accelerator.device, dtype=weight_dtype)
# backbone.register_to_config(**backbone.config)
slot_attn = MultiHeadSTEVESA.from_pretrained(
    args.ckpt_path, subfolder="MultiHeadSTEVESA".lower())
slot_attn = slot_attn.to(device=accelerator.device, dtype=weight_dtype)
# backbone.register_to_config(**load_model.config)

if args.enable_xformers_memory_efficient_attention:
    if is_xformers_available():
        import xformers

        xformers_version = version.parse(xformers.__version__)
        if xformers_version == version.parse("0.0.16"):
            logger.warn(
                "xFormers 0.0.16 cannot be used for training in some GPUs. If you observe problems during training, please update xFormers to at least 0.0.17. See https://huggingface.co/docs/diffusers/main/en/optimization/xformers for more details."
            )
        backbone.enable_xformers_memory_efficient_attention()
    else:
        raise ValueError(
            "xformers is not available. Make sure it is installed correctly")

total_data = 0
ari_list_no_bg = []
# ari_list_with_bg = []
miou_list_no_bg = []
mbo_list_no_bg = []

slots_list = []
# shape_label_list = []
label_continuous_dict = {
    k: [] for k in regresion_label_names
}

label_discrete_dict = {
    k: [] for k in classification_label_names
}

progress_bar = tqdm(
        range(0, len(test_loader)),
        initial=0,
        desc="Steps",
        # Only show the progress bar once on each machine.
        disable=not accelerator.is_local_main_process,
        position=0, leave=True
    )

with torch.no_grad():
    for batch_idx, batch in enumerate(test_loader):
        # image, true_mask, labels = batch_data
        pixel_values = batch["pixel_values"].to(device=accelerator.device, dtype=weight_dtype)
        true_masks = batch["masks"].to(device=accelerator.device, dtype=weight_dtype)
        labels = {k: v.to(device=accelerator.device, dtype=weight_dtype) for k, v in batch["labels"].items()}

        feat = backbone(pixel_values)
        slots, attns = slot_attn.forward_slots(feat[:, None])  # for the time dimension
        slots = slots[:, 0]
        attns = attns[:, 0]

        attns = reduce(attns, 'b num_head (h w) num_s -> b num_s h w', h=int(np.sqrt(attns.shape[2])),
                       w=int(np.sqrt(attns.shape[2])), reduction='sum')
        attns = F.interpolate(attns, size=true_masks.shape[-1], mode='bilinear')
        attn_argmax = attns.argmax(dim=1, keepdim=False)
        attns_one_hot = \
            rearrange(F.one_hot(attn_argmax, num_classes=attns.shape[1]).float(
            ), 'b h w num_slots -> b num_slots h w')

        ari_list_no_bg.append(
            ari(true_masks.cpu(), attn_argmax.cpu(), num_ignored_objects=1))
        true_masks_one_hot = F.one_hot(
            true_masks.long(), num_classes=max_num_obj + 1).float()
        true_masks_one_hot = rearrange(
            true_masks_one_hot, 'b h w num_classes -> b num_classes h w')
 
        cost_matrix = get_mask_cosine_distance(
            true_masks_one_hot[..., None, :, :], attns_one_hot[..., None, :, :])  # attns or attns_one_hot

        if labels['visibility'].shape[1] >= max_num_obj + 1:
            selected_objects = labels['visibility']
            selected_objects[:, 0] = 0
            if len(selected_objects.shape) == 2:
                selected_objects = selected_objects[:, :, None]
            obj_idx_adjustment = 0  # multi_object_datasets
        else:
            selected_objects = \
                torch.cat([torch.zeros_like(labels['visibility'][:, 0:1]).to(labels['visibility'].device),
                           labels['visibility'] > 0], dim=1)[..., None]
            obj_idx_adjustment = 1  # movi series or clevrtex

        cost_matrix = cost_matrix * selected_objects + \
            10000 * (1 - selected_objects)
        _, indices = hungarian_algorithm(cost_matrix)

        for idx_in_batch, num_o in enumerate(labels['num_obj']):
            for gt_idx, pred_idx in zip(indices[idx_in_batch][0], indices[idx_in_batch][1]):
                if selected_objects[idx_in_batch, ..., 0][
                        gt_idx] == 0:  # no gt_idx - 1 here because we added the background to the beginning
                    continue

                gt_map = (true_masks[idx_in_batch] == gt_idx)
                pred_map = (attn_argmax[idx_in_batch] == pred_idx)
                miou_list_no_bg.append(
                    (gt_map & pred_map).sum() / (gt_map | pred_map).sum())

                iou_all = (gt_map[None] & attns_one_hot[idx_in_batch].bool()).sum((-2, -1)) / \
                          (gt_map[None] | attns_one_hot[idx_in_batch].bool()).sum(
                              (-2, -1))
                mbo_list_no_bg.append(iou_all.max())

                # slot_obj = slots[idx_in_batch, 0, pred_idx]
                slot_obj = slots[idx_in_batch, pred_idx]
                slots_list.append(slot_obj)
                for k in label_continuous_dict.keys():
                    label_continuous_dict[k].append(
                        labels[k][idx_in_batch, gt_idx - obj_idx_adjustment])
                for k in label_discrete_dict.keys():
                    label_discrete_dict[k].append(
                        (labels[k][idx_in_batch, gt_idx - obj_idx_adjustment]).long())

        total_data += pixel_values.shape[0]
        progress_bar.update(1)


ari_list_no_bg = torch.cat(ari_list_no_bg)

print('ari_list_no_bg: ', ari_list_no_bg.mean())

miou_list_no_bg = torch.mean(torch.stack(miou_list_no_bg))

print('miou_list_no_bg: ', miou_list_no_bg.mean())

mbo_list_no_bg = torch.mean(torch.stack(mbo_list_no_bg))

print('mbo_list_no_bg: ', mbo_list_no_bg.mean())

slots_list = torch.stack(slots_list, dim=0)

# for training use float32 only
weight_dtype = torch.float32
slots_list = slots_list.to(device=accelerator.device, dtype=weight_dtype)

label_continuous_dict = {
    k: torch.stack(v, dim=0).to(device=accelerator.device, dtype=weight_dtype) for k, v in label_continuous_dict.items()
}
# normalize each dimension to aproximately [-1, 1]
# we didn't perform it when getting the paper results
# but realized we should have it for future researches
for k, v in label_continuous_dict.items():
    v_min = v.min(dim=0, keepdim=True)[0]
    v_max = v.max(dim=0, keepdim=True)[0]
    v = (v - v_min) / (v_max - v_min) * 2 - 1
    label_continuous_dict[k] = v
label_continuous_dict = {
    k: v[..., None] if len(v.shape) == 1 else v for k, v in label_continuous_dict.items()
}

label_discrete_dict = {
    k: torch.stack(v, dim=0).to(device=accelerator.device, dtype=torch.long) for k, v in label_discrete_dict.items()
}
label_discrete_dict = {
    k: v[..., 0] if len(v.shape) == 2 else v for k, v in label_discrete_dict.items()
}

slot_continuous_net_dict = {k: nn.Sequential(
    nn.Linear(slots_list.shape[1], slots_list.shape[1]),
    nn.ReLU(),
    nn.Linear(slots_list.shape[1], v.shape[1]),
).to(device=accelerator.device, dtype=weight_dtype) for k, v in label_continuous_dict.items()}
slot_discrete_net_dict = {k: nn.Linear(slots_list.shape[1], int(v.max() + 1)).to(device=accelerator.device, dtype=weight_dtype)
                          for k, v in label_discrete_dict.items()}

# split train and test data
train_portion = args.linear_prob_train_portion
# shuffle data and label in the same way
torch.random.manual_seed(args.seed)
shuffle_idx = torch.randperm(slots_list.shape[0])
for k, v in label_continuous_dict.items():
    label_continuous_dict[k] = label_continuous_dict[k][shuffle_idx]
for k, v in label_discrete_dict.items():
    label_discrete_dict[k] = label_discrete_dict[k][shuffle_idx]
slots_list = slots_list[shuffle_idx]
slots_list_train = slots_list[:int(slots_list.shape[0] * train_portion)]
slots_list_test = slots_list[int(slots_list.shape[0] * train_portion):]
label_continuous_dict_train = {
    k: v[:int(v.shape[0] * train_portion)] for k, v in label_continuous_dict.items()}
label_continuous_dict_test = {
    k: v[int(v.shape[0] * train_portion):] for k, v in label_continuous_dict.items()}
label_discrete_dict_train = {
    k: v[:int(v.shape[0] * train_portion)] for k, v in label_discrete_dict.items()}
label_discrete_dict_test = {
    k: v[int(v.shape[0] * train_portion):] for k, v in label_discrete_dict.items()}


params = []
for k, v in slot_continuous_net_dict.items():
    params = params + list(v.parameters())
for k, v in slot_discrete_net_dict.items():
    params = params + list(v.parameters())

optimizer = torch.optim.AdamW(params)
slot_continuous_loss = dict()
slot_discrete_loss = dict()

print('start property prediction training')

num_training_steps = 4000

progress_bar = tqdm(
    range(0, num_training_steps),
    initial=0,
    desc="Steps",
    # Only show the progress bar once on each machine.
    disable=not accelerator.is_local_main_process,
    position=0, leave=True
)

for epoch in range(num_training_steps):
    optimizer.zero_grad()
    for k, v in slot_continuous_net_dict.items():
        slot_continuous_loss[k] = F.mse_loss(
            v(slots_list_train), label_continuous_dict_train[k])
    for k, v in slot_discrete_net_dict.items():
        slot_discrete_loss[k] = F.cross_entropy(
            v(slots_list_train), label_discrete_dict_train[k])
    loss = sum(slot_continuous_loss.values()) + \
        sum(slot_discrete_loss.values())
    loss = loss / slots_list_train.shape[0]
    loss.backward()
    optimizer.step()
    progress_bar.update(1)

all_loss = dict()
for k, v in slot_continuous_net_dict.items():
    print(f'continuous {k}:', torch.mean(
        F.mse_loss(v(slots_list_test), label_continuous_dict_test[k])))
for k, v in slot_discrete_net_dict.items():
    print(f'discrete {k}:', torch.mean(
        (torch.argmax(v(slots_list_test), dim=1) == label_discrete_dict_test[k]).float()))
src/eval/eval.py
JindongJiang/latent-slot-diffusion
import os
import json
import glob
import torch
import random
import itertools
import numpy as np
from tqdm import tqdm
from PIL import Image, ImageFile
from torchvision import transforms
from torch.utils.data import Dataset
from torch import Tensor, LongTensor
import torch.nn.functional as F
from einops import rearrange, repeat, reduce
from sklearn.metrics import adjusted_rand_score
from scipy.optimize import linear_sum_assignment
from typing import Optional, Tuple, Dict, Any, List

clevrtex_labels_to_idx = {
    'shape': {'cube': 0, 'cylinder': 1, 'monkey': 2, 'sphere': 3},
    'size': {'large': 0, 'medium': 1, 'small': 2},
    'color': {'blue': 0, 'brown': 1, 'cyan': 2, 'gray': 3, 'green': 4, 'purple': 5, 'red': 6, 'yellow': 7},
    'material': {
        'poliigonbricks01': 0,
        'poliigonbricksflemishred001': 1,
        'poliigonbrickspaintedwhite001': 2,
        'poliigoncarpettwistnatural001': 3,
        'poliigonchainmailcopperroundedthin001': 4,
        'poliigoncitystreetasphaltgenericcracked002': 5,
        'poliigoncitystreetroadasphalttwolaneworn001': 6,
        'poliigoncliffjagged004': 7,
        'poliigoncobblestonearches002': 8,
        'poliigonconcretewall001': 9,
        'poliigonfabricdenim003': 10,
        'poliigonfabricfleece001': 11,
        'poliigonfabricleatherbuffalorustic001': 12,
        'poliigonfabricrope001': 13,
        'poliigonfabricupholsterybrightanglepattern001': 14,
        'poliigongroundclay002': 15,
        'poliigongrounddirtforest014': 16,
        'poliigongrounddirtrocky002': 17,
        'poliigongroundforest003': 18,
        'poliigongroundforest008': 19,
        'poliigongroundforestmulch001': 20,
        'poliigongroundforestroots001': 21,
        'poliigongroundmoss001': 22,
        'poliigongroundsnowpitted003': 23,
        'poliigongroundtiretracks001': 24,
        'poliigoninteriordesignrugstarrynight001': 25,
        'poliigonmarble062': 26, 'poliigonmarble13': 27,
        'poliigonmetalcorrodedheavy001': 28,
        'poliigonmetalcorrugatedironsheet002': 29,
        'poliigonmetaldesignerweavesteel002': 30,
        'poliigonmetalpanelrectangular001': 31,
        'poliigonmetalspottydiscoloration001': 32,
        'poliigonmetalstainlesssteelbrushed': 33,
        'poliigonplaster07': 34, 'poliigonplaster17': 35,
        'poliigonroadcityworn001': 36,
        'poliigonrooftilesterracotta004': 37,
        'poliigonrustmixedonpaint012': 38,
        'poliigonrustplain007': 39,
        'poliigonsolarpanelspolycrystallinetypebframedclean001': 40,
        'poliigonstonebricksbeige015': 41,
        'poliigonstonemarblecalacatta004': 42,
        'poliigonterrazzovenetianmattewhite001': 43,
        'poliigontiles05': 44,
        'poliigontilesmarblechevroncreamgrey001': 45,
        'poliigontilesmarblesagegreenbrickbondhoned001': 46,
        'poliigontilesonyxopaloblack001': 47,
        'poliigontilesrectangularmirrorgray001': 48,
        'poliigonwallmedieval003': 49,
        'poliigonwaterdropletsmixedbubbled001': 50,
        'poliigonwoodfinedark004': 51,
        'poliigonwoodflooring044': 52,
        'poliigonwoodflooring061': 53,
        'poliigonwoodflooringmahoganyafricansanded001': 54,
        'poliigonwoodflooringmerbaubrickbondnatural001': 55,
        'poliigonwoodplanks028': 56,
        'poliigonwoodplanksworn33': 57,
        'poliigonwoodquarteredchiffon001': 58,
        'whitemarble': 59
    }
}


def cosine_similarity(a: Tensor, b: Tensor, eps: float = 1e-6):
    """Computes the cosine similarity between two tensors.
    Args:
        a (Tensor): Tensor with shape (batch size, N_a, D).
        b (Tensor): Tensor with shape (batch size, N_b, D).
        eps (float): Small constant for numerical stability.
    Returns:
        The (batched) cosine similarity between `a` and `b`, with shape (batch size, N_a, N_b).
    """
    dot_products = torch.matmul(a, b.transpose(1, 2))
    norm_a = (a * a).sum(dim=2).sqrt().unsqueeze(2)
    norm_b = (b * b).sum(dim=2).sqrt().unsqueeze(1)
    return dot_products / (torch.matmul(norm_a, norm_b) + eps)


def cosine_distance(a: Tensor, b: Tensor, eps: float = 1e-6):
    """Computes the cosine distance between two tensors, as 1 - cosine_similarity.
    Args:
        a (Tensor): Tensor with shape (batch size, N_a, D).
        b (Tensor): Tensor with shape (batch size, N_b, D).
        eps (float): Small constant for numerical stability.
    Returns:
        The (batched) cosine distance between `a` and `b`, with shape (batch size, N_a, N_b).
    """
    return 1 - cosine_similarity(a, b, eps)


def get_mask_cosine_distance(true_mask: Tensor, pred_mask: Tensor):
    """Computes the cosine distance between the true and predicted masks.
    Args:
        true_mask (Tensor): Tensor of shape (batch size, num objects, 1, H, W).
        pred_mask (Tensor): Tensor of shape (batch size, num slots, 1, H, W).
    Returns:
        The (batched) cosine similarity between the true and predicted masks, with
        shape (batch size, num objects, num slots).
    """
    return cosine_distance(true_mask.flatten(2).detach(), pred_mask.flatten(2).detach())


def hungarian_algorithm(cost_matrix: Tensor) -> Tuple[Tensor, LongTensor]:
    """Batch-applies the hungarian algorithm to find a matching that minimizes the overall cost.
    Returns the matching indices as a LongTensor with shape (batch size, 2, min(num objects, num slots)).
    The first column is the row indices (the indices of the true objects) while the second
    column is the column indices (the indices of the slots). The row indices are always
    in ascending order, while the column indices are not necessarily.
    The outputs are on the same device as `cost_matrix` but gradients are detached.
    A small example:
                | 4, 1, 3 |
                | 2, 0, 5 |
                | 3, 2, 2 |
                | 4, 0, 6 |
    would result in selecting elements (1,0), (2,2) and (3,1). Therefore, the row
    indices will be [1,2,3] and the column indices will be [0,2,1].
    Args:
        cost_matrix: Tensor of shape (batch size, num objects, num slots).
    Returns:
        A tuple containing:
            - a Tensor with shape (batch size, min(num objects, num slots)) with the
              costs of the matches.
            - a LongTensor with shape (batch size, 2, min(num objects, num slots))
              containing the indices for the resulting matching.
    """

    # List of tuples of size 2 containing flat arrays
    indices = list(map(linear_sum_assignment, cost_matrix.cpu().detach().numpy()))
    indices = torch.LongTensor(np.array(indices))
    smallest_cost_matrix = torch.stack(
        [
            cost_matrix[i][indices[i, 0], indices[i, 1]]
            for i in range(cost_matrix.shape[0])
        ]
    )
    device = cost_matrix.device
    return smallest_cost_matrix.to(device), indices.to(device)


def ari(
        true_mask, pred_mask, num_ignored_objects: int
) -> torch.FloatTensor:
    """Computes the ARI score.
    Args:
        true_mask: tensor of shape [batch_size x *] where values go from 0 to the number of objects.
        pred_mask:  tensor of shape [batch_size x *] where values go from 0 to the number of objects.
        num_ignored_objects: number of objects (in ground-truth mask) to be ignored when computing ARI.
    Returns:
        a vector of ARI scores, of shape [batch_size, ].
    """
    true_mask = true_mask.flatten(1)
    pred_mask = pred_mask.flatten(1)
    not_bg = true_mask >= num_ignored_objects
    result = []
    batch_size = len(true_mask)
    for i in range(batch_size):
        ari_value = adjusted_rand_score(
            true_mask[i][not_bg[i]], pred_mask[i][not_bg[i]]
        )
        result.append(ari_value)
    result = torch.FloatTensor(result)  # shape (batch_size, )
    return result


def movi_label_reading(image_path, max_num_obj, keys_to_log=None):
    p_mask = str(image_path).replace('/images/', '/labels/').replace('_image.png', '_segment.png')
    p_labels = str(image_path).replace('/images/', '/labels/').replace('_image.png', '_instances.json')
    segment = np.array(Image.open(p_mask).convert('L'))
    with open(p_labels, 'rb') as f:
        label_dict = json.load(f)

    num_obj = len(label_dict['visibility'])
    labels = {'num_obj': num_obj}
    for k, v in label_dict.items():
        if np.array(v).ndim >= 3:
            # flatten the last two dimensions
            v = np.array(v)
            v = v.reshape(v.shape[0], -1).tolist()
        if k == 'visibility':
            continue
        if keys_to_log is None:
            labels[k] = torch.tensor(
                [i for i in v] + [v[-1] for _ in range(max_num_obj - num_obj)]
            )
        else:
            if k in keys_to_log:
                labels[k] = torch.tensor(
                    [i for i in v] + [v[-1] for _ in range(max_num_obj - num_obj)]
                )
    labels['visibility'] = torch.tensor(
        [i for i in label_dict['visibility']] +
        [0 for _ in range(max_num_obj - num_obj)]
    )
    if 'pixel_coords' in labels:
        labels['pixel_coords'] = labels['pixel_coords'][..., :2]
        labels['pixel_coords'][..., 0] -= 160
        labels['pixel_coords'][..., 0] /= 160

        labels['pixel_coords'][..., 1] -= 120
        labels['pixel_coords'][..., 1] /= 120

    return segment, labels

def clevrtex_label_reading(image_path, max_num_obj, keys_to_log=None):
    seg_path = image_path.replace('.png', '_flat.png')
    json_path = image_path.replace('.png', '.json')
    segment = Image.open(seg_path)
    w, h = segment.size
    min_side_len = min(w, h)
    left = (w - min_side_len) // 2
    top = (h - min_side_len) // 2
    right = (w + min_side_len) // 2
    bottom = (h + min_side_len) // 2
    segment = segment.crop((left, top, right, bottom))
    segment = np.array(segment).astype(np.uint8)
    with open(json_path, 'rb') as f:
        label_dict = json.load(f)
    objects_list = label_dict['objects']
    num_obj = len(objects_list)
    labels = {'num_obj': num_obj}
    for k in keys_to_log:
        labels[k] = []
    labels['visibility'] = []

    for obj_dict in objects_list:
        idx = obj_dict['index']
        labels['visibility'].append(np.any(segment == idx))
        for k in keys_to_log:
            v = obj_dict[k]
            if np.array(v).ndim >= 2:
                # flatten the last two dimensions
                v = np.array(v)
                v = v.reshape(v.shape[0], -1).tolist()
            if k in clevrtex_labels_to_idx:
                v = clevrtex_labels_to_idx[k][v]
            labels[k].append(v)
    for k, v in labels.items():
        if isinstance(v, int):
            continue
        if k == 'visibility':
            labels[k] = torch.tensor(
                [i for i in v] + [0 for _ in range(max_num_obj - num_obj)]
            ).float()
        else:
            labels[k] = torch.tensor([i for i in v] + [v[-1] for _ in range(max_num_obj - num_obj)]).float()

    if 'pixel_coords' in labels:
        labels['pixel_coords'] = labels['pixel_coords'][..., :2]
        labels['pixel_coords'][..., 0] -= w // 2
        labels['pixel_coords'][..., 0] /= w // 2

        labels['pixel_coords'][..., 1] -= h // 2
        labels['pixel_coords'][..., 1] /= h // 2

    return segment, labels

class GlobVideoDatasetWithLabel(Dataset):
    def __init__(self, root, img_size, max_num_obj, random_data_on_portion=True, 
                 img_glob='*.png', data_portion=(), keys_to_log=None, 
                 label_reading=movi_label_reading, return_path=False
                 ):
        
        super().__init__()
        if isinstance(root, str) or not hasattr(root, '__iter__'):
            root = [root]
            img_glob = [img_glob]
        if not all(hasattr(sublist, '__iter__') for sublist in data_portion) or data_portion == (): # if not iterable or empty
            data_portion = [data_portion]
        self.root = root
        self.img_size = img_size
        self.episodes = []

        for n, (r, g) in enumerate(zip(root, img_glob)):
            episodes = glob.glob(os.path.join(r, g), recursive=True)

            episodes = sorted(episodes)

            data_p = data_portion[n]

            assert (len(data_p) == 0 or len(data_p) == 2)
            if len(data_p) == 2:
                assert max(data_p) <= 1.0 and min(data_p) >= 0.0

            if data_p and data_p != (0., 1.):
                if random_data_on_portion:
                    random.Random(42).shuffle(episodes) # fix results
                episodes = \
                    episodes[int(len(episodes)*data_p[0]):int(len(episodes)*data_p[1])]

            self.episodes += episodes
        
        # resize the shortest side to img_size and center crop
        self.transform = transforms.Compose([
            transforms.Resize(img_size, interpolation=transforms.InterpolationMode.BILINEAR),
            transforms.CenterCrop(img_size),
            transforms.ToTensor(),
            transforms.Normalize(mean=[0.5], std=[0.5])
        ])

        self.max_num_obj = max_num_obj
        self.keys_to_log = keys_to_log
        self.label_reading = label_reading
        self.return_path = return_path

    def __len__(self):
        return len(self.episodes)

    def __getitem__(self, idx):
        img_loc = self.episodes[idx]
        if self.keys_to_log:
            masks, labels = self.label_reading(img_loc, self.max_num_obj, keys_to_log=self.keys_to_log)
        else:
            masks, labels = 0, 0
        image = Image.open(img_loc).convert("RGB")
        tensor_image = self.transform(image)
        example = {}
        example['pixel_values'] = tensor_image
        example['labels'] = labels
        example['masks'] = masks
        example['img_loc'] = img_loc
        return example
src/eval/eval_utils.py
JindongJiang/latent-slot-diffusion
import os
import sys
if __name__ == "__main__":
    sys.path.append(os.path.join(os.path.dirname(__file__), '../'))
import torch
from torch import nn
from diffusers.models.unet_2d import UNet2DModel, UNet2DOutput
from typing import Any, Dict, List, Optional, Tuple, Union


class UNetEncoder(UNet2DModel):

    def __init__(
        self,
        sample_size: Optional[Union[int, Tuple[int, int]]] = None,
        in_channels: int = 3,
        out_channels: int = 3,
        down_block_types: Tuple[str] = ("DownBlock2D", "AttnDownBlock2D", "AttnDownBlock2D", "AttnDownBlock2D"),
        up_block_types: Tuple[str] = ("AttnUpBlock2D", "AttnUpBlock2D", "AttnUpBlock2D", "UpBlock2D"),
        block_out_channels: Tuple[int] = (224, 448, 672, 896),
        layers_per_block: int = 2,
        attention_head_dim: Optional[int] = 8,
        norm_num_groups: int = 32,
        input_resolution: int = 256,
        input_channels: int = 3,
        **kwargs
    ):
        super().__init__(
            sample_size=sample_size,
            in_channels=in_channels,
            out_channels=out_channels,
            down_block_types=down_block_types,
            up_block_types=up_block_types,
            block_out_channels=block_out_channels,
            layers_per_block=layers_per_block,
            attention_head_dim=attention_head_dim,
            norm_num_groups=norm_num_groups,
            **kwargs
        )

        self.register_to_config(
            input_resolution=input_resolution,
            input_channels=input_channels,
        )
        downscale_stride = input_resolution // sample_size
        self.downscale_cnn = nn.Conv2d(input_channels, in_channels, kernel_size=downscale_stride, stride=downscale_stride)
        self.original_forward = super().forward

    def forward(
        self,
        sample: torch.FloatTensor,
    ) -> Union[UNet2DOutput, Tuple]:

        sample = self.downscale_cnn(sample)
        return self.original_forward(sample, timestep=0, class_labels=None).sample

if __name__ == "__main__":
    block_out_channels = (128, 128, 256, 512) # each of then should be divisible by the least common multiple of (80, 32, 15)

    unet_encoder = UNetEncoder(
        sample_size= 64,
        in_channels = 128, # input channel of the unet conv, not the input image
        out_channels = 192, 
        down_block_types = ("DownBlock2D", "DownBlock2D", "DownBlock2D", "DownBlock2D"),
        up_block_types = ("UpBlock2D", "UpBlock2D", "UpBlock2D", "UpBlock2D"),
        block_out_channels = block_out_channels,
        layers_per_block = 1,
        attention_head_dim = 8,
        norm_num_groups = 16,
        input_resolution=256,
        input_channels=3,
    )

    # save the config only if pass this test
    image_cond = torch.randn(1, 3, 256, 256)
    out = unet_encoder(image_cond)

    unet_encoder.save_config("./configs/movi-e/backbone/")   
    pass
src/models/backbone.py
JindongJiang/latent-slot-diffusion
import os
import sys
if __name__ == "__main__":
    sys.path.append(os.path.join(os.path.dirname(__file__), '../../'))
import math
import torch
from torch import nn
import torch.nn.functional as F
from einops import rearrange, repeat, reduce
from einops.layers.torch import Rearrange
from diffusers.models import ModelMixin
from diffusers.configuration_utils import ConfigMixin, register_to_config

from src.models.utils import CartesianPositionalEmbedding

def is_square(n: float) -> bool:
    if n < 0:
        return False
    sqrt_n = math.sqrt(n)
    return sqrt_n ** 2 == n

class MultiHeadSTEVESA(ModelMixin, ConfigMixin):

    # enable diffusers style config and model save/load
    @register_to_config
    def __init__(self, num_iterations, num_slots, num_heads,
                 input_size, out_size, slot_size, mlp_hidden_size, 
                 input_resolution, epsilon=1e-8, 
                 learnable_slot_init=False, 
                 bi_level=False):
        super().__init__()

        self.pos = CartesianPositionalEmbedding(input_size, input_resolution)
        self.in_layer_norm = nn.LayerNorm(input_size)
        self.in_mlp = nn.Sequential(
            nn.Linear(input_size, input_size),
            nn.ReLU(),
            nn.Linear(input_size, input_size)
            )
        if bi_level:
            assert learnable_slot_init, 'Bi-level training requires learnable_slot_init=True'

        self.num_iterations = num_iterations
        self.num_slots = num_slots
        self.num_heads = num_heads
        self.input_size = input_size
        self.slot_size = slot_size
        self.mlp_hidden_size = mlp_hidden_size
        self.epsilon = epsilon
        self.learnable_slot_init = learnable_slot_init
        self.bi_level = bi_level

        assert slot_size % num_heads == 0, 'slot_size must be divisible by num_heads'

        if learnable_slot_init:
            self.slot_mu = nn.Parameter(torch.Tensor(1, num_slots, slot_size))
            nn.init.xavier_uniform_(self.slot_mu)
        else:
            # parameters for Gaussian initialization (shared by all slots).
            self.slot_mu = nn.Parameter(torch.Tensor(1, 1, slot_size))
            self.slot_log_sigma = nn.Parameter(torch.Tensor(1, 1, slot_size))
            nn.init.xavier_uniform_(self.slot_mu)
            nn.init.xavier_uniform_(self.slot_log_sigma)

        # norms
        self.norm_inputs = nn.LayerNorm(input_size)
        self.norm_slots = nn.LayerNorm(slot_size)
        self.norm_mlp = nn.LayerNorm(slot_size)

        # linear maps for the attention module.
        self.project_q = nn.Linear(slot_size, slot_size, bias=False)
        self.project_k = nn.Linear(input_size, slot_size, bias=False)
        self.project_v = nn.Linear(input_size, slot_size, bias=False)

        # slot update functions.
        self.gru = nn.GRUCell(slot_size, slot_size)
        self.mlp = nn.Sequential(
            nn.Linear(slot_size, mlp_hidden_size),
            nn.ReLU(),
            nn.Linear(mlp_hidden_size, slot_size))
        
        self.out_layer_norm = nn.LayerNorm(slot_size)
        self.out_linear = nn.Linear(slot_size, out_size)
        
    def forward(self, inputs):
        slots_collect, attns_collect = self.forward_slots(inputs)
        slots_collect = self.out_layer_norm(slots_collect)
        slots_collect = self.out_linear(slots_collect)
        return slots_collect, attns_collect

    def forward_slots(self, inputs):
        """
        inputs: batch_size x seq_len x input_size x h x w
        return: batch_size x num_slots x slot_size
        """
        B, T, input_size, h, w = inputs.size()
        inputs = self.pos(inputs)
        inputs = rearrange(inputs, 'b t n_inp h w -> b t (h w) n_inp')
        inputs = self.in_mlp(self.in_layer_norm(inputs))

        # num_inputs = h * w

        if self.learnable_slot_init:
            slots = repeat(self.slot_mu, '1 num_s d -> b num_s d', b=B)
        else:
            # initialize slots
            slots = inputs.new_empty(B, self.num_slots, self.slot_size).normal_()
            slots = self.slot_mu + torch.exp(self.slot_log_sigma) * slots

        # setup key and value
        inputs = self.norm_inputs(inputs)
        k = rearrange(self.project_k(inputs), 'b t n_inp (h d) -> b t h n_inp d',
                      h=self.num_heads)  # Shape: [batch_size, T, num_heads, num_inputs, slot_size].
        v = rearrange(self.project_v(inputs), 'b t n_inp (h d) -> b t h n_inp d',
                      h=self.num_heads)  # Shape: [batch_size, T, num_heads, num_inputs, slot_size].
        k = (self.slot_size ** (-0.5)) * k

        # loop over frames
        attns_collect = []
        slots_collect = []
        for t in range(T):
            # corrector iterations
            for i in range(self.num_iterations):
                if self.bi_level and i == self.num_iterations - 1:
                    slots = slots.detach() + self.slot_mu - self.slot_mu.detach()
                slots_prev = slots
                slots = self.norm_slots(slots)

                # Attention.
                q = rearrange(self.project_q(slots), 'b n_s (h d) -> b h n_s d',
                              h=self.num_heads)  # Shape: [batch_size, num_heads, num_slots, slot_size].
                attn_logits = torch.einsum('...id,...sd->...is', k[:, t],
                                           q)  # Shape: [batch_size, num_heads, num_inputs, num_slots]
                attn = F.softmax(rearrange(attn_logits, 'b h n_inp n_s -> b n_inp (h n_s)'), -1)
                attn_vis = rearrange(attn, 'b n_inp (h n_s) -> b h n_inp n_s', h=self.num_heads)
                # `attn_vis` has shape: [batch_size, num_inputs, num_slots].

                # Weighted mean.
                attn = attn_vis + self.epsilon
                attn = attn / torch.sum(attn, dim=-2, keepdim=True)  # norm over inputs
                updates = torch.einsum('...is,...id->...sd', attn,
                                       v[:, t])  # Shape: [batch_size, num_heads, num_slots, num_inp].
                updates = rearrange(updates, 'b h n_s d -> b n_s (h d)')
                # `updates` has shape: [batch_size, num_slots, slot_size].

                # Slot update.
                slots = self.gru(updates.view(-1, self.slot_size),
                                 slots_prev.reshape(-1, self.slot_size))
                slots = slots.view(-1, self.num_slots, self.slot_size)

                slots = slots + self.mlp(self.norm_mlp(slots))

            # collect
            attns_collect += [attn_vis]
            slots_collect += [slots]

        attns_collect = torch.stack(attns_collect, dim=1)  # B, T, num_inputs, num_slots
        slots_collect = torch.stack(slots_collect, dim=1)  # B, T, num_slots, slot_size

        return slots_collect, attns_collect

if __name__ == "__main__":
    # test
    slot_attn = MultiHeadSTEVESA(
        num_iterations=3, 
        num_slots=24, 
        num_heads=1,
        input_size=192, # unet_encoder.config.out_channels
        out_size=192, # unet.config.cross_attention_dim
        slot_size=192, 
        mlp_hidden_size=192,
        input_resolution=64, # unet_encoder.config.latent_size
        learnable_slot_init=False
    )
    slot_attn.save_config('./configs/movi-e/slot_attn')
    inputs = torch.randn(2, 1, 192, 64, 64)
    slots_collect, attns_collect = slot_attn(inputs)
    print(slots_collect.shape)
    pass
src/models/slot_attn.py
JindongJiang/latent-slot-diffusion
import os
import sys
if __name__ == "__main__":
    sys.path.append(os.path.join(os.path.dirname(__file__), '../../'))
import torch
import numpy as np
from torch import nn
from einops.layers.torch import Rearrange
from einops import rearrange, repeat, reduce
from typing import Any, Dict, List, Optional, Tuple, Union
from diffusers.models.unet_2d_condition import UNet2DConditionModel
from diffusers.models.unet_2d_blocks import CrossAttnDownBlock2D, CrossAttnUpBlock2D, \
    UNetMidBlock2DCrossAttn, UNetMidBlock2DSimpleCrossAttn
from src.models.utils import CartesianPositionalEmbedding, SinCosPositionalEmbedding2D, \
    LearnedPositionalEmbedding1D
    
def get_pos_emb(dim, resolution, pos_type):

    if pos_type == 'cartesian':
        return nn.Sequential(
                    Rearrange('b (h w) c -> b c h w', h=resolution),
                    CartesianPositionalEmbedding(dim, resolution),
                    Rearrange('b c h w -> b (h w) c')
                )
    elif pos_type == 'sincos':
        return nn.Sequential(
                    Rearrange('b (h w) c -> b c h w', h=resolution),
                    SinCosPositionalEmbedding2D(dim, resolution),
                    Rearrange('b c h w -> b (h w) c')
                )
    elif pos_type == 'learned':
        return LearnedPositionalEmbedding1D(resolution**2, dim)
    elif pos_type == 'none':
        return nn.Identity()
    else:
        raise NotImplementedError

def hack_unet_for_pos(unet, pos_type='cartesian'):
    # hack the unet to add positional embedding to the transformer blocks
    blocks_to_add_pos = ["down_blocks", "up_blocks", "mid_block"]
    resolutions = {}
    resolutions["down_blocks"] = [unet.config.sample_size]
    for block in unet.down_blocks:
        if block.downsamplers is not None:
            resolutions["down_blocks"].append(resolutions["down_blocks"][-1] // 2)
        else:
            resolutions["down_blocks"].append(resolutions["down_blocks"][-1])
    resolutions["mid_block"] = [resolutions["down_blocks"][-1]]
    resolutions["up_blocks"] = [resolutions["down_blocks"][-1]]
    resolutions["down_blocks"] = resolutions["down_blocks"][:-1]
    for block in unet.up_blocks:
        if block.upsamplers is not None:
            resolutions["up_blocks"].append(int(resolutions["up_blocks"][-1] * 2))
        else:
            resolutions["up_blocks"].append(resolutions["up_blocks"][-1])
    resolutions["up_blocks"] = resolutions["up_blocks"][:-1]
        
    for block_name in blocks_to_add_pos:
        blocks = getattr(unet, block_name)
        resolution_blocks = resolutions[block_name]
        if block_name == "mid_block":
            blocks = [blocks]
        for block_idx, block in enumerate(blocks):
            if isinstance(block, CrossAttnDownBlock2D) or isinstance(block, CrossAttnUpBlock2D) or \
                isinstance(block, UNetMidBlock2DCrossAttn) or isinstance(block, UNetMidBlock2DSimpleCrossAttn):
                for attention_block in block.attentions:
                    for transformer_block in attention_block.transformer_blocks:
                        transformer_block.attn1.pos_emb = get_pos_emb(
                            dim=transformer_block.attn1.to_q.in_features,
                            resolution=resolution_blocks[block_idx],
                            pos_type=pos_type
                        )
                        transformer_block.attn1.original_forward = \
                            transformer_block.attn1.forward
                        def forward_with_pos(self, *inputs, **kwargs):
                            hidden_states = inputs[0]
                            hidden_states = self.pos_emb(hidden_states)
                            out = self.original_forward(hidden_states, *inputs[1:], **kwargs)
                            return out

                        bound_method = forward_with_pos.__get__(
                            transformer_block.attn1, 
                            transformer_block.attn1.__class__
                        )
                        setattr(transformer_block.attn1, 'forward', bound_method)
                        
                        transformer_block.attn2.pos_emb = get_pos_emb(
                            dim=transformer_block.attn2.to_q.in_features,
                            resolution=resolution_blocks[block_idx],
                            pos_type=pos_type
                        )
                        transformer_block.attn2.original_forward = \
                            transformer_block.attn2.forward
                        def forward_with_pos(self, *inputs, **kwargs):
                            hidden_states = inputs[0]
                            hidden_states = self.pos_emb(hidden_states)
                            out = self.original_forward(hidden_states, *inputs[1:], **kwargs)
                            return out

                        bound_method = forward_with_pos.__get__(
                            transformer_block.attn2, 
                            transformer_block.attn2.__class__
                        )
                        setattr(transformer_block.attn2, 'forward', bound_method)
    return

class UNet2DConditionModelWithPos(UNet2DConditionModel):

    # since it is inherited from UNet2DConditionModel, it supports options like 
    # enable_gradient_checkpointing and enable_xformers_memory_efficient_attention
    # for efficient training and inference
    _supports_gradient_checkpointing = True

    def __init__(
        self,
        sample_size: Optional[int] = None,
        in_channels: int = 4,
        out_channels: int = 4,
        down_block_types: Tuple[str] = (
            "CrossAttnDownBlock2D",
            "CrossAttnDownBlock2D",
            "CrossAttnDownBlock2D",
            "DownBlock2D",
        ),
        up_block_types: Tuple[str] = (
            "UpBlock2D",
            "CrossAttnUpBlock2D",
            "CrossAttnUpBlock2D",
            "CrossAttnUpBlock2D"
        ),
        block_out_channels: Tuple[int] = (320, 640, 1280, 1280),
        layers_per_block: Union[int, Tuple[int]] = 2,
        norm_num_groups: Optional[int] = 32,
        cross_attention_dim: Union[int, Tuple[int]] = 1280,
        transformer_layers_per_block: Union[int, Tuple[int]] = 1,
        attention_head_dim: Union[int, Tuple[int]] = 8,
        num_attention_heads: Optional[Union[int, Tuple[int]]] = None,
        pos_type: str ='cartesian',
        **kwargs
    ):
        super().__init__(
            sample_size=sample_size,
            in_channels=in_channels,
            out_channels=out_channels,
            down_block_types=down_block_types,
            up_block_types=up_block_types,
            block_out_channels=block_out_channels,
            layers_per_block=layers_per_block,
            norm_num_groups=norm_num_groups,
            cross_attention_dim=cross_attention_dim,
            transformer_layers_per_block=transformer_layers_per_block,
            attention_head_dim=attention_head_dim,
            num_attention_heads=num_attention_heads,
            **kwargs
        )
        self.register_to_config(pos_type=pos_type)
        hack_unet_for_pos(self, pos_type=pos_type)


if __name__ == "__main__":
    unet = UNet2DConditionModelWithPos(
        sample_size = 32,
        in_channels = 4,
        out_channels = 4,
        cross_attention_dim=192, # slot size
        down_block_types = (
            "CrossAttnDownBlock2D",
            "CrossAttnDownBlock2D",
            "CrossAttnDownBlock2D",
            "DownBlock2D",
        ),
        up_block_types = (
            "UpBlock2D", 
            "CrossAttnUpBlock2D", 
            "CrossAttnUpBlock2D", 
            "CrossAttnUpBlock2D"
        ),
        block_out_channels = (192, 384, 768, 768),
        layers_per_block = 2,
        norm_num_groups = 24,
        transformer_layers_per_block = 1,
        attention_head_dim = 24,
        pos_type='cartesian'
    )

    # need to pass this small test
    feat = torch.randn(1, 4, 32, 32)
    encoder_hidden_states = torch.randn(1, 10, 192)
    unet(feat, timestep=0, encoder_hidden_states=encoder_hidden_states)

    unet.save_config("./configs/movi-e/unet/")

    pass
src/models/unet_with_pos.py
JindongJiang/latent-slot-diffusion
import copy
import random
import distinctipy
from PIL import Image

import torch
import numpy as np
from torch import nn

from torchvision import transforms
from einops import rearrange, repeat
import matplotlib.cm as mpl_color_map
from torchvision.utils import make_grid
from einops.layers.torch import Rearrange
from einops import rearrange, repeat, reduce

def get_1d_sincos_pos_embed_from_grid(embed_dim, pos):
    """
    embed_dim: output dimension for each position
    pos: a list of positions to be encoded: size (M,)
    out: (M, D)
    """
    assert embed_dim % 2 == 0
    omega = np.arange(embed_dim // 2, dtype=float)
    omega /= embed_dim / 2.
    omega = 1. / 10000 ** omega  # (D/2,)

    pos = pos.reshape(-1)  # (M,)
    out = np.einsum('m,d->md', pos, omega)  # (M, D/2), outer product

    emb_sin = np.sin(out)  # (M, D/2)
    emb_cos = np.cos(out)  # (M, D/2)

    emb = np.concatenate([emb_sin, emb_cos], axis=1)  # (M, D)
    return emb


def get_2d_sincos_pos_embed_from_grid(embed_dim, grid):
    assert embed_dim % 2 == 0

    # use half of dimensions to encode grid_h
    emb_h = get_1d_sincos_pos_embed_from_grid(embed_dim // 2, grid[0])  # (H*W, D/2)
    emb_w = get_1d_sincos_pos_embed_from_grid(embed_dim // 2, grid[1])  # (H*W, D/2)

    emb = np.concatenate([emb_h, emb_w], axis=1)  # (H*W, D)
    return emb


def get_2d_sincos_pos_embed(embed_dim, grid_size, cls_token=False):
    """
    grid_size: int of the grid height and width
    return:
    pos_embed: [grid_size*grid_size, embed_dim] or [1+grid_size*grid_size, embed_dim] (w/ or w/o cls_token)
    """
    grid_h = np.arange(grid_size, dtype=float)
    grid_w = np.arange(grid_size, dtype=float)
    grid = np.meshgrid(grid_w, grid_h, indexing='xy')  # here w goes first
    grid = np.stack(grid, axis=0)

    grid = grid.reshape([2, 1, grid_size, grid_size])
    pos_embed = get_2d_sincos_pos_embed_from_grid(embed_dim, grid)
    if cls_token:
        pos_embed = np.concatenate([np.zeros([1, embed_dim]), pos_embed], axis=0)
    return pos_embed


class SinCosPositionalEmbedding2D(nn.Module):

    def __init__(self, embed_dim, grid_size, learnable_gamma=True):
        super().__init__()
        pos_embed = get_2d_sincos_pos_embed(embed_dim, grid_size, cls_token=False)
        pos_embed = rearrange(pos_embed, '(h w) d -> 1 d h w', h=grid_size, w=grid_size)
        # self.pos_embed = nn.Parameter(torch.from_numpy(pos_embed).float(), requires_grad=False)
        if learnable_gamma:
            self.gamma = nn.Parameter(torch.tensor(1.), requires_grad=True)
        else:
            self.gamma = 1.

        self.register_buffer('pos_embed', torch.from_numpy(pos_embed).float())

    def forward(self, input):
        """
        input: batch_size x d_model x H x W
        """
        return input + self.gamma * self.pos_embed


class LearnedPositionalEmbedding1D(nn.Module):

    def __init__(self, num_inputs, input_size, dropout=0.0):
        super().__init__()
        self.dropout = nn.Dropout(dropout)
        self.pe = nn.Parameter(torch.zeros(1, num_inputs, input_size), requires_grad=True)
        nn.init.trunc_normal_(self.pe, std=0.02)

    def forward(self, input):
        """
        input: batch_size x seq_len x d_model
        return: batch_size x seq_len x d_model
        """
        return self.dropout(input + self.pe)


class CartesianPositionalEmbedding(nn.Module):

    def __init__(self, channels, image_size):
        super().__init__()

        # self.projection = conv2d(4, channels, 1)
        self.projection = nn.Conv2d(4, channels, 1)
        # self.pe = nn.Parameter(self.build_grid(image_size).unsqueeze(0), requires_grad=False)

        self.register_buffer('pe', self.build_grid(image_size).unsqueeze(0))

    def build_grid(self, side_length):
        coords = torch.linspace(0., 1., side_length + 1)
        coords = 0.5 * (coords[:-1] + coords[1:])
        grid_y, grid_x = torch.meshgrid(coords, coords, indexing='xy')
        return torch.stack((grid_x, grid_y, 1 - grid_x, 1 - grid_y), dim=0)

    def forward(self, inputs):
        # `inputs` has shape: [batch_size, out_channels, height, width].
        # `grid` has shape: [batch_size, in_channels, height, width].
        return inputs + self.projection(self.pe)


def random_colors(N, randomize=False, rng=42):
    colors = distinctipy.get_colors(N, rng=rng)
    if randomize:
        random.shuffle(colors)
    return colors

class ColorMask(object):
    def __init__(self, num_slots, log_img_size, norm_mean, norm_std, rng=42, img_tmp_pth=None):
        self.img_tmp_pth = img_tmp_pth
        self.num_slots = num_slots
        self.log_img_size = log_img_size
        self.color = torch.tensor(random_colors(num_slots, randomize=False, rng=rng))
        self.log_image_resize = transforms.Resize(log_img_size,
                                                  interpolation=transforms.InterpolationMode.BILINEAR,
                                                  antialias=True)
        self.norm_mean = norm_mean
        self.norm_std = norm_std
        self.img_unnorm = transforms.Compose([
            transforms.Normalize(mean=[0., 0., 0.],
                                 std=1 / torch.tensor(norm_std)),
            transforms.Normalize(mean=-torch.tensor(norm_mean),
                                 std=[1., 1., 1.])
        ])

    def apply_colormap_on_image(self, org_im, activation, colormap_name, alpha=0.5):
        """
            Apply heatmap on image
        Args:
            org_img (PIL img): Original image
            activation_map (numpy arr): Activation map (grayscale) 0-255
            colormap_name (str): Name of the colormap
        """
        # Get colormap
        color_map = mpl_color_map.get_cmap(colormap_name)
        no_trans_heatmap = color_map(activation)
        # Change alpha channel in colormap to make sure original image is displayed
        heatmap = copy.copy(no_trans_heatmap)
        heatmap[:, :, 3] = alpha
        heatmap = Image.fromarray((heatmap * 255).astype(np.uint8))
        no_trans_heatmap = Image.fromarray((no_trans_heatmap * 255).astype(np.uint8))

        # Apply heatmap on image
        heatmap_on_image = Image.new("RGBA", org_im.size)
        heatmap_on_image = Image.alpha_composite(heatmap_on_image, org_im.convert('RGBA'))
        heatmap_on_image = Image.alpha_composite(heatmap_on_image, heatmap)
        return no_trans_heatmap, heatmap_on_image

    def _apply_mask(self, image, mask, alpha=0.5, color=None):

        B, C, H, W = image.size()
        B, N, H, W = mask.size()

        image = image.clone()
        mask_only = torch.ones_like(image)

        if color is None:
            color = random_colors(N)

        for n in range(N):
            for c in range(3):
                image[..., c, :, :] = torch.where(
                    mask[:, n] == 1,
                    image[..., c, :, :] * (1 - alpha) + alpha * (color[n][c] if isinstance(color, list) else
                                                                 color[..., n, c][..., None, None]),
                    image[..., c, :, :]
                )
                mask_only[..., c, :, :] = torch.where(
                    mask[:, n] == 1,
                    mask_only[..., c, :, :] * (1 - 0.99) + 0.99 * (color[n][c] if isinstance(color, list) else
                                                                   color[..., n, c][..., None, None]),
                    mask_only[..., c, :, :]
                )
        return image, mask_only

    def get_heatmap(self, img, attn, recon=None, mask_pred_sorted=None, return_all=False):
        '''

                :param img: b, c, h, w
                :param attn: b, s, h, w
                :param name:
                :param global_step:
                :return:
                '''
        img = img.to(torch.device('cpu'), dtype=torch.float32)
        attn = attn.to(torch.device('cpu'), dtype=torch.float32)
        if recon is not None:
            if not isinstance(recon, list):
                recon = [recon]
            recon = [r.to(torch.device('cpu'), dtype=torch.float32) for r in recon]
        if mask_pred_sorted is not None:
            mask_pred_sorted = mask_pred_sorted.to(torch.device('cpu'))
        bs, inp_channel, h, w = img.size()

        img = self.img_unnorm(img).clamp(0., 1.)
        if recon is not None:
            recon = [self.img_unnorm(r).clamp(0., 1.) for r in recon]

        if h > self.log_img_size:
            img = self.log_image_resize(img)
            h, w = img.shape[-2:]

        num_s = attn.size(1)
        # --------------------------------------------------------------------------
        # get color map
        if mask_pred_sorted is None:
            mask_pred = (attn.argmax(1, keepdim=True) == torch.arange(attn.size(1))[None, :, None, None]).float()
        else:
            mask_pred = mask_pred_sorted

        if mask_pred.shape[-2] != img.shape[-2] or mask_pred.shape[-1] != img.shape[-1]:
            mask_pred = transforms.Resize(size=img.shape[-2:], interpolation=transforms.InterpolationMode.NEAREST)(
                mask_pred)

        # b c h w
        img_overlay, color_mask = self._apply_mask(img, mask_pred, alpha=0.5, color=self.color)

        # --------------------------------------------------------------------------

        if attn.shape[-2] != img.shape[-2] or attn.shape[-1] != img.shape[-1]:
            # attn = transforms.Resize(size=img.shape[-2:], interpolation=transforms.InterpolationMode.NEAREST,
            #                          antialias=True)(attn)
            attn = transforms.Resize(size=img.shape[-2:], interpolation=transforms.InterpolationMode.BILINEAR,
                                     antialias=True)(attn)

        attn = rearrange(attn, 'b s h w -> b s (h w)')

        # attn = (attn - attn.min(dim=-1, keepdim=True)[0]) / \
        #        (attn.max(dim=-1, keepdim=True)[0] - attn.min(dim=-1, keepdim=True)[0] + 1e-8)

        attn = rearrange(attn, 'b s h_w -> (b s) h_w').detach().numpy()

        img_reshape = repeat(img, 'b c h w -> c (b s) (h w) ', s=num_s)

        img_pil = transforms.ToPILImage()(img_reshape)

        no_trans_heatmap, heatmap_on_image = self.apply_colormap_on_image(img_pil, attn, 'gray')

        heatmap_on_image = transforms.ToTensor()(heatmap_on_image.convert('RGB'))

        heatmap_on_image = rearrange(heatmap_on_image, 'c (b s) (h w) -> b s c h w', b=bs, c=inp_channel, h=h, w=w)

        grid_image = torch.cat([img[:, None], img_overlay[:, None], heatmap_on_image], dim=1)
        if recon is not None:
            if not isinstance(recon, list):
                recon = [recon]
            recon = [self.log_image_resize(r) if r.shape[-2] != h else r for r in recon]
            grid_image = torch.cat([*[r[:, None] for r in recon if r is not None], grid_image], dim=1)
        grid_image = make_grid(rearrange(grid_image, 'b n c h w -> (b n) c h w'),
                               nrow=grid_image.size(1), padding=1,
                               pad_value=0.8)  # clamp to prevent overshooting caused by bicubic interpolation
        if return_all:
            return grid_image, img_overlay, color_mask, heatmap_on_image
        return grid_image

    def log_heatmap(self, img, attn, recon=None, mask_pred_sorted=None, path=None):
        assert path is not None or self.img_tmp_pth is not None, 'path is None and img_tmp_pth is None'

        grid_image = self.get_heatmap(img, attn, recon, mask_pred_sorted)
        # save_image(grid_image, self.img_tmp_pth)
        ndarr = grid_image.mul(255).add_(0.5).clamp_(0, 255).permute(1, 2, 0).to('cpu', torch.uint8).numpy()
        im = Image.fromarray(ndarr)
        img_path = path if path is not None else self.img_tmp_pth
        im.save(img_path, optimize=True, quality=95)
        return img_path
src/models/utils.py
JindongJiang/latent-slot-diffusion
import argparse
import copy
import gc
import hashlib
import importlib
import itertools
import logging
import math
import os
import shutil
import warnings
from pathlib import Path

import numpy as np
import torch
import torch.nn.functional as F
import torch.utils.checkpoint
import transformers
from accelerate import Accelerator
from accelerate.logging import get_logger
from accelerate.utils import ProjectConfiguration, set_seed
from huggingface_hub import create_repo
from packaging import version
from PIL import Image
from tqdm.auto import tqdm

import diffusers
from diffusers import (
    AutoencoderKL,
    DDPMScheduler,
    StableDiffusionPipeline,
)
# from diffusers.training_utils import compute_snr # diffusers is still working on this yet, uncomment in the future version
from diffusers.utils import check_min_version, is_wandb_available
from diffusers.utils.import_utils import is_xformers_available

from einops import rearrange, reduce, repeat

from src.models.backbone import UNetEncoder
from src.models.slot_attn import MultiHeadSTEVESA
from src.models.unet_with_pos import UNet2DConditionModelWithPos
from src.data.dataset import GlobDataset

from configs.parser import parse_args

from src.models.utils import ColorMask

if is_wandb_available():
    import wandb

logger = get_logger(__name__)


@torch.no_grad()
def log_validation(
    val_dataset,
    backbone,
    slot_attn,
    unet,
    vae,
    scheduler,
    args,
    accelerator,
    weight_dtype,
    global_step,
):
    logger.info(
        f"Running validation... \n."
    )
    unet = accelerator.unwrap_model(unet)
    backbone = accelerator.unwrap_model(backbone)
    slot_attn = accelerator.unwrap_model(slot_attn)

    colorizer = ColorMask(
        num_slots=slot_attn.config.num_slots,
        log_img_size=256,
        norm_mean=0,
        norm_std=1,
    )

    val_dataloader = torch.utils.data.DataLoader(
        val_dataset,
        batch_size=args.val_batch_size,
        shuffle=False,
        num_workers=args.dataloader_num_workers,
    )

    # We train on the simplified learning objective. If we were previously predicting a variance, we need the scheduler to ignore it
    scheduler_args = {}

    if "variance_type" in scheduler.config:
        variance_type = scheduler.config.variance_type

        if variance_type in ["learned", "learned_range"]:
            variance_type = "fixed_small"

        scheduler_args["variance_type"] = variance_type

    # use a more efficient scheduler at test time
    module = importlib.import_module("diffusers")
    scheduler_class = getattr(module, args.validation_scheduler)
    scheduler = scheduler_class.from_config(
        scheduler.config, **scheduler_args)

    pipeline = StableDiffusionPipeline(
        vae=vae,
        text_encoder=None,
        tokenizer=None,
        unet=unet,
        scheduler=scheduler,
        safety_checker=None,
        feature_extractor=None,
        requires_safety_checker=None,
    )

    pipeline = pipeline.to(accelerator.device)
    pipeline.set_progress_bar_config(disable=True)

    # run inference
    generator = None if args.seed is None else torch.Generator(
        device=accelerator.device).manual_seed(args.seed)

    num_digits = len(str(args.max_train_steps))
    folder_name = f"image_logging_{global_step:0{num_digits}}"
    image_log_dir = os.path.join(accelerator.logging_dir, folder_name, )
    os.makedirs(image_log_dir, exist_ok=True)

    images = []
    image_count = 0

    for batch_idx, batch in enumerate(val_dataloader):

        pixel_values = batch["pixel_values"].to(
            device=accelerator.device, dtype=weight_dtype)

        with torch.autocast("cuda"):
            model_input = vae.encode(pixel_values).latent_dist.sample()
            pixel_values_recon = vae.decode(model_input).sample

            if args.backbone_config == "pretrain_dino":
                pixel_values_vit = batch["pixel_values_vit"].to(device=accelerator.device, 
                                                                dtype=weight_dtype)
                feat = backbone(pixel_values_vit)
            else:
                feat = backbone(pixel_values)
            slots, attn = slot_attn(feat[:, None])  # for the time dimension
            slots = slots[:, 0]
            images_gen = pipeline(
                prompt_embeds=slots,
                height=args.resolution,
                width=args.resolution,
                num_inference_steps=25,
                generator=generator,
                guidance_scale=1.,
                output_type="pt",
            ).images

        grid_image = colorizer.get_heatmap(img=(pixel_values * 0.5 + 0.5),
                                           attn=reduce(
                                               attn[:, 0], 'b num_h (h w) s -> b s h w', h=int(np.sqrt(attn.shape[-2])), 
                                               reduction='mean'
                                           ),
                                           recon=[pixel_values_recon * 0.5 + 0.5, images_gen])
        ndarr = grid_image.mul(255).add_(0.5).clamp_(0, 255).permute(1, 2, 0).to('cpu', torch.uint8).numpy()
        im = Image.fromarray(ndarr)
        images.append(im)
        img_path = os.path.join(image_log_dir, f"image_{batch_idx:02}.jpg")
        im.save(img_path, optimize=True, quality=95)
        image_count += pixel_values.shape[0]
        if image_count >= args.num_validation_images:
            break

    for tracker in accelerator.trackers:
        if tracker.name == "tensorboard":
            np_images = np.stack([np.asarray(img) for img in images])
            tracker.writer.add_images(
                "validation", np_images, global_step, dataformats="NHWC")
        if tracker.name == "wandb":
            tracker.log(
                {
                    "validation": [
                        wandb.Image(image, caption=f"{i}") for i, image in enumerate(images)
                    ]
                }
            )

    del pipeline
    torch.cuda.empty_cache()

    return images


def main(args):
    logging_dir = Path(args.output_dir, args.logging_dir)

    accelerator_project_config = ProjectConfiguration(
        project_dir=args.output_dir, logging_dir=logging_dir
    )

    accelerator = Accelerator(
        gradient_accumulation_steps=args.gradient_accumulation_steps,
        mixed_precision=args.mixed_precision,
        log_with=args.report_to,
        project_config=accelerator_project_config,
    )

    if args.report_to == "wandb":
        if not is_wandb_available():
            raise ImportError(
                "Make sure to install wandb if you want to use it for logging during training.")

    # Make one log on every process with the configuration for debugging.
    logging.basicConfig(
        format="%(asctime)s - %(levelname)s - %(name)s - %(message)s",
        datefmt="%m/%d/%Y %H:%M:%S",
        level=logging.INFO,
    )
    logger.info(accelerator.state, main_process_only=False)
    if accelerator.is_local_main_process:
        transformers.utils.logging.set_verbosity_warning()
        diffusers.utils.logging.set_verbosity_info()
    else:
        transformers.utils.logging.set_verbosity_error()
        diffusers.utils.logging.set_verbosity_error()

    # If passed along, set the training seed now.
    if args.seed is not None:
        set_seed(args.seed)

    # Handle the repository creation
    if accelerator.is_main_process:
        if args.output_dir is not None:
            os.makedirs(args.output_dir, exist_ok=True)

        if args.push_to_hub:
            repo_id = create_repo(
                repo_id=args.hub_model_id or Path(args.output_dir).name, exist_ok=True, token=args.hub_token
            ).repo_id

    # Load scheduler and models
    noise_scheduler_config = DDPMScheduler.load_config(args.scheduler_config)
    noise_scheduler = DDPMScheduler.from_config(noise_scheduler_config)

    vae = AutoencoderKL.from_pretrained(
        args.pretrained_model_name, subfolder="vae")
    
    backbone_config = UNetEncoder.load_config(args.backbone_config)
    backbone = UNetEncoder.from_config(backbone_config)

    slot_attn_config = MultiHeadSTEVESA.load_config(args.slot_attn_config)
    slot_attn = MultiHeadSTEVESA.from_config(slot_attn_config)

    unet_config = UNet2DConditionModelWithPos.load_config(args.unet_config)
    unet = UNet2DConditionModelWithPos.from_config(unet_config)

    # create custom saving & loading hooks so that `accelerator.save_state(...)` serializes in a nice format

    def save_model_hook(models, weights, output_dir):
        if accelerator.is_main_process:
            for model in models:

                # continue if not one of [UNetEncoder, MultiHeadSTEVESA, UNet2DConditionModelWithPos]
                if not isinstance(model, (UNetEncoder, MultiHeadSTEVESA, UNet2DConditionModelWithPos)):
                    continue

                sub_dir = model._get_name().lower()

                model.save_pretrained(os.path.join(output_dir, sub_dir))

                # make sure to pop weight so that corresponding model is not saved again
                weights.pop()

    def load_model_hook(models, input_dir):
        while len(models) > 0:
            # pop models so that they are not loaded again
            model = models.pop()

            sub_dir = model._get_name().lower()

            if isinstance(model, UNetEncoder):
                # load diffusers style into model
                load_model = UNetEncoder.from_pretrained(
                    input_dir, subfolder=sub_dir)
                model.register_to_config(**load_model.config)
            elif isinstance(model, MultiHeadSTEVESA):
                load_model = MultiHeadSTEVESA.from_pretrained(
                    input_dir, subfolder=sub_dir)
                model.register_to_config(**load_model.config)
            elif isinstance(model, UNet2DConditionModelWithPos):
                load_model = UNet2DConditionModelWithPos.from_pretrained(
                    input_dir, subfolder=sub_dir)
                model.register_to_config(**load_model.config)
            else:
                raise ValueError(
                    f"Unknown model type {type(model)}")

            model.load_state_dict(load_model.state_dict())
            del load_model

    accelerator.register_save_state_pre_hook(save_model_hook)
    accelerator.register_load_state_pre_hook(load_model_hook)

    vae.requires_grad_(False)
    if args.enable_xformers_memory_efficient_attention:
        if is_xformers_available():
            import xformers

            xformers_version = version.parse(xformers.__version__)
            if xformers_version == version.parse("0.0.16"):
                logger.warn(
                    "xFormers 0.0.16 cannot be used for training in some GPUs. If you observe problems during training, please update xFormers to at least 0.0.17. See https://huggingface.co/docs/diffusers/main/en/optimization/xformers for more details."
                )
            unet.enable_xformers_memory_efficient_attention()
            try:
                backbone.enable_xformers_memory_efficient_attention()
            except AttributeError:
                pass
        else:
            raise ValueError(
                "xformers is not available. Make sure it is installed correctly")

    if args.gradient_checkpointing:
        unet.enable_gradient_checkpointing()
        backbone.enable_gradient_checkpointing()

    # Check that all trainable models are in full precision
    low_precision_error_string = (
        "Please make sure to always have all model weights in full float32 precision when starting training - even if"
        " doing mixed precision training. copy of the weights should still be float32."
    )

    if accelerator.unwrap_model(unet).dtype != torch.float32:
        raise ValueError(
            f"Unet loaded as datatype {accelerator.unwrap_model(unet).dtype}. {low_precision_error_string}"
        )

    if accelerator.unwrap_model(backbone).dtype != torch.float32:
        raise ValueError(
            f"Backbone loaded as datatype {accelerator.unwrap_model(backbone).dtype}. {low_precision_error_string}"
        )

    if accelerator.unwrap_model(slot_attn).dtype != torch.float32:
        raise ValueError(
            f"Slot Attn loaded as datatype {accelerator.unwrap_model(slot_attn).dtype}. {low_precision_error_string}"
        )

    # Enable TF32 for faster training on Ampere GPUs,
    # cf https://pytorch.org/docs/stable/notes/cuda.html#tensorfloat-32-tf32-on-ampere-devices
    if args.allow_tf32:
        torch.backends.cuda.matmul.allow_tf32 = True

    if args.scale_lr:
        args.learning_rate = (
            args.learning_rate * args.gradient_accumulation_steps *
            args.train_batch_size * accelerator.num_processes
        )

    # Use 8-bit Adam for lower memory usage or to fine-tune the model in 16GB GPUs
    if args.use_8bit_adam:
        try:
            import bitsandbytes as bnb
        except ImportError:
            raise ImportError(
                "To use 8-bit Adam, please install the bitsandbytes library: `pip install bitsandbytes`."
            )

        optimizer_class = bnb.optim.AdamW8bit
    else:
        optimizer_class = torch.optim.AdamW


    params_to_optimize = list(slot_attn.parameters()) + \
        list(backbone.parameters()) + list(unet.up_blocks.parameters())
    params_group = [
        {'params': list(slot_attn.parameters()) + list(backbone.parameters()),
         'lr': args.learning_rate * args.encoder_lr_scale}
    ]
    params_group.append(
        {'params': unet.parameters(), "lr": args.learning_rate}
    )

    optimizer = optimizer_class(
        params_group,
        lr=args.learning_rate,
        betas=(args.adam_beta1, args.adam_beta2),
        weight_decay=args.adam_weight_decay,
        eps=args.adam_epsilon,
    )

    lr_scheduler = torch.optim.lr_scheduler.LambdaLR(
        optimizer, lr_lambda=[lambda _: 1, lambda _: 1]
        )

    train_dataset = GlobDataset(
        root=args.dataset_root,
        img_size=args.resolution,
        img_glob=args.dataset_glob,
        data_portion=(0.0, args.train_split_portion),
        random_flip=args.flip_images,
    )

    train_dataloader = torch.utils.data.DataLoader(
        train_dataset,
        batch_size=args.train_batch_size,
        shuffle=True,
        num_workers=args.dataloader_num_workers,
    )

    val_dataset = GlobDataset(
        root=args.dataset_root,
        img_size=args.resolution,
        img_glob=args.dataset_glob,
        data_portion=(args.train_split_portion if args.train_split_portion < 1. else 0.9, 1.0),
    )

    # Scheduler and math around the number of training steps.
    overrode_max_train_steps = False
    num_update_steps_per_epoch = math.ceil(
        len(train_dataloader) / args.gradient_accumulation_steps)
    if args.max_train_steps is None:
        args.max_train_steps = args.num_train_epochs * num_update_steps_per_epoch
        overrode_max_train_steps = True

    # Prepare everything with our `accelerator`.
    backbone, unet, slot_attn, optimizer, train_dataloader, lr_scheduler = accelerator.prepare(
        backbone, unet, slot_attn, optimizer, train_dataloader, lr_scheduler
    )

    # For mixed precision training we cast all non-trainable weigths (vae, non-lora text_encoder and non-lora unet) to half-precision
    # as these weights are only used for inference, keeping weights in full precision is not required.
    weight_dtype = torch.float32
    if accelerator.mixed_precision == "fp16":
        weight_dtype = torch.float16
    elif accelerator.mixed_precision == "bf16":
        weight_dtype = torch.bfloat16

    # Move vae device and cast to weight_dtype
    vae.to(accelerator.device, dtype=weight_dtype)

    # We need to recalculate our total training steps as the size of the training dataloader may have changed.
    num_update_steps_per_epoch = math.ceil(
        len(train_dataloader) / args.gradient_accumulation_steps)
    if overrode_max_train_steps:
        args.max_train_steps = args.num_train_epochs * num_update_steps_per_epoch
    # Afterwards we recalculate our number of training epochs
    args.num_train_epochs = math.ceil(
        args.max_train_steps / num_update_steps_per_epoch)

    # We need to initialize the trackers we use, and also store our configuration.
    # The trackers initializes automatically on the main process.
    if accelerator.is_main_process:
        tracker_config = vars(copy.deepcopy(args))
        accelerator.init_trackers(
            args.tracker_project_name, config=tracker_config
        )

    # Train!
    total_batch_size = args.train_batch_size * \
        accelerator.num_processes * args.gradient_accumulation_steps

    logger.info("***** Running training *****")
    logger.info(f"  Num examples = {len(train_dataset)}")
    logger.info(f"  Num batches each epoch = {len(train_dataloader)}")
    logger.info(f"  Num Epochs = {args.num_train_epochs}")
    logger.info(
        f"  Instantaneous batch size per device = {args.train_batch_size}")
    logger.info(
        f"  Total train batch size (w. parallel, distributed & accumulation) = {total_batch_size}")
    logger.info(
        f"  Gradient Accumulation steps = {args.gradient_accumulation_steps}")
    logger.info(f"  Total optimization steps = {args.max_train_steps}")
    global_step = 0
    first_epoch = 0
    accumulate_steps = 0 # necessary for args.gradient_accumulation_steps > 1

    # Potentially load in the weights and states from a previous save
    if args.resume_from_checkpoint:
        if args.resume_from_checkpoint != "latest":
            path = os.path.basename(args.resume_from_checkpoint.rstrip('/')) # only the checkpoint folder name is needed, not the full path
        else:
            # Get the most recent checkpoint
            dirs = os.listdir(args.output_dir)
            dirs = [d for d in dirs if d.startswith("checkpoint")]
            dirs = sorted(dirs, key=lambda x: int(x.split("-")[1]))
            path = dirs[-1] if len(dirs) > 0 else None

        if path is None:
            accelerator.print(
                f"Checkpoint '{args.resume_from_checkpoint}' does not exist. Starting a new training run."
            )
            args.resume_from_checkpoint = None
            initial_global_step = 0
        else:
            accelerator.print(f"Resuming from checkpoint {path}")
            accelerator.load_state(os.path.join(args.output_dir, path))
            global_step = int(path.split("-")[1])

            initial_global_step = global_step
            accumulate_steps = global_step * args.gradient_accumulation_steps
            first_epoch = global_step // num_update_steps_per_epoch
    else:
        initial_global_step = 0

    # Only show the progress bar once on each machine.
    progress_bar = tqdm(
        range(0, args.max_train_steps),
        initial=initial_global_step,
        desc="Steps",
        # Only show the progress bar once on each machine.
        disable=not accelerator.is_local_main_process,
        position=0, leave=True
    )

    for epoch in range(first_epoch, args.num_train_epochs):
        unet.train()
        backbone.train()
        slot_attn.train()
        for step, batch in enumerate(train_dataloader):
            pixel_values = batch["pixel_values"].to(dtype=weight_dtype)

            # Convert images to latent space
            model_input = vae.encode(pixel_values).latent_dist.sample()
            model_input = model_input * vae.config.scaling_factor

            # Sample noise that we'll add to the model input
            if args.offset_noise:
                noise = torch.randn_like(model_input) + 0.1 * torch.randn(
                    model_input.shape[0], model_input.shape[1], 1, 1, device=model_input.device
                )
            else:
                noise = torch.randn_like(model_input)
            bsz, channels, height, width = model_input.shape
            # Sample a random timestep for each image
            timesteps = torch.randint(
                0, noise_scheduler.config.num_train_timesteps, (bsz,), device=model_input.device
            )
            timesteps = timesteps.long()

            # Add noise to the model input according to the noise magnitude at each timestep
            # (this is the forward diffusion process)
            noisy_model_input = noise_scheduler.add_noise(
                model_input, noise, timesteps)

            # timestep is not used, but should we?
            if args.backbone_config == "pretrain_dino":
                pixel_values_vit = batch["pixel_values_vit"].to(dtype=weight_dtype)
                feat = backbone(pixel_values_vit)
            else:
                feat = backbone(pixel_values)
            slots, attn = slot_attn(feat[:, None])  # for the time dimension
            slots = slots[:, 0]

            # Predict the noise residual
            model_pred = unet(
                noisy_model_input, timesteps, slots,
            ).sample

            # Get the target for loss depending on the prediction type
            if noise_scheduler.config.prediction_type == "epsilon":
                target = noise
            elif noise_scheduler.config.prediction_type == "v_prediction":
                target = noise_scheduler.get_velocity(
                    model_input, noise, timesteps)
            else:
                raise ValueError(
                    f"Unknown prediction type {noise_scheduler.config.prediction_type}")

            # Compute instance loss
            if args.snr_gamma is None:
                loss = F.mse_loss(model_pred.float(),
                                  target.float(), reduction="mean")
            else:
                # Compute loss-weights as per Section 3.4 of https://arxiv.org/abs/2303.09556.
                # Since we predict the noise instead of x_0, the original formulation is slightly changed.
                # This is discussed in Section 4.2 of the same paper.
                snr = compute_snr(noise_scheduler, timesteps)
                base_weight = (
                    torch.stack(
                        [snr, args.snr_gamma * torch.ones_like(timesteps)], dim=1).min(dim=1)[0] / snr
                )

                if noise_scheduler.config.prediction_type == "v_prediction":
                    # Velocity objective needs to be floored to an SNR weight of one.
                    mse_loss_weights = base_weight + 1
                else:
                    # Epsilon and sample both use the same loss weights.
                    mse_loss_weights = base_weight
                loss = F.mse_loss(model_pred.float(),
                                  target.float(), reduction="none")
                loss = loss.mean(
                    dim=list(range(1, len(loss.shape)))) * mse_loss_weights
                loss = loss.mean()

            loss = loss / args.gradient_accumulation_steps

            # if args.with_prior_preservation:
            #     # Add the prior loss to the instance loss.
            #     loss = loss + args.prior_loss_weight * prior_loss

            accelerator.backward(loss)
            accumulate_steps += 1
            # if accelerator.sync_gradients:
            if (accumulate_steps+1) % args.gradient_accumulation_steps == 0:
                params_to_clip = params_to_optimize
                accelerator.clip_grad_norm_(
                    params_to_clip, args.max_grad_norm)
                optimizer.step()
                lr_scheduler.step()
                optimizer.zero_grad(set_to_none=args.set_grads_to_none)

            # Checks if the accelerator has performed an optimization step behind the scenes
            if (accumulate_steps+1) % args.gradient_accumulation_steps == 0:
                progress_bar.update(1)
                global_step += 1

                if accelerator.is_main_process:
                    if global_step % args.checkpointing_steps == 0:
                        # _before_ saving state, check if this save would set us over the `checkpoints_total_limit`
                        if args.checkpoints_total_limit is not None:
                            checkpoints = os.listdir(args.output_dir)
                            checkpoints = [
                                d for d in checkpoints if d.startswith("checkpoint")]
                            checkpoints = sorted(
                                checkpoints, key=lambda x: int(x.split("-")[1]))

                            # before we save the new checkpoint, we need to have at _most_ `checkpoints_total_limit - 1` checkpoints
                            if len(checkpoints) >= args.checkpoints_total_limit:
                                num_to_remove = len(
                                    checkpoints) - args.checkpoints_total_limit + 1
                                removing_checkpoints = checkpoints[0:num_to_remove]

                                logger.info(
                                    f"{len(checkpoints)} checkpoints already exist, removing {len(removing_checkpoints)} checkpoints"
                                )
                                logger.info(
                                    f"removing checkpoints: {', '.join(removing_checkpoints)}")

                                for removing_checkpoint in removing_checkpoints:
                                    removing_checkpoint = os.path.join(
                                        args.output_dir, removing_checkpoint)
                                    shutil.rmtree(removing_checkpoint)

                        save_path = os.path.join(
                            args.output_dir, f"checkpoint-{global_step}")
                        accelerator.save_state(save_path)
                        logger.info(f"Saved state to {save_path}")

                    images = []

                    if global_step % args.validation_steps == 0:
                        images = log_validation(
                            val_dataset=val_dataset,
                            backbone=backbone,
                            slot_attn=slot_attn,
                            unet=unet,
                            vae=vae,
                            scheduler=noise_scheduler,
                            args=args,
                            accelerator=accelerator,
                            weight_dtype=weight_dtype,
                            global_step=global_step,
                        )

            logs = {"loss": loss.detach().item(
            ), "lr": lr_scheduler.get_last_lr()[0]}
            progress_bar.set_postfix(**logs)
            accelerator.log(logs, step=global_step)

            if global_step >= args.max_train_steps:
                break

    # Create the pipeline using using the trained modules and save it.
    accelerator.wait_for_everyone()
    if accelerator.is_main_process:
        save_path = os.path.join(
            args.output_dir, f"checkpoint-{global_step}-last")
        accelerator.save_state(save_path)
        logger.info(f"Saved state to {save_path}")

    accelerator.end_training()


if __name__ == "__main__":
    args = parse_args()
    main(args)

train_lsd.py
JindongJiang/latent-slot-diffusion