|
{ |
|
"imports": [ |
|
"$import torch", |
|
"$from datetime import datetime", |
|
"$from pathlib import Path" |
|
], |
|
"bundle_root": ".", |
|
"model_dir": "$@bundle_root + '/models'", |
|
"output_dir": "$@bundle_root + '/output'", |
|
"create_output_dir": "$Path(@output_dir).mkdir(exist_ok=True)", |
|
"device": "$torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')", |
|
"output_postfix": "$datetime.now().strftime('sample_%Y%m%d_%H%M%S')", |
|
"spatial_dims": 3, |
|
"image_channels": 1, |
|
"latent_channels": 8, |
|
"latent_shape": [ |
|
8, |
|
36, |
|
44, |
|
28 |
|
], |
|
"autoencoder_def": { |
|
"_target_": "generative.networks.nets.AutoencoderKL", |
|
"spatial_dims": "@spatial_dims", |
|
"in_channels": "@image_channels", |
|
"out_channels": "@image_channels", |
|
"latent_channels": "@latent_channels", |
|
"num_channels": [ |
|
64, |
|
128, |
|
256 |
|
], |
|
"num_res_blocks": 2, |
|
"norm_num_groups": 32, |
|
"norm_eps": 1e-06, |
|
"attention_levels": [ |
|
false, |
|
false, |
|
false |
|
], |
|
"with_encoder_nonlocal_attn": false, |
|
"with_decoder_nonlocal_attn": false |
|
}, |
|
"network_def": { |
|
"_target_": "generative.networks.nets.DiffusionModelUNet", |
|
"spatial_dims": "@spatial_dims", |
|
"in_channels": "@latent_channels", |
|
"out_channels": "@latent_channels", |
|
"num_channels": [ |
|
256, |
|
256, |
|
512 |
|
], |
|
"attention_levels": [ |
|
false, |
|
true, |
|
true |
|
], |
|
"num_head_channels": [ |
|
0, |
|
64, |
|
64 |
|
], |
|
"num_res_blocks": 2 |
|
}, |
|
"load_autoencoder_path": "$@bundle_root + '/models/model_autoencoder.pt'", |
|
"load_autoencoder": "$@autoencoder_def.load_state_dict(torch.load(@load_autoencoder_path))", |
|
"autoencoder": "$@autoencoder_def.to(@device)", |
|
"load_diffusion_path": "$@model_dir + '/model.pt'", |
|
"load_diffusion": "$@network_def.load_state_dict(torch.load(@load_diffusion_path))", |
|
"diffusion": "$@network_def.to(@device)", |
|
"noise_scheduler": { |
|
"_target_": "generative.networks.schedulers.DDIMScheduler", |
|
"_requires_": [ |
|
"@load_diffusion", |
|
"@load_autoencoder" |
|
], |
|
"num_train_timesteps": 1000, |
|
"beta_start": 0.0015, |
|
"beta_end": 0.0195, |
|
"beta_schedule": "scaled_linear", |
|
"clip_sample": false |
|
}, |
|
"noise": "$torch.randn([1]+@latent_shape).to(@device)", |
|
"set_timesteps": "$@noise_scheduler.set_timesteps(num_inference_steps=50)", |
|
"inferer": { |
|
"_target_": "scripts.ldm_sampler.LDMSampler", |
|
"_requires_": "@set_timesteps" |
|
}, |
|
"sample": "$@inferer.sampling_fn(@noise, @autoencoder, @diffusion, @noise_scheduler)", |
|
"saver": { |
|
"_target_": "SaveImage", |
|
"_requires_": "@create_output_dir", |
|
"output_dir": "@output_dir", |
|
"output_postfix": "@output_postfix" |
|
}, |
|
"generated_image": "$@sample", |
|
"run": [ |
|
"$@saver(@generated_image[0])" |
|
] |
|
} |
|
|