Spaces:
Running
on
Zero
Running
on
Zero
File size: 1,421 Bytes
ea19453 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 |
import torch
from diffusers.configuration_utils import ConfigMixin, register_to_config
from diffusers import ModelMixin
from torch import Tensor
from .temporaltrans.temptrans import SimpleTemperalPointModel, SimpleTransModel
class PointModel(ModelMixin, ConfigMixin):
@register_to_config
def __init__(
self,
model_type: str = 'pvcnn',
in_channels: int = 3,
out_channels: int = 3,
embed_dim: int = 64,
dropout: float = 0.1,
width_multiplier: int = 1,
voxel_resolution_multiplier: int = 1,
):
super().__init__()
self.model_type = model_type
if self.model_type == 'simple':
self.autocast_context = torch.autocast('cuda', dtype=torch.float32)
self.model = SimpleTransModel(
embed_dim=embed_dim,
num_classes=out_channels,
extra_feature_channels=(in_channels - 3),
)
self.model.output_projection.bias.data.normal_(0, 1e-6)
self.model.output_projection.weight.data.normal_(0, 1e-6)
else:
raise NotImplementedError()
def forward(self, inputs: Tensor, t: Tensor, context=None) -> Tensor:
""" Receives input of shape (B, N, in_channels) and returns output
of shape (B, N, out_channels) """
with self.autocast_context:
return self.model(inputs, t, context)
|