|
from typing import Dict, List |
|
import os |
|
|
|
import torch |
|
import torch.nn as nn |
|
import torch.nn.functional as F |
|
import numpy as np |
|
import numpy.random as npr |
|
import copy |
|
from functools import partial |
|
from contextlib import contextmanager |
|
|
|
from .common.get_model import get_model, register |
|
from .sd import DDPM |
|
|
|
version = '0' |
|
symbol = 'codi' |
|
|
|
|
|
@register('codi', version) |
|
class CoDi(DDPM): |
|
def __init__(self, |
|
audioldm_cfg=None, |
|
autokl_cfg=None, |
|
optimus_cfg=None, |
|
clip_cfg=None, |
|
clap_cfg=None, |
|
vision_scale_factor=0.1812, |
|
text_scale_factor=4.3108, |
|
audio_scale_factor=0.9228, |
|
scale_by_std=False, |
|
*args, |
|
**kwargs): |
|
super().__init__(*args, **kwargs) |
|
|
|
if audioldm_cfg is not None: |
|
self.audioldm = get_model()(audioldm_cfg) |
|
|
|
if autokl_cfg is not None: |
|
self.autokl = get_model()(autokl_cfg) |
|
|
|
if optimus_cfg is not None: |
|
self.optimus = get_model()(optimus_cfg) |
|
|
|
if clip_cfg is not None: |
|
self.clip = get_model()(clip_cfg) |
|
|
|
if clap_cfg is not None: |
|
self.clap = get_model()(clap_cfg) |
|
|
|
if not scale_by_std: |
|
self.vision_scale_factor = vision_scale_factor |
|
self.text_scale_factor = text_scale_factor |
|
self.audio_scale_factor = audio_scale_factor |
|
else: |
|
self.register_buffer("text_scale_factor", torch.tensor(text_scale_factor)) |
|
self.register_buffer("audio_scale_factor", torch.tensor(audio_scale_factor)) |
|
self.register_buffer('vision_scale_factor', torch.tensor(vision_scale_factor)) |
|
|
|
@property |
|
def device(self): |
|
return next(self.parameters()).device |
|
|
|
@torch.no_grad() |
|
def autokl_encode(self, image): |
|
encoder_posterior = self.autokl.encode(image) |
|
z = encoder_posterior.sample().to(image.dtype) |
|
return self.vision_scale_factor * z |
|
|
|
@torch.no_grad() |
|
def autokl_decode(self, z): |
|
z = 1. / self.vision_scale_factor * z |
|
return self.autokl.decode(z) |
|
|
|
@torch.no_grad() |
|
def optimus_encode(self, text): |
|
if isinstance(text, List): |
|
token = [self.optimus.tokenizer_encoder.tokenize(sentence.lower()) for sentence in text] |
|
token_id = [] |
|
for tokeni in token: |
|
token_sentence = [self.optimus.tokenizer_encoder._convert_token_to_id(i) for i in tokeni] |
|
token_sentence = self.optimus.tokenizer_encoder.add_special_tokens_single_sentence(token_sentence) |
|
token_id.append(torch.LongTensor(token_sentence)) |
|
token_id = torch._C._nn.pad_sequence(token_id, batch_first=True, padding_value=0.0)[:, :512] |
|
else: |
|
token_id = text |
|
token_id = token_id.to(self.device) |
|
z = self.optimus.encoder(token_id, attention_mask=(token_id > 0))[1] |
|
z_mu, z_logvar = self.optimus.encoder.linear(z).chunk(2, -1) |
|
return z_mu.squeeze(1) * self.text_scale_factor |
|
|
|
@torch.no_grad() |
|
def optimus_decode(self, z, temperature=1.0, max_length=30): |
|
z = 1.0 / self.text_scale_factor * z |
|
z = z.to(self.device) |
|
return self.optimus.decode(z, temperature, max_length=max_length) |
|
|
|
@torch.no_grad() |
|
def audioldm_encode(self, audio, time=2.0): |
|
encoder_posterior = self.audioldm.encode(audio, time=time) |
|
z = encoder_posterior.sample().to(audio.dtype) |
|
return z * self.audio_scale_factor |
|
|
|
@torch.no_grad() |
|
def audioldm_decode(self, z): |
|
if torch.max(torch.abs(z)) > 1e2: |
|
z = torch.clip(z, min=-10, max=10) |
|
z = 1.0 / self.audio_scale_factor * z |
|
return self.audioldm.decode(z) |
|
|
|
@torch.no_grad() |
|
def mel_spectrogram_to_waveform(self, mel): |
|
|
|
if len(mel.size()) == 4: |
|
mel = mel.squeeze(1) |
|
mel = mel.permute(0, 2, 1) |
|
waveform = self.audioldm.vocoder(mel) |
|
waveform = waveform.cpu().detach().numpy() |
|
return waveform |
|
|
|
@torch.no_grad() |
|
def clip_encode_text(self, text, encode_type='encode_text'): |
|
swap_type = self.clip.encode_type |
|
self.clip.encode_type = encode_type |
|
embedding = self.clip(text, encode_type) |
|
self.clip.encode_type = swap_type |
|
return embedding |
|
|
|
@torch.no_grad() |
|
def clip_encode_vision(self, vision, encode_type='encode_vision'): |
|
swap_type = self.clip.encode_type |
|
self.clip.encode_type = encode_type |
|
embedding = self.clip(vision, encode_type) |
|
self.clip.encode_type = swap_type |
|
return embedding |
|
|
|
@torch.no_grad() |
|
def clap_encode_audio(self, audio): |
|
embedding = self.clap(audio) |
|
return embedding |
|
|
|
def forward(self, x=None, c=None, noise=None, xtype='image', ctype='prompt', u=None, return_algined_latents=False): |
|
if isinstance(x, list): |
|
t = torch.randint(0, self.num_timesteps, (x[0].shape[0],), device=x[0].device).long() |
|
else: |
|
t = torch.randint(0, self.num_timesteps, (x.shape[0],), device=x.device).long() |
|
return self.p_losses(x, c, t, noise, xtype, ctype, u, return_algined_latents) |
|
|
|
def apply_model(self, x_noisy, t, cond, xtype='image', ctype='text', u=None, return_algined_latents=False): |
|
return self.model.diffusion_model(x_noisy, t, cond, xtype, ctype, u, return_algined_latents) |
|
|
|
def get_pixel_loss(self, pred, target, mean=True): |
|
if self.loss_type == 'l1': |
|
loss = (target - pred).abs() |
|
if mean: |
|
loss = loss.mean() |
|
elif self.loss_type == 'l2': |
|
if mean: |
|
loss = torch.nn.functional.mse_loss(target, pred) |
|
else: |
|
loss = torch.nn.functional.mse_loss(target, pred, reduction='none') |
|
else: |
|
raise NotImplementedError("unknown loss type '{loss_type}'") |
|
loss = torch.nan_to_num(loss, nan=0.0, posinf=0.0, neginf=-0.0) |
|
return loss |
|
|
|
def get_text_loss(self, pred, target): |
|
if self.loss_type == 'l1': |
|
loss = (target - pred).abs() |
|
elif self.loss_type == 'l2': |
|
loss = torch.nn.functional.mse_loss(target, pred, reduction='none') |
|
loss = torch.nan_to_num(loss, nan=0.0, posinf=0.0, neginf=0.0) |
|
return loss |
|
|
|
def p_losses(self, x_start, cond, t, noise=None, xtype='image', ctype='prompt', u=None, return_algined_latents=False): |
|
if isinstance(x_start, list): |
|
noise = [torch.randn_like(x_start_i) for x_start_i in x_start] if noise is None else noise |
|
x_noisy = [self.q_sample(x_start=x_start_i, t=t, noise=noise_i) for x_start_i, noise_i in zip(x_start, noise)] |
|
model_output = self.apply_model(x_noisy, t, cond, xtype, ctype, u, return_algined_latents) |
|
if return_algined_latents: |
|
return model_output |
|
|
|
loss_dict = {} |
|
|
|
if self.parameterization == "x0": |
|
target = x_start |
|
elif self.parameterization == "eps": |
|
target = noise |
|
else: |
|
raise NotImplementedError() |
|
|
|
loss = 0.0 |
|
for model_output_i, target_i, xtype_i in zip(model_output, target, xtype): |
|
if xtype_i == 'image': |
|
loss_simple = self.get_pixel_loss(model_output_i, target_i, mean=False).mean([1, 2, 3]) |
|
elif xtype_i == 'video': |
|
loss_simple = self.get_pixel_loss(model_output_i, target_i, mean=False).mean([1, 2, 3, 4]) |
|
elif xtype_i == 'text': |
|
loss_simple = self.get_text_loss(model_output_i, target_i).mean([1]) |
|
elif xtype_i == 'audio': |
|
loss_simple = self.get_pixel_loss(model_output_i, target_i, mean=False).mean([1, 2, 3]) |
|
loss += loss_simple.mean() |
|
return loss / len(xtype) |
|
|
|
else: |
|
noise = torch.randn_like(x_start) if noise is None else noise |
|
x_noisy = self.q_sample(x_start=x_start, t=t, noise=noise) |
|
model_output = self.apply_model(x_noisy, t, cond, xtype, ctype) |
|
|
|
loss_dict = {} |
|
|
|
if self.parameterization == "x0": |
|
target = x_start |
|
elif self.parameterization == "eps": |
|
target = noise |
|
else: |
|
raise NotImplementedError() |
|
|
|
if xtype == 'image': |
|
loss_simple = self.get_pixel_loss(model_output, target, mean=False).mean([1, 2, 3]) |
|
elif xtype == 'video': |
|
loss_simple = self.get_pixel_loss(model_output, target, mean=False).mean([1, 2, 3, 4]) |
|
elif xtype == 'text': |
|
loss_simple = self.get_text_loss(model_output, target).mean([1]) |
|
elif xtype == 'audio': |
|
loss_simple = self.get_pixel_loss(model_output, target, mean=False).mean([1, 2, 3]) |
|
loss = loss_simple.mean() |
|
return loss |
|
|
