add no curvature option

2023-01-26 12:03:50 +00:00
parent 2ac1a812fc
commit 44e1d02e5a
4 changed files with 32 additions and 19 deletions
@@ -803,12 +803,14 @@ if __name__ == '__main__':
    if len(config['pyramid']['resolutions']) == 1:
        model.static_enc.trainable = False
        model.body_enc.trainable = False
-        model.curv_enc.trainable = False
+        if config['use_curvature']:
+            model.curv_enc.trainable = False
    else:
        load_model(model, config['pyramid']['pretrained_pyramid'], iteration=(config['pyramid']['pretrained_pyramid_iteration'] if 'pretrained_pyramid_iteration' in config.keys() else ''))
        model.static_enc.trainable = False
        model.body_enc.trainable = False
-        model.curv_enc.trainable = False
+        if config['use_curvature']:
+            model.curv_enc.trainable = False
        model.cond_encoder.trainable = False

    # load_model(model.levels[0].upscale_enhancer, config['pyramid']['levels'][0]['upscale_enhancer'])