configs/inference.json

{
    "imports": [
        "$import torch",
        "$from datetime import datetime",
        "$from pathlib import Path",
        "$from transformers import CLIPTextModel",
        "$from transformers import CLIPTokenizer"
    ],
    "bundle_root": ".",
    "dataset_dir": "",
    "dataset": "",
    "evaluator": "",
    "inferer": "",
    "load_old": 1,
    "model_dir": "$@bundle_root + '/models'",
    "output_dir": "$@bundle_root + '/output'",
    "create_output_dir": "$Path(@output_dir).mkdir(exist_ok=True)",
    "prompt": "Big right-sided pleural effusion",
    "prompt_list": "$['', @prompt]",
    "guidance_scale": 7.0,
    "device": "$torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')",
    "tokenizer": "$CLIPTokenizer.from_pretrained(\"stabilityai/stable-diffusion-2-1-base\", subfolder=\"tokenizer\")",
    "text_encoder": "$CLIPTextModel.from_pretrained(\"stabilityai/stable-diffusion-2-1-base\", subfolder=\"text_encoder\")",
    "tokenized_prompt": "$@tokenizer(@prompt_list, padding=\"max_length\", max_length=@tokenizer.model_max_length, truncation=True,return_tensors=\"pt\")",
    "prompt_embeds": "$@text_encoder(@tokenized_prompt.input_ids.squeeze(1))[0].to(@device)",
    "out_file": "$datetime.now().strftime('sample_%H%M%S_%d%m%Y')",
    "autoencoder_def": {
        "_target_": "monai.networks.nets.AutoencoderKL",
        "spatial_dims": 2,
        "in_channels": 1,
        "out_channels": 1,
        "latent_channels": 3,
        "channels": [
            64,
            128,
            128,
            128
        ],
        "num_res_blocks": 2,
        "norm_num_groups": 32,
        "norm_eps": 1e-06,
        "attention_levels": [
            false,
            false,
            false,
            false
        ],
        "with_encoder_nonlocal_attn": false,
        "with_decoder_nonlocal_attn": false
    },
    "network_def": "@diffusion_def",
    "load_autoencoder_path": "$@model_dir + '/autoencoder.pt'",
    "load_autoencoder_func": "$@autoencoder_def.load_old_state_dict if bool(@load_old) else @autoencoder_def.load_state_dict",
    "load_autoencoder": "$@load_autoencoder_func(torch.load(@load_autoencoder_path))",
    "autoencoder": "$@autoencoder_def.to(@device)",
    "diffusion_def": {
        "_target_": "monai.networks.nets.DiffusionModelUNet",
        "spatial_dims": 2,
        "in_channels": 3,
        "out_channels": 3,
        "channels": [
            256,
            512,
            768
        ],
        "num_res_blocks": 2,
        "attention_levels": [
            false,
            true,
            true
        ],
        "norm_num_groups": 32,
        "norm_eps": 1e-06,
        "resblock_updown": false,
        "num_head_channels": [
            0,
            512,
            768
        ],
        "with_conditioning": true,
        "transformer_num_layers": 1,
        "cross_attention_dim": 1024
    },
    "load_diffusion_path": "$@model_dir + '/model.pt'",
    "load_diffusion_func": "$@diffusion_def.load_old_state_dict if bool(@load_old) else @diffusion_def.load_state_dict",
    "load_diffusion": "$@load_diffusion_func(torch.load(@load_diffusion_path))",
    "diffusion": "$@diffusion_def.to(@device)",
    "scheduler": {
        "_target_": "monai.networks.schedulers.DDIMScheduler",
        "_requires_": [
            "@load_diffusion",
            "@load_autoencoder"
        ],
        "beta_start": 0.0015,
        "beta_end": 0.0205,
        "num_train_timesteps": 1000,
        "schedule": "scaled_linear_beta",
        "prediction_type": "v_prediction",
        "clip_sample": false
    },
    "noise": "$torch.randn((1, 3, 64, 64)).to(@device)",
    "set_timesteps": "$@scheduler.set_timesteps(num_inference_steps=50)",
    "sampler": {
        "_target_": "scripts.sampler.Sampler",
        "_requires_": "@set_timesteps"
    },
    "sample": "$@sampler.sampling_fn(@noise, @autoencoder, @diffusion, @scheduler, @prompt_embeds)",
    "saver": {
        "_target_": "scripts.saver.JPGSaver",
        "_requires_": "@create_output_dir",
        "output_dir": "@output_dir"
    },
    "run": "$@saver.save(@sample, @out_file)",
    "save": "$torch.save(@sample, @output_dir + '/' + @out_file + '.pt')"
}