mirror of
https://github.com/nerfstudio-project/nerfstudio.git
synced 2026-11-11 18:44:21 +08:00
Cleanup SD code
This commit is contained in:
Matt Tancik
parent
a45c6ca962
commit
ab1301ac3c
@ -16,9 +16,9 @@
|
||||
|
||||
import sys
|
||||
from pathlib import Path
|
||||
from typing import List, Optional, Union
|
||||
from typing import List, Optional, Tuple, Union
|
||||
|
||||
import matplotlib.pyplot as plt
|
||||
import mediapy
|
||||
import numpy as np
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
@ -48,14 +48,20 @@ CLIP_SOURCE = "openai/clip-vit-large-patch14"
|
||||
|
||||
|
||||
class StableDiffusion(nn.Module):
|
||||
"""Stable Diffusion implementation"""
|
||||
"""Stable Diffusion implementation
|
||||
|
||||
def __init__(self, device) -> None:
|
||||
Args:
|
||||
device: device to use
|
||||
num_train_timesteps: number of training timesteps
|
||||
|
||||
"""
|
||||
|
||||
def __init__(self, device: Union[torch.device, str], num_train_timesteps: int = 1000) -> None:
|
||||
super().__init__()
|
||||
|
||||
self.device = device
|
||||
self.num_train_timesteps = num_train_timesteps
|
||||
|
||||
self.num_train_timesteps = 1000
|
||||
self.min_step = int(self.num_train_timesteps * 0.02)
|
||||
self.max_step = int(self.num_train_timesteps * 0.98)
|
||||
|
||||
@ -65,7 +71,7 @@ class StableDiffusion(nn.Module):
|
||||
beta_schedule="scaled_linear",
|
||||
num_train_timesteps=self.num_train_timesteps,
|
||||
)
|
||||
self.alphas = self.scheduler.alphas_cumprod.to(self.device)
|
||||
self.alphas = self.scheduler.alphas_cumprod.to(self.device) # type: ignore
|
||||
|
||||
pipe = StableDiffusionPipeline.from_pretrained(SD_SOURCE, torch_dtype=torch.float16)
|
||||
assert pipe is not None
|
||||
@ -127,65 +133,82 @@ class StableDiffusion(nn.Module):
|
||||
|
||||
return text_embeddings
|
||||
|
||||
def sds_loss(self, text_embeddings, nerf_output, guidance_scale=100):
|
||||
nerf_output_np = nerf_output.squeeze(0).permute(1, 2, 0).detach().cpu().numpy()
|
||||
nerf_output_np = np.clip(nerf_output_np, 0.0, 1.0)
|
||||
plt.imsave("nerf_output.png", nerf_output_np)
|
||||
nerf_output = F.interpolate(nerf_output, (IMG_DIM, IMG_DIM), mode="bilinear")
|
||||
def sds_loss(
|
||||
self,
|
||||
text_embeddings: TensorType["N", "max_length", "embed_dim"],
|
||||
image: TensorType["BS", 3, "H", "W"],
|
||||
guidance_scale: float = 100.0,
|
||||
) -> Tuple[torch.Tensor, TensorType["BS", "2N", "H", "W"], TensorType["BS", "2N", "H", "W"]]:
|
||||
"""Score Distilation Sampling loss proposed in DreamFusion paper (https://dreamfusion3d.github.io/)
|
||||
|
||||
Args:
|
||||
text_embeddings: Text embeddings
|
||||
image: Rendered image
|
||||
guidance_scale: How much to weigh the guidance
|
||||
|
||||
Returns:
|
||||
A tuple of Loss, latent, and gradient
|
||||
"""
|
||||
image = F.interpolate(image, (IMG_DIM, IMG_DIM), mode="bilinear")
|
||||
t = torch.randint(self.min_step, self.max_step + 1, [1], dtype=torch.long, device=self.device)
|
||||
latents = self.imgs_to_latent(nerf_output)
|
||||
latents = self.imgs_to_latent(image)
|
||||
|
||||
# predict the noise residual with unet, NO grad!
|
||||
with torch.no_grad():
|
||||
# add noise
|
||||
noise = torch.randn_like(latents)
|
||||
latents_noisy = self.scheduler.add_noise(latents, noise, t)
|
||||
latents_noisy = self.scheduler.add_noise(latents, noise, t) # type: ignore
|
||||
# pred noise
|
||||
latent_model_input = torch.cat([latents_noisy] * 2)
|
||||
noise_pred = self.unet(latent_model_input, t, encoder_hidden_states=text_embeddings).sample
|
||||
|
||||
# perform guidance (high scale from paper!)
|
||||
# perform guidance
|
||||
noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)
|
||||
noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond)
|
||||
# with torch.no_grad():
|
||||
# diffused_img = self.latents_to_img(latents_noisy - noise_pred)
|
||||
# diffused_img = diffused_img.detach().cpu().permute(0, 2, 3, 1).numpy().reshape((512, 512, 3))
|
||||
# diffused_img = (diffused_img * 255).round().astype("uint8")
|
||||
# plt.imsave("sdimg.png", diffused_img)
|
||||
|
||||
# w(t), sigma_t^2
|
||||
w = 1 - self.alphas[t]
|
||||
# w = self.alphas[t] ** 0.5 * (1 - self.alphas[t])
|
||||
grad = w * (noise_pred - noise)
|
||||
|
||||
# clip grad for stable training?
|
||||
# grad = grad.clamp(-10, 10)
|
||||
grad = w * (noise_pred - noise)
|
||||
grad = torch.nan_to_num(grad)
|
||||
|
||||
# noise_difference = self.latents_to_img(grad)
|
||||
# plt.imsave('noise_difference.png', noise_difference.view(3, 512, 512).permute(1, 2, 0).detach().cpu().numpy())
|
||||
|
||||
# manually backward, since we omitted an item in grad and cannot simply autodiff.
|
||||
# _t = time.time()
|
||||
|
||||
noise_loss = torch.mean(torch.nan_to_num(torch.square(noise_pred - noise)))
|
||||
|
||||
# print('SDS LOSS:', noise_loss)
|
||||
|
||||
return noise_loss, latents, grad
|
||||
|
||||
def produce_latents(
|
||||
self, text_embeddings, height=IMG_DIM, width=IMG_DIM, num_inference_steps=50, guidance_scale=7.5, latents=None
|
||||
):
|
||||
self,
|
||||
text_embeddings: TensorType["N", "max_length", "embed_dim"],
|
||||
height: int = IMG_DIM,
|
||||
width: int = IMG_DIM,
|
||||
num_inference_steps: int = 50,
|
||||
guidance_scale: float = 7.5,
|
||||
latents: Optional[TensorType["BS", "2N", "H", "W"]] = None,
|
||||
) -> TensorType["BS", "2N", "H", "W"]:
|
||||
"""Produce latents for a given text embedding
|
||||
|
||||
Args:
|
||||
text_embeddings: Text embeddings
|
||||
height: Height of the image
|
||||
width: Width of the image
|
||||
num_inference_steps: Number of inference steps
|
||||
guidance_scale: How much to weigh the guidance
|
||||
latents: Latents to start with
|
||||
|
||||
Returns:
|
||||
Latents
|
||||
"""
|
||||
|
||||
if latents is None:
|
||||
latents = torch.randn(
|
||||
(text_embeddings.shape[0] // 2, self.unet.in_channels, height // 8, width // 8), device=self.device
|
||||
)
|
||||
|
||||
self.scheduler.set_timesteps(num_inference_steps)
|
||||
self.scheduler.set_timesteps(num_inference_steps) # type: ignore
|
||||
|
||||
with torch.autocast("cuda"):
|
||||
for t in self.scheduler.timesteps:
|
||||
for t in self.scheduler.timesteps: # type: ignore
|
||||
assert latents is not None
|
||||
# expand the latents if we are doing classifier-free guidance to avoid doing two forward passes.
|
||||
latent_model_input = torch.cat([latents] * 2)
|
||||
|
||||
@ -198,11 +221,18 @@ class StableDiffusion(nn.Module):
|
||||
noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond)
|
||||
|
||||
# compute the previous noisy sample x_t -> x_t-1
|
||||
latents = self.scheduler.step(noise_pred, t, latents)["prev_sample"]
|
||||
|
||||
latents = self.scheduler.step(noise_pred, t, latents)["prev_sample"] # type: ignore
|
||||
return latents
|
||||
|
||||
def latents_to_img(self, latents):
|
||||
def latents_to_img(self, latents: TensorType["BS", "2N", "H", "W"]) -> TensorType["BS", 3, "H", "W"]:
|
||||
"""Convert latents to images
|
||||
|
||||
Args:
|
||||
latents: Latents to convert
|
||||
|
||||
Returns:
|
||||
Images
|
||||
"""
|
||||
|
||||
latents = 1 / CONST_SCALE * latents
|
||||
|
||||
@ -213,8 +243,15 @@ class StableDiffusion(nn.Module):
|
||||
|
||||
return imgs
|
||||
|
||||
def imgs_to_latent(self, imgs):
|
||||
# imgs: [B, 3, H, W]
|
||||
def imgs_to_latent(self, imgs: TensorType["BS", 3, "H", "W"]) -> TensorType["BS", "2N", "H", "W"]:
|
||||
"""Convert images to latents
|
||||
|
||||
Args:
|
||||
imgs: Images to convert
|
||||
|
||||
Returns:
|
||||
Latents
|
||||
"""
|
||||
imgs = 2 * imgs - 1
|
||||
|
||||
posterior = self.auto_encoder.encode(imgs).latent_dist
|
||||
@ -297,7 +334,7 @@ def generate_image(
|
||||
with torch.no_grad():
|
||||
sd = StableDiffusion(cuda_device)
|
||||
imgs = sd.prompt_to_img(prompt, negative, steps)
|
||||
plt.imsave(str(save_path), imgs[0])
|
||||
mediapy.write_image(str(save_path), imgs[0])
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
|
||||
Loading…
Reference in New Issue
Block a user