video3d/diffusion/vsd.py

import os
os.environ['HUGGINGFACE_HUB_CACHE'] = '/viscam/u/zzli'
os.environ['HF_HOME'] = '/viscam/u/zzli'

from transformers import CLIPTextModel, CLIPTokenizer, logging
from diffusers import AutoencoderKL, UNet2DConditionModel, PNDMScheduler, DDIMScheduler

from diffusers.loaders import AttnProcsLayers
from diffusers.models.attention_processor import LoRAAttnProcessor
from diffusers.models.embeddings import TimestepEmbedding
from diffusers.utils.import_utils import is_xformers_available

# Suppress partial model loading warning
logging.set_verbosity_error()

import gc
import random
import torch
import torch.nn as nn
import torch.nn.functional as F
import tinycudann as tcnn
from video3d.diffusion.sd import StableDiffusion
from torch.cuda.amp import custom_bwd, custom_fwd 


def seed_everything(seed):
    torch.manual_seed(seed)
    torch.cuda.manual_seed(seed)

def cleanup():
    gc.collect()
    torch.cuda.empty_cache()
    tcnn.free_temporary_memory()

class StableDiffusion_VSD(StableDiffusion):
    def __init__(self, device, sd_version='2.1', hf_key=None, torch_dtype=torch.float32, lora_n_timestamp_samples=1):
        super().__init__(device, sd_version=sd_version, hf_key=hf_key, torch_dtype=torch_dtype)

        # self.device = device
        # self.sd_version = sd_version
        # self.torch_dtype = torch_dtype
        
        if hf_key is not None:
            print(f'[INFO] using hugging face custom model key: {hf_key}')
            model_key = hf_key
        elif self.sd_version == '2.1':
            model_key = "stabilityai/stable-diffusion-2-1-base"
        elif self.sd_version == '2.0':
            model_key = "stabilityai/stable-diffusion-2-base"
        elif self.sd_version == '1.5':
            model_key = "runwayml/stable-diffusion-v1-5"
        else:
            raise ValueError(f'Stable-diffusion version {self.sd_version} not supported.')

        # # Create model
        # self.vae = AutoencoderKL.from_pretrained(model_key, subfolder="vae", torch_dtype=torch_dtype).to(self.device)
        # self.tokenizer = CLIPTokenizer.from_pretrained(model_key, subfolder="tokenizer")
        # self.text_encoder = CLIPTextModel.from_pretrained(model_key, subfolder="text_encoder").to(self.device)
        # self.unet = UNet2DConditionModel.from_pretrained(model_key, subfolder="unet", torch_dtype=torch_dtype).to(self.device)
        
        # self.scheduler = DDIMScheduler.from_pretrained(model_key, subfolder="scheduler")
        # # self.scheduler = PNDMScheduler.from_pretrained(model_key, subfolder="scheduler")

        # self.num_train_timesteps = self.scheduler.config.num_train_timesteps
        # self.alphas = self.scheduler.alphas_cumprod.to(self.device) # for convenience

        print(f'[INFO] loading stable diffusion VSD modules...')

        self.unet_lora = UNet2DConditionModel.from_pretrained(model_key, subfolder="unet", torch_dtype=torch_dtype).to(self.device)
        cleanup()
        
        for p in self.vae.parameters():
            p.requires_grad_(False)
        for p in self.text_encoder.parameters():
            p.requires_grad_(False)
        for p in self.unet.parameters():
            p.requires_grad_(False)
        for p in self.unet_lora.parameters():
            p.requires_grad_(False)
        
        # set up LoRA layers
        lora_attn_procs = {}
        for name in self.unet_lora.attn_processors.keys():
            cross_attention_dim = (
                None
                if name.endswith("attn1.processor")
                else self.unet_lora.config.cross_attention_dim
            )
            if name.startswith("mid_block"):
                hidden_size = self.unet_lora.config.block_out_channels[-1]
            elif name.startswith("up_blocks"):
                block_id = int(name[len("up_blocks.")])
                hidden_size = list(reversed(self.unet_lora.config.block_out_channels))[
                    block_id
                ]
            elif name.startswith("down_blocks"):
                block_id = int(name[len("down_blocks.")])
                hidden_size = self.unet_lora.config.block_out_channels[block_id]

            lora_attn_procs[name] = LoRAAttnProcessor(
                hidden_size=hidden_size, cross_attention_dim=cross_attention_dim
            )

        self.unet_lora.set_attn_processor(lora_attn_procs)

        self.lora_layers = AttnProcsLayers(self.unet_lora.attn_processors).to(
            self.device
        )
        self.lora_layers._load_state_dict_pre_hooks.clear()
        self.lora_layers._state_dict_hooks.clear()
        self.lora_n_timestamp_samples = lora_n_timestamp_samples
        self.scheduler_lora = DDIMScheduler.from_pretrained(model_key, subfolder="scheduler")

        print(f'[INFO] loaded stable diffusion VSD modules!')

    def train_lora(
        self,
        latents,
        text_embeddings,
        camera_condition
    ):
        B = latents.shape[0]
        lora_n_timestamp_samples = self.lora_n_timestamp_samples
        latents = latents.detach().repeat(lora_n_timestamp_samples, 1, 1, 1)

        t = torch.randint(
            int(self.num_train_timesteps * 0.0),
            int(self.num_train_timesteps * 1.0),
            [B * lora_n_timestamp_samples],
            dtype=torch.long,
            device=self.device,
        )

        noise = torch.randn_like(latents)
        noisy_latents = self.scheduler_lora.add_noise(latents, noise, t)
        if self.scheduler_lora.config.prediction_type == "epsilon":
            target = noise
        elif self.scheduler_lora.config.prediction_type == "v_prediction":
            target = self.scheduler_lora.get_velocity(latents, noise, t)
        else:
            raise ValueError(
                f"Unknown prediction type {self.scheduler_lora.config.prediction_type}"
            )
        
        # use view-independent text embeddings in LoRA
        _, text_embeddings_cond = text_embeddings.chunk(2)

        if random.random() < 0.1:
            camera_condition = torch.zeros_like(camera_condition)
        
        noise_pred = self.unet_lora(
            noisy_latents,
            t,
            encoder_hidden_states=text_embeddings_cond.repeat(
                lora_n_timestamp_samples, 1, 1
            ),
            class_labels=camera_condition.reshape(B, -1).repeat(
                lora_n_timestamp_samples, 1
            ),
            cross_attention_kwargs={"scale": 1.0}
        ).sample

        loss_lora = 0.5 * F.mse_loss(noise_pred.float(), target.float(), reduction="mean")
        return loss_lora

    
    def train_step(
        self,
        text_embeddings,
        text_embeddings_vd, 
        pred_rgb,
        camera_condition,
        im_features,
        guidance_scale=7.5,
        guidance_scale_lora=7.5,
        loss_weight=1.0,
        min_step_pct=0.02,
        max_step_pct=0.98, 
        return_aux=False
    ):
        pred_rgb = pred_rgb.to(self.torch_dtype)
        text_embeddings = text_embeddings.to(self.torch_dtype)
        text_embeddings_vd = text_embeddings_vd.to(self.torch_dtype)
        camera_condition = camera_condition.to(self.torch_dtype)
        im_features = im_features.to(self.torch_dtype)

        # condition_label = camera_condition
        condition_label = im_features
        
        b = pred_rgb.shape[0]
        
        # interp to 512x512 to be fed into vae.
        # _t = time.time()
        pred_rgb_512 = F.interpolate(pred_rgb, (512, 512), mode='bilinear', align_corners=False)
        # torch.cuda.synchronize(); print(f'[TIME] guiding: interp {time.time() - _t:.4f}s')

        # timestep ~ U(0.02, 0.98) to avoid very high/low noise level
        min_step = int(self.num_train_timesteps * min_step_pct)
        max_step = int(self.num_train_timesteps * max_step_pct)
        t = torch.randint(min_step, max_step + 1, [b], dtype=torch.long, device=self.device)

        # encode image into latents with vae, requires grad!
        # _t = time.time()
        latents = self.encode_imgs(pred_rgb_512)
        # torch.cuda.synchronize(); print(f'[TIME] guiding: vae enc {time.time() - _t:.4f}s')

        # predict the noise residual with unet, NO grad!
        # _t = time.time()
        with torch.no_grad():
            # add noise
            noise = torch.randn_like(latents)
            latents_noisy = self.scheduler.add_noise(latents, noise, t)
            # pred noise
            latent_model_input = torch.cat([latents_noisy] * 2)
            
            # disable unet class embedding here
            cls_embedding = self.unet.class_embedding
            self.unet.class_embedding = None

            cross_attention_kwargs = None
            noise_pred_pretrain = self.unet(
                latent_model_input, 
                torch.cat([t, t]), 
                encoder_hidden_states=text_embeddings_vd, 
                class_labels=None, 
                cross_attention_kwargs=cross_attention_kwargs
            ).sample

            self.unet.class_embedding = cls_embedding
            
            # use view-independent text embeddings in LoRA
            _, text_embeddings_cond = text_embeddings.chunk(2)
            
            noise_pred_est = self.unet_lora(
                latent_model_input,
                torch.cat([t, t]),
                encoder_hidden_states=torch.cat([text_embeddings_cond] * 2),
                class_labels=torch.cat(
                    [
                        condition_label.reshape(b, -1),
                        torch.zeros_like(condition_label.reshape(b, -1)),
                    ],
                    dim=0,
                ),
                cross_attention_kwargs={"scale": 1.0},
            ).sample

        noise_pred_pretrain_uncond, noise_pred_pretrain_text = noise_pred_pretrain.chunk(2)
        
        noise_pred_pretrain = noise_pred_pretrain_uncond + guidance_scale * (
            noise_pred_pretrain_text - noise_pred_pretrain_uncond
        )

        assert self.scheduler.config.prediction_type == "epsilon"
        if self.scheduler_lora.config.prediction_type == "v_prediction":
            alphas_cumprod = self.scheduler_lora.alphas_cumprod.to(
                device=latents_noisy.device, dtype=latents_noisy.dtype
            )
            alpha_t = alphas_cumprod[t] ** 0.5
            sigma_t = (1 - alphas_cumprod[t]) ** 0.5

            noise_pred_est = latent_model_input * torch.cat([sigma_t] * 2, dim=0).reshape(
                -1, 1, 1, 1
            ) + noise_pred_est * torch.cat([alpha_t] * 2, dim=0).reshape(-1, 1, 1, 1)
        
        noise_pred_est_uncond, noise_pred_est_camera = noise_pred_est.chunk(2)

        noise_pred_est = noise_pred_est_uncond + guidance_scale_lora * (
            noise_pred_est_camera - noise_pred_est_uncond
        )

        # w(t), sigma_t^2
        w = (1 - self.alphas[t])
        # w = self.alphas[t] ** 0.5 * (1 - self.alphas[t])
        grad = loss_weight * w[:, None, None, None] * (noise_pred_pretrain - noise_pred_est)

        grad = torch.nan_to_num(grad)

        targets = (latents - grad).detach()
        loss_vsd = 0.5 * F.mse_loss(latents.float(), targets, reduction='sum') / latents.shape[0]

        loss_lora = self.train_lora(latents, text_embeddings, condition_label)
        
        loss = {
            'loss_vsd': loss_vsd,
            'loss_lora': loss_lora
        }

        if return_aux:
            aux = {'grad': grad, 't': t, 'w': w}
            return loss, aux
        else:
            return loss 
    


if __name__ == '__main__':
    import argparse
    import matplotlib.pyplot as plt

    parser = argparse.ArgumentParser()
    parser.add_argument('prompt', type=str)
    parser.add_argument('--negative', default='', type=str)
    parser.add_argument('--sd_version', type=str, default='2.1', choices=['1.5', '2.0', '2.1'], help="stable diffusion version")
    parser.add_argument('--hf_key', type=str, default=None, help="hugging face Stable diffusion model key")
    parser.add_argument('-H', type=int, default=512)
    parser.add_argument('-W', type=int, default=512)
    parser.add_argument('--seed', type=int, default=0)
    parser.add_argument('--steps', type=int, default=50)
    opt = parser.parse_args()

    seed_everything(opt.seed)

    device = torch.device('cuda')

    sd = StableDiffusion_VSD(device, opt.sd_version, opt.hf_key)

    imgs = sd.prompt_to_img(opt.prompt, opt.negative, opt.H, opt.W, opt.steps)

    # visualize image
    plt.imshow(imgs[0])
    plt.show()
    plt.savefig(f'{opt.prompt}.png')