File size: 2,043 Bytes
bf8981a |
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 48 49 50 51 52 53 54 55 56 57 58 59 60 |
from .latent_diffusion import LatentDiffusion
from .model_sdf import Diffpro_SDF
from .architecture.unet import UNetModel
import os
def init_ldm_model(params, debug_mode=False):
unet_model = UNetModel(
in_channels=params.in_channels,
out_channels=params.out_channels,
channels=params.channels,
attention_levels=params.attention_levels,
n_res_blocks=params.n_res_blocks,
channel_multipliers=params.channel_multipliers,
n_heads=params.n_heads,
tf_layers=params.tf_layers,
#d_cond=params.d_cond,
)
ldm_model = LatentDiffusion(
unet_model=unet_model,
#autoencoder=None,
#autoreg_cond_enc=autoreg_cond_enc,
#external_cond_enc=external_cond_enc,
latent_scaling_factor=params.latent_scaling_factor,
n_steps=params.n_steps,
linear_start=params.linear_start,
linear_end=params.linear_end,
debug_mode=debug_mode
)
return ldm_model
def init_diff_pro_sdf(ldm_model, params, device):
return Diffpro_SDF(ldm_model).to(device)
def get_model_path(model_dir, model_id=None):
model_desc = os.path.basename(model_dir)
if model_id is None:
model_path = os.path.join(model_dir, 'chkpts', 'weights.pt')
# retrieve real model_id from the actual file weights.pt is pointing to
model_id = os.path.basename(os.path.realpath(model_path)).split('-')[1].split('.')[0]
elif model_id == 'best':
model_path = os.path.join(model_dir, 'chkpts', 'weights_best.pt')
# retrieve real model_id from the actual file weights.pt is pointing to
model_id = os.path.basename(os.path.realpath(model_path)).split('-')[1].split('.')[0]
elif model_id == 'default':
model_path = os.path.join(model_dir, 'chkpts', 'weights_default.pt')
if not os.path.exists(model_path):
return get_model_path(model_dir, 'best')
else:
model_path = f"{model_dir}/chkpts/weights-{model_id}.pt"
return model_path, model_id, model_desc
|