update pyramid
This commit is contained in:
@@ -10,7 +10,7 @@ from custom_utils.custom_utils import load_model
|
||||
from models.base import AEEncoder
|
||||
from models.cvae_model import CVAE
|
||||
from models.pyramid_level import PyramidLevel
|
||||
from models.upscaler import UpscalerBasic, UpscalerBasicNew
|
||||
from models.upscaler import UpscalerBasic, UpscalerBasicConv, UpscalerBasicNew
|
||||
|
||||
|
||||
class Pyramid(tf.keras.Model):
|
||||
@@ -22,7 +22,7 @@ class Pyramid(tf.keras.Model):
|
||||
if config['pyramid']['static_model'] != '':
|
||||
load_model(self.static_enc, config['pyramid']['static_model'] + '_enc')
|
||||
|
||||
self.body_enc : AEEncoder = AEEncoder(config['cond_ae']['latent_dim'], config['cond_ae']['encoder_depths'], config['cond_ae']['encoder_dims'], config['cond_ae']['bottleneck_size'], config['cond_ae']['kernel_size'])
|
||||
self.body_enc : AEEncoder = AEEncoder(config['cond_ae_body']['latent_dim'], config['cond_ae_body']['encoder_depths'], config['cond_ae_body']['encoder_dims'], config['cond_ae_body']['bottleneck_size'], config['cond_ae_body']['kernel_size'])
|
||||
if config['pyramid']['body_model'] != '':
|
||||
load_model(self.body_enc, config['pyramid']['body_model'] + '_enc')
|
||||
|
||||
@@ -45,7 +45,7 @@ class Pyramid(tf.keras.Model):
|
||||
|
||||
self.upscale_enhancer = None
|
||||
if config['pyramid']['final_upscale_enhancer'] != '':
|
||||
self.upscale_enhancer = UpscalerBasicNew(low_res=config['pyramid']['resolutions'][-1])
|
||||
self.upscale_enhancer = UpscalerBasicConv(low_res=config['pyramid']['resolutions'][-1])
|
||||
load_model(self.upscale_enhancer, config['pyramid']['final_upscale_enhancer'])
|
||||
self.upscale_enhancer.trainable = False
|
||||
|
||||
|
||||
Reference in New Issue
Block a user