Diffusers Implementation of Controlling Text-to-Image Diffusion by Orthogonal Finetuning
import torch
import torch.nn as nn

# based on LoRACrossAttnProcessor


class OFTCrossAttnProcessor(nn.Module):
    def __init__(self, hidden_size, r=8, constraint=1e-3):
        super().__init__()

        self.hidden_size = hidden_size
        self.r = r
        self.block_size = int(hidden_size / r)
        self.constraint = constraint * hidden_size

        # block diagonals that turn into the identity matrix
        self.W_Q = nn.Parameter(torch.zeros(r, self.block_size, self.block_size))
        self.W_K = nn.Parameter(torch.zeros(r, self.block_size, self.block_size))
        self.W_V = nn.Parameter(torch.zeros(r, self.block_size, self.block_size))

    def __call__(
        self, attn, hidden_states, encoder_hidden_states=None, attention_mask=None
    ):
        batch_size, sequence_length, _ = hidden_states.shape
        attention_mask = attn.prepare_attention_mask(
            attention_mask, sequence_length, batch_size)

        encoder_hidden_states = encoder_hidden_states if encoder_hidden_states is not None else hidden_states

        query = attn.to_q(hidden_states)
        key = attn.to_k(encoder_hidden_states)
        value = attn.to_v(encoder_hidden_states)

        # OFT changes
        block_Q_Q = self.W_Q - self.W_Q.transpose(1, 2)
        block_Q_K = self.W_K - self.W_K.transpose(1, 2)
        block_Q_V = self.W_V - self.W_V.transpose(1, 2)

        norm_Q_Q = torch.norm(block_Q_Q.flatten())
        norm_Q_K = torch.norm(block_Q_K.flatten())
        norm_Q_V = torch.norm(block_Q_V.flatten())

        new_norm_Q_Q = torch.clamp(norm_Q_Q, max=self.constraint)
        new_norm_Q_K = torch.clamp(norm_Q_K, max=self.constraint)
        new_norm_Q_V = torch.clamp(norm_Q_V, max=self.constraint)

        block_Q_Q = block_Q_Q * ((new_norm_Q_Q + 1e-8) / (norm_Q_Q + 1e-8))
        block_Q_K = block_Q_K * ((new_norm_Q_K + 1e-8) / (norm_Q_K + 1e-8))
        block_Q_V = block_Q_V * ((new_norm_Q_V + 1e-8) / (norm_Q_V + 1e-8))

        I = torch.eye(self.block_size, device=key.device).unsqueeze(0).repeat(
            self.r, 1, 1
        )

        block_R_Q = torch.matmul(I + block_Q_Q, (I - block_Q_Q).inverse())
        block_R_K = torch.matmul(I + block_Q_K, (I - block_Q_K).inverse())
        block_R_V = torch.matmul(I + block_Q_V, (I - block_Q_V).inverse())

        R_Q = torch.block_diag(*block_R_Q).to(key.dtype)
        R_K = torch.block_diag(*block_R_K).to(key.dtype)
        R_V = torch.block_diag(*block_R_V).to(key.dtype)

        query = torch.matmul(query, R_Q)
        key = torch.matmul(key, R_K)
        value = torch.matmul(value, R_V)
        # end OFT changes

        key = attn.head_to_batch_dim(key)
        value = attn.head_to_batch_dim(value)
        query = attn.head_to_batch_dim(query)

        attention_probs = attn.get_attention_scores(query, key, attention_mask)
        hidden_states = torch.bmm(attention_probs, value)
        hidden_states = attn.batch_to_head_dim(hidden_states)

        # linear proj
        hidden_states = attn.to_out[0](hidden_states)
        # dropout
        hidden_states = attn.to_out[1](hidden_states)

        return hidden_states

# save to file


def save_attn_processors(unet, device, dtype, save_path):
    attn_processors = unet.attn_processors
    keys = list(attn_processors.keys())
    weights_dict = {}
    parameters_dict = {}

    for key in keys:
        processor = attn_processors[key].to(device).to(dtype)
        weights_dict[key] = processor.state_dict()
        parameters_dict[key] = {
            'hidden_size': processor.hidden_size,
            'r': processor.r,
            'constraint': processor.constraint
        }

    output_dict = {
        'weights': weights_dict,
        'parameters': parameters_dict
    }

    torch.save(output_dict, save_path)


