과제4 stable-dreamfusion

image

DreamFusion (ICLR 2023 Oral) Text가 입력으로 주어졌을 때 3D를 생성할 수 있는 방법론입니다! 3D represntation은 Instant-NGP 사용(다양한 모델 활용 가능)합니다. 저는 NeRF를 사용했습니다. Loss function은 Score Distillation (SDS) loss를 사용합니다.

SDS Loss는 무엇일까요 ?

기존 diffusion models은 수식 전개를 통해, noise prediction 만으로 diffusion model 업데이트가 가능했습니다! (e.g., DDPM).

image

condition y를 입력으로 받아, conditioning input을 cross-attention layers에 적용해 이미지 생성 시 condition을 반영합니다(e.g., Stable Diffusion).

image

SDS loss는 pre-trained Stable Diffusion(구현에서는 Stable Diffusion을 사용하지만 아래 figure에서는 Imagen 모델 사용)을 활용하여 text prompt 만으로 3D 생성이 가능하게 하는 방법론이랍니다.

image

먼저 필요한 package와 dependency들을 설치해줍니다. DreamFusion의 오픈소스 버전인 Stable-Dreamfusion 코드를 설치해주고, diffusion을 활용하기 위한 추가적인 라이브러리(e.g., diffusers)도 설치해줍니다.

1. 클래스 초기화 (Initializing the Class with Super)

StableDiffusion 클래스의 __init__ 메소드에서는 부모 클래스의 생성자를 호출해야 한다. 이렇게 하는 이유는 상속된 속성을 초기화하고, 클래스가 올바르게 설정되도록 하기 위해서다.

# 클래스 초기화 코드
super(StableDiffusion, self).__init__()

2. min_step과 max_step 정의 (Defining min_step and max_step)

min_step과 max_step는 t_range와 self.num_train_timesteps에서 파생된 정수 값이다. t_range의 값에 self.num_train_timesteps를 곱한 후 정수로 변환해야 한다.

# min_step과 max_step 정의 코드
self.min_step = int(self.num_train_timesteps * t_range[0])
self.max_step = int(self.num_train_timesteps * t_range[1])

3. pred_rgb를 Latents로 보간하기 (Interpolating pred_rgb to Latents)

as_latent이 True일 때, pred_rgb 이미지를 (64, 64) 해상도로 리사이징하고, 값을 [-1, 1] 범위로 조정해서 바로 latent로 처리해야 한다.

# as_latent이 True일 때 pred_rgb를 latents로 보간하기
latents = F.interpolate(pred_rgb, size=(64, 64), mode='bilinear', align_corners=False)
latents = 2 * latents - 1

4. as_latent이 False일 때 이미지 인코딩 (Encoding Images when as_latent is False)

as_latent이 False일 때는 pred_rgb 이미지를 (512, 512) 해상도로 리사이징한 후, encode_imgs 메소드를 사용해 latents로 인코딩해야 한다.

# as_latent이 False일 때 pred_rgb 리사이징 및 인코딩
pred_rgb_512 = F.interpolate(pred_rgb, size=(512, 512), mode='bilinear', align_corners=False)
latents = self.encode_imgs(pred_rgb_512)

5. Timestep t 샘플링 (Sampling Timestep t)

min_step과 max_step 범위 내에서 무작위로 timestep t를 샘플링해야 한다. 이 작업은 torch.randint를 사용해 수행할 수 있다.

# Timestep t 샘플링 코드
t = torch.randint(self.min_step, self.max_step + 1, (1,), device=self.device)

6. Gradient 계산하기 (Calculating Gradient)

Gradient 계산을 위해, 예측된 노이즈와 실제 노이즈의 차이를 계산하고, 여기에 w 값을 곱해 최종 gradient를 구해야 한다.

# Gradient 계산 코드
grad = w * (noise_pred - noise)

7. SDS Loss를 위한 Targets 설정하기 (Setting Targets for SDS Loss)

targets 변수는 latents에서 예측된 노이즈를 뺀 값으로 설정해야 한다. 이 작업은 보통 noisy latents에서 노이즈 예측 값을 빼고, alpha 값으로 스케일링하여 수행된다.

# SDS Loss를 위한 targets 정의
targets = latents - grad

SDS loss를 활용해서 3D로 생성하고자 하는 text prompt를 입력하는데 나는 햄버거를 입력했다!

text = "a hamburger" # text prompt
workspace = "trial" # workspace, output이 저장되는 폴더명
epochs = 100  # 전체 학습 epoch

에폭은 100으로 설정했다. device는 cuda다 !

sd = StableDiffusion(device, opt.fp16, opt.vram_O, opt.sd_version, opt.hf_key)

H = 512  # image height
W = 512  # image width
steps = 50  # diffusion timestep
imgs = sd.prompt_to_img(opt.text, opt.negative, H, W, steps)

# visualize image
import matplotlib.pyplot as plt
plt.imshow(imgs[0])

image

이제 3D representation인 NeRF모델을 불러온다. SDS loss는 매번 diffusion을 inference 해야하기 때문에, 시간이 오래 걸린다. 나는 100에폭으로 돌려서 77분이 걸렸다! 디버깅 시에는 이보다 작은 epoch으로 진행하고, 코드 완성된 후 큰 epoch으로 실험해보면 좋다고한다.

엄청난 GPU 사용.. image

for i in [1, 50, 100]:
  for j in [1, 3, 5, 7]:
    images.append(f"./trial/validation/df_ep0{i:03d}_000{j}_rgb.png")

labels = ['Epoch 1', 'Epoch 50', 'Epoch 100']

image

depth rendering 결과도 visualize 해보면,

for i in [1, 50, 100]:
  for j in [1, 3, 5, 7]:
    images.append(f"./trial/validation/df_ep0{i:03d}_000{j}_depth.png")

image