import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
from mvdream.camera_utils import get_camera, convert_opengl_to_blender, normalize_camera
from mvdream.model_zoo import build_model
from mvdream.ldm.models.diffusion.ddim import DDIMSampler
from diffusers import DDIMScheduler
class MVDream(nn.Module):
def __init__(
self,
device,
model_name='sd-v2.1-base-4view',
ckpt_path=None,
t_range=[0.02, 0.98],
):
super().__init__()
self.device = device
self.model_name = model_name
self.ckpt_path = ckpt_path
self.model = build_model(self.model_name, ckpt_path=self.ckpt_path).eval().to(self.device)
self.model.device = device
for p in self.model.parameters():
p.requires_grad_(False)
self.dtype = torch.float32
self.num_train_timesteps = 1000
self.min_step = int(self.num_train_timesteps * t_range[0])
self.max_step = int(self.num_train_timesteps * t_range[1])
self.embeddings = None
self.scheduler = DDIMScheduler.from_pretrained(
"stabilityai/stable-diffusion-2-1-base", subfolder="scheduler", torch_dtype=self.dtype
)
@torch.no_grad()
def get_text_embeds(self, prompts, negative_prompts):
pos_embeds = self.encode_text(prompts).repeat(4,1,1) # [1, 77, 768]
neg_embeds = self.encode_text(negative_prompts).repeat(4,1,1)
self.embeddings = torch.cat([neg_embeds, pos_embeds], dim=0) # [2, 77, 768]
def encode_text(self, prompt):
# prompt: [str]
embeddings = self.model.get_learned_conditioning(prompt).to(self.device)
return embeddings
@torch.no_grad()
def refine(self, pred_rgb, camera,
guidance_scale=100, steps=50, strength=0.8,
):
batch_size = pred_rgb.shape[0]
pred_rgb_256 = F.interpolate(pred_rgb, (256, 256), mode='bilinear', align_corners=False)
latents = self.encode_imgs(pred_rgb_256.to(self.dtype))
# latents = torch.randn((1, 4, 64, 64), device=self.device, dtype=self.dtype)
self.scheduler.set_timesteps(steps)
init_step = int(steps * strength)
latents = self.scheduler.add_noise(latents, torch.randn_like(latents), self.scheduler.timesteps[init_step])
camera = camera[:, [0, 2, 1, 3]] # to blender convention (flip y & z axis)
camera[:, 1] *= -1
camera = normalize_camera(camera).view(batch_size, 16)
camera = camera.repeat(2, 1)
context = {"context": self.embeddings, "camera": camera, "num_frames": 4}
for i, t in enumerate(self.scheduler.timesteps[init_step:]):
latent_model_input = torch.cat([latents] * 2)
tt = torch.cat([t.unsqueeze(0).repeat(batch_size)] * 2).to(self.device)
noise_pred = self.model.apply_model(latent_model_input, tt, context)
noise_pred_uncond, noise_pred_cond = noise_pred.chunk(2)
noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_cond - noise_pred_uncond)
latents = self.scheduler.step(noise_pred, t, latents).prev_sample
imgs = self.decode_latents(latents) # [1, 3, 512, 512]
return imgs
def train_step(
self,
pred_rgb, # [B, C, H, W], B is multiples of 4
camera, # [B, 4, 4]
step_ratio=None,
guidance_scale=50,
as_latent=False,
):
batch_size = pred_rgb.shape[0]
pred_rgb = pred_rgb.to(self.dtype)
if as_latent:
latents = F.interpolate(pred_rgb, (32, 32), mode="bilinear", align_corners=False) * 2 - 1
else:
# interp to 256x256 to be fed into vae.
pred_rgb_256 = F.interpolate(pred_rgb, (256, 256), mode="bilinear", align_corners=False)
# encode image into latents with vae, requires grad!
latents = self.encode_imgs(pred_rgb_256)
if step_ratio is not None:
# dreamtime-like
# t = self.max_step - (self.max_step - self.min_step) * np.sqrt(step_ratio)
t = np.round((1 - step_ratio) * self.num_train_timesteps).clip(self.min_step, self.max_step)
t = torch.full((batch_size,), t, dtype=torch.long, device=self.device)
else:
t = torch.randint(self.min_step, self.max_step + 1, (batch_size,), dtype=torch.long, device=self.device)
# camera = convert_opengl_to_blender(camera)
# flip_yz = torch.tensor([[1, 0, 0, 0], [0, 0, 1, 0], [0, 1, 0, 0], [0, 0, 0, 1]]).unsqueeze(0)
# camera = torch.matmul(flip_yz.to(camera), camera)
camera = camera[:, [0, 2, 1, 3]] # to blender convention (flip y & z axis)
camera[:, 1] *= -1
camera = normalize_camera(camera).view(batch_size, 16)
###############
# sampler = DDIMSampler(self.model)
# shape = [4, 32, 32]
# c_ = {"context": self.embeddings[4:]}
# uc_ = {"context": self.embeddings[:4]}
# # print(camera)
# # camera = get_camera(4, elevation=0, azimuth_start=0)
# # camera = camera.repeat(batch_size // 4, 1).to(self.device)
# # print(camera)
# c_["camera"] = uc_["camera"] = camera
# c_["num_frames"] = uc_["num_frames"] = 4
# latents_, _ = sampler.sample(S=30, conditioning=c_,
# batch_size=batch_size, shape=shape,
# verbose=False,
# unconditional_guidance_scale=guidance_scale,
# unconditional_conditioning=uc_,
# eta=0, x_T=None)
# # Img latents -> imgs
# imgs = self.decode_latents(latents_) # [4, 3, 256, 256]
# import kiui
# kiui.vis.plot_image(imgs)
###############
camera = camera.repeat(2, 1)
context = {"context": self.embeddings, "camera": camera, "num_frames": 4}
# predict the noise residual with unet, NO grad!
with torch.no_grad():
# add noise
noise = torch.randn_like(latents)
latents_noisy = self.model.q_sample(latents, t, noise)
# pred noise
latent_model_input = torch.cat([latents_noisy] * 2)
tt = torch.cat([t] * 2)
# import kiui
# kiui.lo(latent_model_input, t, context['context'], context['camera'])
noise_pred = self.model.apply_model(latent_model_input, tt, context)
# perform guidance (high scale from paper!)
noise_pred_uncond, noise_pred_pos = noise_pred.chunk(2)
noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_pos - noise_pred_uncond)
grad = (noise_pred - noise)
grad = torch.nan_to_num(grad)
# seems important to avoid NaN...
# grad = grad.clamp(-1, 1)
target = (latents - grad).detach()
loss = 0.5 * F.mse_loss(latents.float(), target, reduction='sum') / latents.shape[0]
return loss
def decode_latents(self, latents):
imgs = self.model.decode_first_stage(latents)
imgs = ((imgs + 1) / 2).clamp(0, 1)
return imgs
def encode_imgs(self, imgs):
# imgs: [B, 3, 256, 256]
imgs = 2 * imgs - 1
latents = self.model.get_first_stage_encoding(self.model.encode_first_stage(imgs))
return latents # [B, 4, 32, 32]
@torch.no_grad()
def prompt_to_img(
self,
prompts,
negative_prompts="",
height=256,
width=256,
num_inference_steps=50,
guidance_scale=7.5,
latents=None,
elevation=0,
azimuth_start=0,
):
if isinstance(prompts, str):
prompts = [prompts]
if isinstance(negative_prompts, str):
negative_prompts = [negative_prompts]
batch_size = len(prompts) * 4
# Text embeds -> img latents
sampler = DDIMSampler(self.model)
shape = [4, height // 8, width // 8]
c_ = {"context": self.encode_text(prompts).repeat(4,1,1)}
uc_ = {"context": self.encode_text(negative_prompts).repeat(4,1,1)}
camera = get_camera(4, elevation=elevation, azimuth_start=azimuth_start)
camera = camera.repeat(batch_size // 4, 1).to(self.device)
c_["camera"] = uc_["camera"] = camera
c_["num_frames"] = uc_["num_frames"] = 4
latents, _ = sampler.sample(S=num_inference_steps, conditioning=c_,
batch_size=batch_size, shape=shape,
verbose=False,
unconditional_guidance_scale=guidance_scale,
unconditional_conditioning=uc_,
eta=0, x_T=None)
# Img latents -> imgs
imgs = self.decode_latents(latents) # [4, 3, 256, 256]
# Img to Numpy
imgs = imgs.detach().cpu().permute(0, 2, 3, 1).numpy()
imgs = (imgs * 255).round().astype("uint8")
return imgs
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("--steps", type=int, default=30)
opt = parser.parse_args()
device = torch.device("cuda")
sd = MVDream(device)
while True:
imgs = sd.prompt_to_img(opt.prompt, opt.negative, num_inference_steps=opt.steps)
grid = np.concatenate([
np.concatenate([imgs[0], imgs[1]], axis=1),
np.concatenate([imgs[2], imgs[3]], axis=1),
], axis=0)
# visualize image
plt.imshow(grid)
plt.show()