cnywt/SyncTalk

1def get_encoder(encoding, input_dim=3,2                multires=6, 3                degree=4,4                num_levels=16, level_dim=2, base_resolution=16, log2_hashmap_size=19, desired_resolution=2048, align_corners=False,5                **kwargs):6 7    if encoding == 'None':8        return lambda x, **kwargs: x, input_dim9    10    elif encoding == 'frequency':11        from freqencoder import FreqEncoder12        encoder = FreqEncoder(input_dim=input_dim, degree=multires)13 14    elif encoding == 'spherical_harmonics':15        from shencoder import SHEncoder16        encoder = SHEncoder(input_dim=input_dim, degree=degree)17 18    elif encoding == 'hashgrid':19        from gridencoder import GridEncoder20        encoder = GridEncoder(input_dim=input_dim, num_levels=num_levels, level_dim=level_dim, base_resolution=base_resolution, log2_hashmap_size=log2_hashmap_size, desired_resolution=desired_resolution, gridtype='hash', align_corners=align_corners)21    22    elif encoding == 'tiledgrid':23        from gridencoder import GridEncoder24        encoder = GridEncoder(input_dim=input_dim, num_levels=num_levels, level_dim=level_dim, base_resolution=base_resolution, log2_hashmap_size=log2_hashmap_size, desired_resolution=desired_resolution, gridtype='tiled', align_corners=align_corners)25    26    elif encoding == 'ash':27        from ashencoder import AshEncoder28        encoder = AshEncoder(input_dim=input_dim, output_dim=16, log2_hashmap_size=log2_hashmap_size, resolution=desired_resolution)29 30    else:31        raise NotImplementedError('Unknown encoding mode, choose from [None, frequency, spherical_harmonics, hashgrid, tiledgrid]')32 33    return encoder, encoder.output_dim