# load from file
def load_attn_processors(unet, device, dtype, save_path):
    input_dict = torch.load(save_path)
    weights_dict = input_dict['weights']
    parameters_dict = input_dict['parameters']

    keys = list(weights_dict.keys())

    attn_processors = {}

    for key in keys:
        attn_processors[key] = OFTCrossAttnProcessor(
            hidden_size=parameters_dict[key]['hidden_size'],
            r=parameters_dict[key]['r'],
            constraint=parameters_dict[key]['constraint']
        ).to(device).to(dtype)
        attn_processors[key].load_state_dict(weights_dict[key])

    unet.set_attn_processor(attn_processors)

attention_processor.py
tripplyons/oft

import argparse
import torch

argparser = argparse.ArgumentParser()
argparser.add_argument('--processor_path', type=str)
argparser.add_argument('--output_path', type=str)
argparser.add_argument('--model_path', type=str)

args = argparser.parse_args()

processor_path = args.processor_path
output_path = args.output_path
model_path = args.model_path

processor = torch.load(processor_path, map_location='cpu')
state_dict = torch.load(model_path, map_location='cpu')

# from https://gist.github.com/jachiam/8a5c0b607e38fcc585168b90c686eb05

layer_types_to_processor = {
    'to_q': 'W_Q',
    'to_k': 'W_K',
    'to_v': 'W_V',
}

unet_conversion_map = [
    # (stable-diffusion, HF Diffusers)
    ("time_embed.0.weight", "time_embedding.linear_1.weight"),
    ("time_embed.0.bias", "time_embedding.linear_1.bias"),
    ("time_embed.2.weight", "time_embedding.linear_2.weight"),
    ("time_embed.2.bias", "time_embedding.linear_2.bias"),
    ("input_blocks.0.0.weight", "conv_in.weight"),
    ("input_blocks.0.0.bias", "conv_in.bias"),
    ("out.0.weight", "conv_norm_out.weight"),
    ("out.0.bias", "conv_norm_out.bias"),
    ("out.2.weight", "conv_out.weight"),
    ("out.2.bias", "conv_out.bias"),
]

unet_conversion_map_resnet = [
    # (stable-diffusion, HF Diffusers)
    ("in_layers.0", "norm1"),
    ("in_layers.2", "conv1"),
    ("out_layers.0", "norm2"),
    ("out_layers.3", "conv2"),
    ("emb_layers.1", "time_emb_proj"),
    ("skip_connection", "conv_shortcut"),
]

unet_conversion_map_layer = []
# hardcoded number of downblocks and resnets/attentions...
# would need smarter logic for other networks.
for i in range(4):
    # loop over downblocks/upblocks

    for j in range(2):
        # loop over resnets/attentions for downblocks
        hf_down_res_prefix = f"down_blocks.{i}.resnets.{j}."
        sd_down_res_prefix = f"input_blocks.{3*i + j + 1}.0."
        unet_conversion_map_layer.append((sd_down_res_prefix, hf_down_res_prefix))

        if i < 3:
            # no attention layers in down_blocks.3
            hf_down_atn_prefix = f"down_blocks.{i}.attentions.{j}."
            sd_down_atn_prefix = f"input_blocks.{3*i + j + 1}.1."
            unet_conversion_map_layer.append((sd_down_atn_prefix, hf_down_atn_prefix))

    for j in range(3):
        # loop over resnets/attentions for upblocks
        hf_up_res_prefix = f"up_blocks.{i}.resnets.{j}."
        sd_up_res_prefix = f"output_blocks.{3*i + j}.0."
        unet_conversion_map_layer.append((sd_up_res_prefix, hf_up_res_prefix))

        if i > 0:
            # no attention layers in up_blocks.0
            hf_up_atn_prefix = f"up_blocks.{i}.attentions.{j}."
            sd_up_atn_prefix = f"output_blocks.{3*i + j}.1."
            unet_conversion_map_layer.append((sd_up_atn_prefix, hf_up_atn_prefix))

    if i < 3:
        # no downsample in down_blocks.3
        hf_downsample_prefix = f"down_blocks.{i}.downsamplers.0.conv."
        sd_downsample_prefix = f"input_blocks.{3*(i+1)}.0.op."
        unet_conversion_map_layer.append((sd_downsample_prefix, hf_downsample_prefix))

        # no upsample in up_blocks.3
        hf_upsample_prefix = f"up_blocks.{i}.upsamplers.0."
        sd_upsample_prefix = f"output_blocks.{3*i + 2}.{1 if i == 0 else 2}."
        unet_conversion_map_layer.append((sd_upsample_prefix, hf_upsample_prefix))

hf_mid_atn_prefix = "mid_block.attentions.0."
sd_mid_atn_prefix = "middle_block.1."
unet_conversion_map_layer.append((sd_mid_atn_prefix, hf_mid_atn_prefix))

