Spaces:
Running
on
T4
Running
on
T4
| import ipdb # noqa: F401 | |
| import numpy as np | |
| import torch | |
| import torch.nn as nn | |
| from diffusionsfm.model.dit import DiT | |
| from diffusionsfm.model.feature_extractors import PretrainedVAE, SpatialDino | |
| from diffusionsfm.model.scheduler import NoiseScheduler | |
| class RayDiffuser(nn.Module): | |
| def __init__( | |
| self, | |
| model_type="dit", | |
| depth=8, | |
| width=16, | |
| hidden_size=1152, | |
| P=1, | |
| max_num_images=1, | |
| noise_scheduler=None, | |
| freeze_encoder=True, | |
| feature_extractor="dino", | |
| append_ndc=True, | |
| use_unconditional=False, | |
| diffuse_depths=False, | |
| depth_resolution=1, | |
| use_homogeneous=False, | |
| cond_depth_mask=False, | |
| ): | |
| super().__init__() | |
| if noise_scheduler is None: | |
| self.noise_scheduler = NoiseScheduler() | |
| else: | |
| self.noise_scheduler = noise_scheduler | |
| self.diffuse_depths = diffuse_depths | |
| self.depth_resolution = depth_resolution | |
| self.use_homogeneous = use_homogeneous | |
| self.ray_dim = 3 | |
| if self.use_homogeneous: | |
| self.ray_dim += 1 | |
| self.ray_dim += self.ray_dim * self.depth_resolution**2 | |
| if self.diffuse_depths: | |
| self.ray_dim += 1 | |
| self.append_ndc = append_ndc | |
| self.width = width | |
| self.max_num_images = max_num_images | |
| self.model_type = model_type | |
| self.use_unconditional = use_unconditional | |
| self.cond_depth_mask = cond_depth_mask | |
| if feature_extractor == "dino": | |
| self.feature_extractor = SpatialDino( | |
| freeze_weights=freeze_encoder, num_patches_x=width, num_patches_y=width | |
| ) | |
| self.feature_dim = self.feature_extractor.feature_dim | |
| elif feature_extractor == "vae": | |
| self.feature_extractor = PretrainedVAE( | |
| freeze_weights=freeze_encoder, num_patches_x=width, num_patches_y=width | |
| ) | |
| self.feature_dim = self.feature_extractor.feature_dim | |
| else: | |
| raise Exception(f"Unknown feature extractor {feature_extractor}") | |
| if self.use_unconditional: | |
| self.register_parameter( | |
| "null_token", nn.Parameter(torch.randn(self.feature_dim, 1, 1)) | |
| ) | |
| self.input_dim = self.feature_dim * 2 | |
| if self.append_ndc: | |
| self.input_dim += 2 | |
| if model_type == "dit": | |
| self.ray_predictor = DiT( | |
| in_channels=self.input_dim, | |
| out_channels=self.ray_dim, | |
| width=width, | |
| depth=depth, | |
| hidden_size=hidden_size, | |
| max_num_images=max_num_images, | |
| P=P, | |
| ) | |
| self.scratch = nn.Module() | |
| self.scratch.input_conv = nn.Linear(self.ray_dim + int(self.cond_depth_mask), self.feature_dim) | |
| def forward_noise( | |
| self, x, t, epsilon=None, zero_out_mask=None | |
| ): | |
| """ | |
| Applies forward diffusion (adds noise) to the input. | |
| If a mask is provided, the noise is only applied to the masked inputs. | |
| """ | |
| t = t.reshape(-1, 1, 1, 1, 1) | |
| if epsilon is None: | |
| epsilon = torch.randn_like(x) | |
| else: | |
| epsilon = epsilon.reshape(x.shape) | |
| alpha_bar = self.noise_scheduler.alphas_cumprod[t] | |
| x_noise = torch.sqrt(alpha_bar) * x + torch.sqrt(1 - alpha_bar) * epsilon | |
| if zero_out_mask is not None and self.cond_depth_mask: | |
| x_noise = x_noise * zero_out_mask | |
| return x_noise, epsilon | |
| def forward( | |
| self, | |
| features=None, | |
| images=None, | |
| rays=None, | |
| rays_noisy=None, | |
| t=None, | |
| ndc_coordinates=None, | |
| unconditional_mask=None, | |
| return_dpt_activations=False, | |
| depth_mask=None, | |
| ): | |
| """ | |
| Args: | |
| images: (B, N, 3, H, W). | |
| t: (B,). | |
| rays: (B, N, 6, H, W). | |
| rays_noisy: (B, N, 6, H, W). | |
| ndc_coordinates: (B, N, 2, H, W). | |
| unconditional_mask: (B, N) or (B,). Should be 1 for unconditional samples | |
| and 0 else. | |
| """ | |
| if features is None: | |
| # VAE expects 256x256 images while DINO expects 224x224 images. | |
| # Both feature extractors support autoresize=True, but ideally we should | |
| # set this to be false and handle in the dataloader. | |
| features = self.feature_extractor(images, autoresize=True) | |
| B = features.shape[0] | |
| if ( | |
| unconditional_mask is not None | |
| and self.use_unconditional | |
| ): | |
| null_token = self.null_token.reshape(1, 1, self.feature_dim, 1, 1) | |
| unconditional_mask = unconditional_mask.reshape(B, -1, 1, 1, 1) | |
| features = ( | |
| features * (1 - unconditional_mask) + null_token * unconditional_mask | |
| ) | |
| if isinstance(t, int) or isinstance(t, np.int64): | |
| t = torch.ones(1, dtype=int).to(features.device) * t | |
| else: | |
| t = t.reshape(B) | |
| if rays_noisy is None: | |
| if self.cond_depth_mask: | |
| rays_noisy, epsilon = self.forward_noise(rays, t, zero_out_mask=depth_mask.unsqueeze(2)) | |
| else: | |
| rays_noisy, epsilon = self.forward_noise(rays, t) | |
| else: | |
| epsilon = None | |
| if self.cond_depth_mask: | |
| if depth_mask is None: | |
| depth_mask = torch.ones_like(rays_noisy[:, :, 0]) | |
| ray_repr = torch.cat([rays_noisy, depth_mask.unsqueeze(2)], dim=2) | |
| else: | |
| ray_repr = rays_noisy | |
| ray_repr = ray_repr.permute(0, 1, 3, 4, 2) | |
| ray_repr = self.scratch.input_conv(ray_repr).permute(0, 1, 4, 2, 3).contiguous() | |
| scene_features = torch.cat([features, ray_repr], dim=2) | |
| if self.append_ndc: | |
| scene_features = torch.cat([scene_features, ndc_coordinates], dim=2) | |
| epsilon_pred = self.ray_predictor( | |
| scene_features, | |
| t, | |
| return_dpt_activations=return_dpt_activations, | |
| ) | |
| if return_dpt_activations: | |
| return epsilon_pred, rays_noisy, epsilon | |
| return epsilon_pred, epsilon | |