Spaces:
Running on Zero
Running on Zero
multimodalart HF Staff
Learn2Splat ZeroGPU demo: app, vendored optgs, prebuilt sm_120 wheels
9196862 verified | from torch import nn as nn | |
| from torchvision.models import resnet18 | |
| class BasicBlock(nn.Module): | |
| expansion = 1 | |
| def __init__(self, in_channels, out_channels, stride=1, downsample=None, num_groups=8): | |
| super(BasicBlock, self).__init__() | |
| self.conv1 = nn.Conv2d(in_channels, out_channels, kernel_size=3, | |
| stride=stride, padding=1, bias=False) | |
| self.gn1 = nn.GroupNorm(num_groups, out_channels) | |
| self.conv2 = nn.Conv2d(out_channels, out_channels, kernel_size=3, | |
| stride=1, padding=1, bias=False) | |
| self.gn2 = nn.GroupNorm(num_groups, out_channels) | |
| self.downsample = downsample | |
| self.gelu = nn.GELU() | |
| def forward(self, x): | |
| identity = x | |
| out = self.gelu(self.gn1(self.conv1(x))) | |
| out = self.gn2(self.conv2(out)) | |
| if self.downsample is not None: | |
| identity = self.downsample(x) | |
| out += identity | |
| out = self.gelu(out) | |
| return out | |
| class ResNetFeatureWarpper(nn.Module): | |
| def __init__(self, shallow_resnet_feature=False): | |
| super(ResNetFeatureWarpper, self).__init__() | |
| self.shallow_resnet_feature = shallow_resnet_feature | |
| resnet = resnet18(pretrained=True) | |
| self.conv1 = resnet.conv1 | |
| self.bn1 = resnet.bn1 | |
| self.relu = resnet.relu | |
| self.maxpool = resnet.maxpool | |
| self.layer1 = resnet.layer1 | |
| if not shallow_resnet_feature: | |
| self.layer2 = resnet.layer2 | |
| def forward(self, x): | |
| out = [] | |
| x = self.conv1(x) | |
| out.append(x) | |
| x = self.bn1(x) | |
| x = self.relu(x) | |
| x = self.maxpool(x) | |
| x = self.layer1(x) | |
| out.append(x) | |
| if not self.shallow_resnet_feature: | |
| x = self.layer2(x) | |
| out.append(x) | |
| return out | |