for j in range(2):
    hf_mid_res_prefix = f"mid_block.resnets.{j}."
    sd_mid_res_prefix = f"middle_block.{2*j}."
    unet_conversion_map_layer.append((sd_mid_res_prefix, hf_mid_res_prefix))

def get_name_for_processor(name):
    for sd_name, hf_name in unet_conversion_map:
        name = name.replace(sd_name, hf_name)
    for sd_name, hf_name in unet_conversion_map_layer:
        name = name.replace(sd_name, hf_name)
    return name

for k in state_dict.keys():
    if ".attn2." in k and 'to_out' not in k and '.weight' in k and 'model.diffusion_model.' in k:
        # print(k, v)
        layer = '.'.join(k.split(".")[:-2]).replace('model.diffusion_model.', '')
        processor_key = get_name_for_processor(layer) + '.processor'
        weights = processor['weights'][processor_key]
        r = processor["parameters"][processor_key]["r"]
        hidden_size = processor["parameters"][processor_key]["hidden_size"]
        block_size = hidden_size // r

        layer_type = k.split(".")[-2]
        processor_name = layer_types_to_processor[layer_type]

        W = weights[processor_name]
        block_Q = W - W.transpose(1, 2)
        norm_Q = torch.norm(block_Q.flatten())
        new_norm_Q = torch.clamp(norm_Q, max=processor["parameters"][processor_key]["constraint"])
        block_Q = block_Q * ((new_norm_Q + 1e-8) / (norm_Q + 1e-8))
        I = torch.eye(block_size, device=W.device).unsqueeze(0).repeat(r, 1, 1)
        block_R = torch.matmul(I + block_Q, (I - block_Q).inverse())
        R = torch.block_diag(*block_R).to(torch.float32)
        state_dict[k] = torch.matmul(R.T, state_dict[k].to(torch.float32)).to(state_dict[k].dtype)

torch.save({
    'state_dict': state_dict
}, output_path)

merge_to_original.py
tripplyons/oft
# based on https://github.com/huggingface/diffusers/blob/main/examples/text_to_image/train_text_to_image_lora.py

import os
import json
import torch
import math
from tqdm.auto import tqdm
from accelerate import Accelerator
from datasets import Dataset
from torchvision import transforms
from diffusers import AutoencoderKL, DDPMScheduler, DiffusionPipeline, UNet2DConditionModel
from diffusers.loaders import AttnProcsLayers
from transformers import CLIPTextModel, CLIPTokenizer
from diffusers.optimization import get_scheduler
import numpy as np
import random
import pandas as pd
from PIL import Image
import torch.nn.functional as F

from attention_processor import OFTCrossAttnProcessor, save_attn_processors, load_attn_processors

# parameters

config_path = 'config'
output_dir = 'output'
gradient_accumulation_steps = 1
revision = None
mixed_precision = 'fp16'
weight_dtype = torch.float16
learning_rate = 1e-4
adam_beta1 = 0.9
adam_beta2 = 0.999
adam_weight_decay = 0
adam_epsilon = 1e-8
resolution = 512
train_batch_size = 1
dataloader_num_workers = 1
max_grad_norm = 1.0
lr_warmup_steps = 100
lr_scheduler_type = 'constant'
num_validation_images = 4
resume_load_path = None
use_8bit_optimizer = True

with open(os.path.join(config_path, "config.json")) as f:
    config = json.load(f)

validation_prompt = config['validation_prompt']
validation_negative_prompt = config['validation_negative_prompt']
validation_cfg = config['validation_cfg']
model_name = config['model_name']
num_train_epochs = config['num_train_epochs']
checkpointing_steps = config['checkpointing_steps']
validation_epochs = config['validation_epochs']

