Spaces:
Build error
Build error
File size: 1,212 Bytes
e78c13e |
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 40 41 42 43 44 45 46 47 |
# Copyright (c) Microsoft Corporation.
# Licensed under the MIT License.
import torch
def create_model(opt):
if opt.model == "pix2pixHD":
from .pix2pixHD_model import Pix2PixHDModel, InferenceModel
if opt.isTrain:
model = Pix2PixHDModel()
else:
model = InferenceModel()
else:
from .ui_model import UIModel
model = UIModel()
model.initialize(opt)
if opt.verbose:
print("model [%s] was created" % (model.name()))
if opt.isTrain and len(opt.gpu_ids) > 1:
# pass
model = torch.nn.DataParallel(model, device_ids=opt.gpu_ids)
return model
def create_da_model(opt):
if opt.model == 'pix2pixHD':
from .pix2pixHD_model_DA import Pix2PixHDModel, InferenceModel
if opt.isTrain:
model = Pix2PixHDModel()
else:
model = InferenceModel()
else:
from .ui_model import UIModel
model = UIModel()
model.initialize(opt)
if opt.verbose:
print("model [%s] was created" % (model.name()))
if opt.isTrain and len(opt.gpu_ids) > 1:
#pass
model = torch.nn.DataParallel(model, device_ids=opt.gpu_ids)
return model |