lama

1
import logging
2

3
from saicinpainting.training.modules.ffc import FFCResNetGenerator
4
from saicinpainting.training.modules.pix2pixhd import GlobalGenerator, MultiDilatedGlobalGenerator, \
5
    NLayerDiscriminator, MultidilatedNLayerDiscriminator
6

7
def make_generator(config, kind, **kwargs):
8
    logging.info(f'Make generator {kind}')
9

10
    if kind == 'pix2pixhd_multidilated':
11
        return MultiDilatedGlobalGenerator(**kwargs)
12
    
13
    if kind == 'pix2pixhd_global':
14
        return GlobalGenerator(**kwargs)
15

16
    if kind == 'ffc_resnet':
17
        return FFCResNetGenerator(**kwargs)
18

19
    raise ValueError(f'Unknown generator kind {kind}')
20

21

22
def make_discriminator(kind, **kwargs):
23
    logging.info(f'Make discriminator {kind}')
24

25
    if kind == 'pix2pixhd_nlayer_multidilated':
26
        return MultidilatedNLayerDiscriminator(**kwargs)
27

28
    if kind == 'pix2pixhd_nlayer':
29
        return NLayerDiscriminator(**kwargs)
30

31
    raise ValueError(f'Unknown discriminator kind {kind}')
32