def main():
    accelerator = Accelerator(
        gradient_accumulation_steps=gradient_accumulation_steps,
        mixed_precision=mixed_precision
    )

    # models

    noise_scheduler = DDPMScheduler.from_pretrained(
        model_name, subfolder="scheduler"
    )
    tokenizer = CLIPTokenizer.from_pretrained(
        model_name, subfolder="tokenizer", revision=revision
    )
    text_encoder = CLIPTextModel.from_pretrained(
        model_name, subfolder="text_encoder", revision=revision
    )
    vae = AutoencoderKL.from_pretrained(
        model_name, subfolder="vae", revision=revision
    )
    unet = UNet2DConditionModel.from_pretrained(
        model_name, subfolder="unet", revision=revision
    )

    # only train the attention processors

    unet.requires_grad_(False)
    vae.requires_grad_(False)
    text_encoder.requires_grad_(False)

    # move to GPU

    unet.to(accelerator.device, dtype=weight_dtype)
    vae.to(accelerator.device, dtype=weight_dtype)
    text_encoder.to(accelerator.device, dtype=weight_dtype)

    # create or load attention processors

    if resume_load_path is not None:
        load_attn_processors(unet, 'cuda', torch.float32, resume_load_path)
    else:
        oft_attn_procs = {}
        for name in unet.attn_processors.keys():
            if name.startswith("mid_block"):
                hidden_size = unet.config.block_out_channels[-1]
            elif name.startswith("up_blocks"):
                block_id = int(name[len("up_blocks.")])
                hidden_size = list(reversed(unet.config.block_out_channels))[
                    block_id]
            elif name.startswith("down_blocks"):
                block_id = int(name[len("down_blocks.")])
                hidden_size = unet.config.block_out_channels[block_id]

            oft_attn_procs[name] = OFTCrossAttnProcessor(
                hidden_size=hidden_size,
            ).to(torch.float32)

        unet.set_attn_processor(oft_attn_procs)

    if use_8bit_optimizer:
        import bitsandbytes as bnb
        optimizer_cls = bnb.optim.AdamW8bit
    else:
        optimizer_cls = torch.optim.AdamW

    oft_layers = AttnProcsLayers(unet.attn_processors)

    optimizer = optimizer_cls(
        oft_layers.parameters(),
        lr=learning_rate,
        betas=(adam_beta1, adam_beta2),
        weight_decay=adam_weight_decay,
        eps=adam_epsilon,
    )

    # load dataset and transforms

    # dataset

    def tokenize_captions(examples):
        captions = [example['caption'] for example in examples]
        inputs = tokenizer(
            captions, max_length=tokenizer.model_max_length, padding="max_length", truncation=True, return_tensors="pt"
        )
        return inputs.input_ids

    train_transforms = transforms.Compose(
        [
            transforms.Resize(
                resolution, interpolation=transforms.InterpolationMode.BILINEAR
            ),
            transforms.CenterCrop(resolution),
            # transforms.RandomHorizontalFlip(),
            transforms.ToTensor(),
            transforms.Normalize([0.5], [0.5]),
        ]
    )

    def preprocess_train(examples):
        dataset = {}
        image_paths = [os.path.join(config_path, example['path']) for example in examples]
        images = [Image.open(image_path) for image_path in image_paths]
        images = [image.convert("RGB") for image in images]
        dataset["pixel_values"] = torch.stack([
            train_transforms(image) for image in images
        ])
        dataset["input_ids"] = tokenize_captions(examples)
        return dataset

    with accelerator.main_process_first():
        # Set the training transforms
        examples = config['examples']
        train_dataset = preprocess_train(examples)

    tensor_dataset = torch.utils.data.TensorDataset(
        train_dataset["pixel_values"], train_dataset["input_ids"]
    )

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

    num_update_steps_per_epoch = math.ceil(
        len(train_dataloader) / gradient_accumulation_steps)
    max_train_steps = num_train_epochs * num_update_steps_per_epoch

    lr_scheduler = get_scheduler(
        lr_scheduler_type,
        optimizer=optimizer,
        num_warmup_steps=lr_warmup_steps * gradient_accumulation_steps,
        num_training_steps=max_train_steps * gradient_accumulation_steps,
    )

    oft_layers, optimizer, train_dataloader, lr_scheduler = accelerator.prepare(
        oft_layers, optimizer, train_dataloader, lr_scheduler
    )

    if accelerator.is_main_process:
        accelerator.init_trackers("text2image-fine-tune")

    global_step = 0
    first_epoch = 0

    total_batch_size = train_batch_size * \
        accelerator.num_processes * gradient_accumulation_steps
    progress_bar = tqdm(range(global_step, max_train_steps),
                        disable=not accelerator.is_local_main_process)
    progress_bar.set_description("Steps")

    # training loop

    for epoch in range(first_epoch, num_train_epochs):
        unet.train()
        train_loss = 0.0
        for step, (pixel_values, input_ids) in enumerate(train_dataloader):
            with accelerator.accumulate(unet):
                # Convert images to latent space
                latents = vae.encode(pixel_values.to(
                    dtype=weight_dtype)).latent_dist.sample()
                latents = latents * vae.config.scaling_factor

                # Sample noise that we'll add to the latents
                noise = torch.randn_like(latents)
                bsz = latents.shape[0]
                # Sample a random timestep for each image
                timesteps = torch.randint(
                    0, noise_scheduler.num_train_timesteps, (bsz,), device=latents.device)
                timesteps = timesteps.long()

                # Add noise to the latents according to the noise magnitude at each timestep
                # (this is the forward diffusion process)
                noisy_latents = noise_scheduler.add_noise(
                    latents, noise, timesteps)

                # Get the text embedding for conditioning
                encoder_hidden_states = text_encoder(input_ids)[0]

                # 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(
                        latents, noise, timesteps)
                else:
                    raise ValueError(
                        f"Unknown prediction type {noise_scheduler.config.prediction_type}")

                # Predict the noise residual and compute loss
                model_pred = unet(noisy_latents, timesteps,
                                  encoder_hidden_states).sample
                loss = F.mse_loss(model_pred.float(),
                                  target.float(), reduction="mean")

                # Gather the losses across all processes for logging (if we use distributed training).
                avg_loss = accelerator.gather(
                    loss.repeat(train_batch_size)).mean()
                train_loss += avg_loss.item() / gradient_accumulation_steps

                # Backpropagate
                accelerator.backward(loss)
                if accelerator.sync_gradients:
                    params_to_clip = oft_layers.parameters()
                    accelerator.clip_grad_norm_(params_to_clip, max_grad_norm)
                optimizer.step()
                lr_scheduler.step()
                optimizer.zero_grad()

            # Checks if the accelerator has performed an optimization step behind the scenes
            if accelerator.sync_gradients:
                progress_bar.update(1)
                global_step += 1
                accelerator.log({"train_loss": train_loss}, step=global_step)
                train_loss = 0.0

                if global_step % checkpointing_steps == 0:
                    if accelerator.is_main_process:
                        save_path = os.path.join(
                            output_dir, f"checkpoint-{global_step}")
                        accelerator.save_state(save_path)

                        # also save the attention processors by themselves
                        save_attn_processors(unet, 'cuda', torch.float32,
                             os.path.join(save_path, "attn_processors.pt"))


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

            if global_step >= max_train_steps:
                break

        # validation check

        if accelerator.is_main_process:
            if validation_prompt is not None and (epoch + 1) % validation_epochs == 0:
                # create pipeline
                pipeline = DiffusionPipeline.from_pretrained(
                    model_name,
                    unet=accelerator.unwrap_model(unet),
                    revision=revision,
                    torch_dtype=weight_dtype,
                    safety_checker=None
                )
                pipeline = pipeline.to(accelerator.device)
                pipeline.set_progress_bar_config(disable=True)

                # run inference
                generator = torch.Generator(
                    device=accelerator.device)
                images = []
                for _ in range(num_validation_images):
                    images.append(
                        pipeline(
                            validation_prompt, negative_prompt=validation_negative_prompt, guidance_scale=validation_cfg,
                            num_inference_steps=30, generator=generator).images[0]
                    )

                # save images
                save_path = os.path.join(
                    output_dir, f"validation-images-{global_step}")
                os.makedirs(save_path, exist_ok=True)
                for i, image in enumerate(images):
                    image.save(os.path.join(save_path, f"{i}.png"))

                del pipeline
                torch.cuda.empty_cache()

    accelerator.wait_for_everyone()

    # save attention processors
    if accelerator.is_main_process:
        unet = unet.to(torch.float32)
        save_attn_processors(unet, 'cuda', torch.float32,
                             os.path.join(output_dir, "attn_processors.pt"))

    # free memory
    del unet
    del vae
    del text_encoder
    del optimizer
    torch.cuda.empty_cache()

    # Final inference
    # Load previous pipeline
    pipeline = DiffusionPipeline.from_pretrained(
        model_name, revision=revision, torch_dtype=weight_dtype, safety_checker=None
    )
    pipeline = pipeline.to(accelerator.device)

    # load attention processors
    load_attn_processors(pipeline.unet, 'cuda', torch.float32,
                         os.path.join(output_dir, "attn_processors.pt"))

    # run inference
    generator = torch.Generator(device=accelerator.device)
    images = []
    for _ in range(num_validation_images):
        images.append(pipeline(validation_prompt, negative_prompt=validation_negative_prompt, guidance_scale=validation_cfg,
                      num_inference_steps=30, generator=generator).images[0])

    # save images
    save_path = os.path.join(output_dir, f"final-images")
    os.makedirs(save_path, exist_ok=True)
    for i, image in enumerate(images):
        image.save(os.path.join(save_path, f"{i}.png"))

    accelerator.end_training()


if __name__ == "__main__":
    main()

train.py
tripplyons/oft