{ "cells": [ { "cell_type": "markdown", "id": "05fc7b5c", "metadata": {}, "source": [ "Copyright (c) MONAI Consortium \n", "Licensed under the Apache License, Version 2.0 (the \"License\"); \n", "you may not use this file except in compliance with the License. \n", "You may obtain a copy of the License at \n", "    http://www.apache.org/licenses/LICENSE-2.0 \n", "Unless required by applicable law or agreed to in writing, software \n", "distributed under the License is distributed on an \"AS IS\" BASIS, \n", "WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. \n", "See the License for the specific language governing permissions and \n", "limitations under the License." ] }, { "cell_type": "markdown", "id": "777b7dcb", "metadata": {}, "source": [ "# Training a 3D Diffusion Model for Generating 3D Images with Various Sizes and Spacings\n", "\n", "![Generated image examples](https://developer-blogs.nvidia.com/wp-content/uploads/2024/06/image3.png)\n", "\n", "In this notebook, we detail the procedure for training a 3D latent diffusion model to generate high-dimensional 3D medical images. Due to the potential for out-of-memory issues on most GPUs when generating large images (e.g., those with dimensions of 512 x 512 x 512 or greater), we have structured the training process into two primary steps: 1) generating image embeddings and 2) training 3D latent diffusion models. The subsequent sections will demonstrate the entire process using a simulated dataset.\n", "\n", "`[Release Note (March 2025)]:` We are excited to announce the new MAISI Version `'maisi3d-rflow'`. Compared with the previous version `'maisi3d-ddpm'`, it accelerated latent diffusion model inference by 33x. Please see the detailed difference in the following section." ] }, { "cell_type": "markdown", "id": "c9ecfb90", "metadata": {}, "source": [ "## Setup environment" ] }, { "cell_type": "code", "execution_count": 2, "id": "58cbde9b", "metadata": {}, "outputs": [], "source": [ "!python -c \"import monai\" || pip install -q \"monai-weekly[pillow, tqdm]\"" ] }, { "cell_type": "markdown", "id": "d655b95c", "metadata": {}, "source": [ "## Setup imports" ] }, { "cell_type": "code", "execution_count": 5, "id": "e3bf0346", "metadata": {}, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "MONAI version: 1.6.dev2537\n", "Numpy version: 2.2.6\n", "Pytorch version: 2.8.0+cu128\n", "MONAI flags: HAS_EXT = False, USE_COMPILED = False, USE_META_DICT = False\n", "MONAI rev id: ca9fa44c33f3b7c073263938e7a70b2fa6123044\n", "MONAI __file__: /home//miniconda3/envs/monai/lib/python3.10/site-packages/monai/__init__.py\n", "\n", "Optional dependencies:\n", "Pytorch Ignite version: NOT INSTALLED or UNKNOWN VERSION.\n", "ITK version: NOT INSTALLED or UNKNOWN VERSION.\n", "Nibabel version: 5.3.2\n", "scikit-image version: 0.25.2\n", "scipy version: 1.15.3\n", "Pillow version: 11.3.0\n", "Tensorboard version: NOT INSTALLED or UNKNOWN VERSION.\n", "gdown version: NOT INSTALLED or UNKNOWN VERSION.\n", "TorchVision version: NOT INSTALLED or UNKNOWN VERSION.\n", "tqdm version: 4.67.1\n", "lmdb version: NOT INSTALLED or UNKNOWN VERSION.\n", "psutil version: 5.9.1\n", "pandas version: NOT INSTALLED or UNKNOWN VERSION.\n", "einops version: NOT INSTALLED or UNKNOWN VERSION.\n", "transformers version: NOT INSTALLED or UNKNOWN VERSION.\n", "mlflow version: NOT INSTALLED or UNKNOWN VERSION.\n", "pynrrd version: NOT INSTALLED or UNKNOWN VERSION.\n", "clearml version: NOT INSTALLED or UNKNOWN VERSION.\n", "\n", "For details about installing the optional dependencies, please visit:\n", " https://docs.monai.io/en/latest/installation.html#installing-the-recommended-dependencies\n", "\n" ] } ], "source": [ "import copy\n", "import os\n", "import json\n", "import numpy as np\n", "import nibabel as nib\n", "import subprocess\n", "from IPython.display import Image, display\n", "\n", "from monai.apps import download_url\n", "from monai.data import create_test_image_3d\n", "from monai.config import print_config\n", "\n", "from scripts.diff_model_setting import setup_logging\n", "\n", "print_config()\n", "\n", "logger = setup_logging(\"notebook\")" ] }, { "cell_type": "markdown", "id": "c2389853", "metadata": {}, "source": [ "## Set up the MAISI version\n", "\n", "Choose between `'maisi3d-ddpm'` and `'maisi3d-rflow'`. The differences are:\n", "- The maisi version `'maisi3d-ddpm'` uses basic noise scheduler DDPM. `'maisi3d-rflow'` uses Rectified Flow scheduler, can be 33 times faster during inference.\n", "- The maisi version `'maisi3d-ddpm'` requires training images to be labeled with body regions (`\"top_region_index\"` and `\"bottom_region_index\"`), while `'maisi3d-rflow'` does not have such requirement. In other words, it is easier to prepare training data for `'maisi3d-rflow'`.\n", "- For the released model weights, `'maisi3d-rflow'` can generate images with better quality for head region and small output volumes, and comparable quality for other cases compared with `'maisi3d-ddpm'`." ] }, { "cell_type": "code", "execution_count": 6, "id": "31684f74", "metadata": {}, "outputs": [ { "name": "stderr", "output_type": "stream", "text": [ "[2025-09-22 04:32:05.110][ INFO](notebook) - MAISI version is maisi3d-rflow, whether to use body_region is False\n" ] } ], "source": [ "maisi_version = \"maisi3d-rflow\"\n", "if maisi_version == \"maisi3d-ddpm\":\n", " model_def_path = \"./configs/config_maisi3d-ddpm.json\"\n", "elif maisi_version == \"maisi3d-rflow\":\n", " model_def_path = \"./configs/config_maisi3d-rflow.json\"\n", "else:\n", " raise ValueError(f\"maisi_version has to be chosen from ['maisi3d-ddpm', 'maisi3d-rflow'], yet got {maisi_version}.\")\n", "with open(model_def_path, \"r\") as f:\n", " model_def = json.load(f)\n", "include_body_region = model_def[\"include_body_region\"]\n", "logger.info(f\"MAISI version is {maisi_version}, whether to use body_region is {include_body_region}\")" ] }, { "cell_type": "markdown", "id": "d8e29c23", "metadata": {}, "source": [ "### Simulate a special dataset\n", "\n", "It is widely recognized that training AI models is a time-intensive process. In this instance, we will simulate a small dataset and conduct training over multiple epochs. While the performance may not reach optimal levels due to the abbreviated training duration, the entire pipeline will be completed within minutes.\n", "\n", "`sim_datalist` provides the information of the simulated datasets. It lists 2 training images. The size of the dimension is defined by the `sim_dim`." ] }, { "cell_type": "code", "execution_count": 9, "id": "fc32a7fe", "metadata": {}, "outputs": [], "source": [ "sim_datalist = {\"training\": [{\"image\": \"tr_image_001.nii.gz\"}, {\"image\": \"tr_image_002.nii.gz\"}]}\n", "\n", "sim_dim = (224, 224, 96)" ] }, { "cell_type": "markdown", "id": "b9ac7677", "metadata": {}, "source": [ "### Generate simulated images\n", "\n", "Now we can use MONAI `create_test_image_3d` and `nib.Nifti1Image` functions to generate the 3D simulated images under the `work_dir`." ] }, { "cell_type": "code", "execution_count": 10, "id": "1b199078", "metadata": {}, "outputs": [ { "name": "stderr", "output_type": "stream", "text": [ "[2025-09-22 04:33:43.832][ INFO](notebook) - Generated simulated images.\n" ] } ], "source": [ "work_dir = \"./temp_work_dir\"\n", "if not os.path.isdir(work_dir):\n", " os.makedirs(work_dir)\n", "\n", "dataroot_dir = os.path.join(work_dir, \"sim_dataroot\")\n", "if not os.path.isdir(dataroot_dir):\n", " os.makedirs(dataroot_dir)\n", "\n", "datalist_file = os.path.join(work_dir, \"sim_datalist.json\")\n", "with open(datalist_file, \"w\") as f:\n", " json.dump(sim_datalist, f)\n", "\n", "for d in sim_datalist[\"training\"]:\n", " im, _ = create_test_image_3d(\n", " sim_dim[0], sim_dim[1], sim_dim[2], rad_max=10, num_seg_classes=1, random_state=np.random.RandomState(42)\n", " )\n", " image_fpath = os.path.join(dataroot_dir, d[\"image\"])\n", " nib.save(nib.Nifti1Image(im, affine=np.eye(4)), image_fpath)\n", "\n", "logger.info(\"Generated simulated images.\")" ] }, { "cell_type": "markdown", "id": "a059ddcf-8525-4241-9fe3-b661c4bdd336", "metadata": {}, "source": [ "### Set up directories and configurations\n", "\n", "To optimize the demonstration for time efficiency, we have adjusted the training epochs to 2. Additionally, we modified the `num_splits` parameter in [AutoencoderKlMaisi](https://github.com/Project-MONAI/MONAI/blob/dev/monai/apps/generation/maisi/networks/autoencoderkl_maisi.py#L881) from its default value of 16 to 4. This adjustment reduces the spatial splitting of feature maps in convolutions, which is particularly beneficial given the smaller input size. (This change helps convert convolutions to a for-loop based approach, thereby conserving GPU memory resources.)" ] }, { "cell_type": "code", "execution_count": 11, "id": "6c7b434c", "metadata": {}, "outputs": [ { "name": "stderr", "output_type": "stream", "text": [ "autoencoder_epoch273.pt: 80.0MB [00:13, 6.15MB/s] \n" ] }, { "name": "stdout", "output_type": "stream", "text": [ "2025-09-22 04:34:48,896 - INFO - Downloaded: temp_work_dir/models/autoencoder_epoch273.pt\n", "2025-09-22 04:34:48,897 - INFO - Expected md5 is None, skip md5 check for file temp_work_dir/models/autoencoder_epoch273.pt.\n" ] }, { "name": "stderr", "output_type": "stream", "text": [ "[2025-09-22 04:34:53.651][ INFO](notebook) - files and folders under work_dir: ['config_maisi.json', 'config_maisi_diff_model.json', 'embeddings', 'environment_maisi_diff_model.json', 'models', 'predictions', 'sim_datalist.json', 'sim_dataroot'].\n", "[2025-09-22 04:34:53.653][ INFO](notebook) - number of GPUs: 1.\n" ] } ], "source": [ "env_config_path = \"./configs/environment_maisi_diff_model.json\"\n", "model_config_path = \"./configs/config_maisi_diff_model.json\"\n", "\n", "# Load environment configuration, model configuration and model definition\n", "with open(env_config_path, \"r\") as f:\n", " env_config = json.load(f)\n", "\n", "with open(model_config_path, \"r\") as f:\n", " model_config = json.load(f)\n", "\n", "env_config_out = copy.deepcopy(env_config)\n", "model_config_out = copy.deepcopy(model_config)\n", "model_def_out = copy.deepcopy(model_def)\n", "\n", "# Set up directories based on configurations\n", "env_config_out[\"data_base_dir\"] = dataroot_dir\n", "env_config_out[\"embedding_base_dir\"] = os.path.join(work_dir, env_config_out[\"embedding_base_dir\"])\n", "env_config_out[\"json_data_list\"] = datalist_file\n", "env_config_out[\"model_dir\"] = os.path.join(work_dir, env_config_out[\"model_dir\"])\n", "env_config_out[\"output_dir\"] = os.path.join(work_dir, env_config_out[\"output_dir\"])\n", "trained_autoencoder_path = os.path.join(work_dir, \"models/autoencoder_epoch273.pt\")\n", "env_config_out[\"trained_autoencoder_path\"] = trained_autoencoder_path\n", "trained_autoencoder_path_url = (\n", " \"https://developer.download.nvidia.com/assets/Clara/monai/tutorials/\"\n", " \"model_zoo/model_maisi_autoencoder_epoch273_alternative.pt\"\n", ")\n", "if not os.path.exists(trained_autoencoder_path):\n", " download_url(url=trained_autoencoder_path_url, filepath=trained_autoencoder_path)\n", "\n", "# Create necessary directories\n", "os.makedirs(env_config_out[\"embedding_base_dir\"], exist_ok=True)\n", "os.makedirs(env_config_out[\"model_dir\"], exist_ok=True)\n", "os.makedirs(env_config_out[\"output_dir\"], exist_ok=True)\n", "\n", "env_config_filepath = os.path.join(work_dir, \"environment_maisi_diff_model.json\")\n", "with open(env_config_filepath, \"w\") as f:\n", " json.dump(env_config_out, f, sort_keys=True, indent=4)\n", "\n", "# Update model configuration for demo\n", "max_epochs = 2\n", "model_config_out[\"diffusion_unet_train\"][\"n_epochs\"] = max_epochs\n", "\n", "model_config_filepath = os.path.join(work_dir, \"config_maisi_diff_model.json\")\n", "with open(model_config_filepath, \"w\") as f:\n", " json.dump(model_config_out, f, sort_keys=True, indent=4)\n", "\n", "# Update model definition for demo\n", "model_def_out[\"autoencoder_def\"][\"num_splits\"] = 2\n", "model_def_filepath = os.path.join(work_dir, \"config_maisi.json\")\n", "with open(model_def_filepath, \"w\") as f:\n", " json.dump(model_def_out, f, sort_keys=True, indent=4)\n", "\n", "# Print files and folders under work_dir\n", "logger.info(f\"files and folders under work_dir: {os.listdir(work_dir)}.\")\n", "\n", "# Adjust based on the number of GPUs you want to use\n", "num_gpus = 1\n", "logger.info(f\"number of GPUs: {num_gpus}.\")" ] }, { "cell_type": "code", "execution_count": 8, "id": "95ea6972", "metadata": {}, "outputs": [], "source": [ "def run_torchrun(module, module_args, num_gpus=1):\n", " # Define the arguments for torchrun\n", " num_nodes = 1\n", "\n", " # Build the torchrun command\n", " torchrun_command = [\n", " \"torchrun\",\n", " \"--nproc_per_node\",\n", " str(num_gpus),\n", " \"--nnodes\",\n", " str(num_nodes),\n", " \"-m\",\n", " module,\n", " ] + module_args\n", "\n", " # Set the OMP_NUM_THREADS environment variable\n", " env = os.environ.copy()\n", " env[\"OMP_NUM_THREADS\"] = \"1\"\n", "\n", " # Execute the command\n", " process = subprocess.Popen(torchrun_command, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True, env=env)\n", "\n", " # Print the output in real-time\n", " try:\n", " while True:\n", " output = process.stdout.readline()\n", " if output == \"\" and process.poll() is not None:\n", " break\n", " if output:\n", " print(output.strip())\n", " except Exception as e:\n", " print(f\"An error occurred: {e}\")\n", " finally:\n", " # Capture and print any remaining output\n", " stdout, stderr = process.communicate()\n", " print(stdout)\n", " if stderr:\n", " print(stderr)\n", " return" ] }, { "cell_type": "markdown", "id": "1c904f52", "metadata": {}, "source": [ "## Step 1: Create Training Data\n", "\n", "To train the latent diffusion model, we first store the latent features produced by the autoencoder's encoder in local storage. This allows the latent diffusion model to directly utilize these features, thereby conserving both time and GPU memory during the training process. Additionally, we have provided the script for multi-GPU processing to save latent features from all training images, significantly accelerating the creation of the entire training set.\n", "\n", "The diffusion model utilizes a U-shaped convolutional neural network architecture, requiring matching input and output dimensions. Therefore, it is advisable to resample the input image dimensions to be multiples of 2 for compatibility. In this case, we have chosen dimensions that are multiples of 128." ] }, { "cell_type": "code", "execution_count": 8, "id": "f45ea863", "metadata": {}, "outputs": [ { "name": "stderr", "output_type": "stream", "text": [ "[2025-03-14 16:14:22.326][ INFO](notebook) - Creating training data...\n" ] }, { "name": "stdout", "output_type": "stream", "text": [ "\n", "[2025-03-14 16:14:29.646][ INFO](creating training data) - Using device cuda:0\n", "[2025-03-14 16:14:30.160][ INFO](creating training data) - filenames_raw: ['tr_image_001.nii.gz', 'tr_image_002.nii.gz']\n", "\n" ] } ], "source": [ "logger.info(\"Creating training data...\")\n", "\n", "# Define the arguments for torchrun\n", "module = \"scripts.diff_model_create_training_data\"\n", "module_args = [\n", " \"--env_config\",\n", " env_config_filepath,\n", " \"--model_config\",\n", " model_config_filepath,\n", " \"--model_def\",\n", " model_def_filepath,\n", " \"--num_gpus\",\n", " str(num_gpus),\n", "]\n", "\n", "run_torchrun(module, module_args, num_gpus=num_gpus)" ] }, { "cell_type": "markdown", "id": "ec5c0c4a", "metadata": {}, "source": [ "### Create .json files for embedding files\n", "\n", "The diffusion model necessitates additional input attributes, including output dimension, output spacing, and top/bottom body region. These dimensions and spacing can be extracted from the header information of the training images. The top and bottom body region inputs can be determined through manual examination or by utilizing segmentation masks from tools such as [TotalSegmentator](https://github.com/wasserth/TotalSegmentator) or [MONAI VISTA](https://github.com/Project-MONAI/VISTA). The body regions are formatted as 4-dimensional one-hot vectors: the head and neck region is represented by [1,0,0,0], the chest region by [0,1,0,0], the abdomen region by [0,0,1,0], and the lower body region (below the abdomen) by [0,0,0,1]. The additional input attributes are saved in a separate .json file. In the following example, we assume that the images cover the chest and abdomen regions." ] }, { "cell_type": "code", "execution_count": 9, "id": "0221a658", "metadata": {}, "outputs": [ { "name": "stderr", "output_type": "stream", "text": [ "[2025-03-14 16:14:32.560][ INFO](notebook) - data: {'dim': (64, 64, 32), 'spacing': [0.875, 0.875, 0.75]}.\n", "[2025-03-14 16:14:32.562][ INFO](notebook) - Save json file to ./temp_work_dir/./embeddings/tr_image_001_emb.nii.gz.json\n", "[2025-03-14 16:14:32.563][ INFO](notebook) - data: {'dim': (64, 64, 32), 'spacing': [0.875, 0.875, 0.75]}.\n", "[2025-03-14 16:14:32.564][ INFO](notebook) - Save json file to ./temp_work_dir/./embeddings/tr_image_002_emb.nii.gz.json\n", "[2025-03-14 16:14:32.565][ INFO](notebook) - Completed creating .json files for all embedding files.\n" ] } ], "source": [ "def list_gz_files(folder_path):\n", " \"\"\"List all .gz files in the folder and its subfolders.\"\"\"\n", " gz_files = []\n", " for root, _, files in os.walk(folder_path):\n", " for file in files:\n", " if file.endswith(\".gz\"):\n", " gz_files.append(os.path.join(root, file))\n", " return gz_files\n", "\n", "\n", "def create_json_files(gz_files):\n", " \"\"\"Create .json files for each .gz file with the specified keys and values.\"\"\"\n", " for gz_file in gz_files:\n", " # Load the NIfTI image\n", " img = nib.load(gz_file)\n", "\n", " # Get the dimensions and spacing\n", " dimensions = img.shape\n", " dimensions = dimensions[:3]\n", " spacing = img.header.get_zooms()[:3]\n", " spacing = spacing[:3]\n", " spacing = [float(_item) for _item in spacing]\n", "\n", " # Create the dictionary with the specified keys and values\n", " data = {\"dim\": dimensions, \"spacing\": spacing}\n", " if include_body_region:\n", " # The region can be selected from one of four regions from top to bottom.\n", " # [1,0,0,0] is the head and neck, [0,1,0,0] is the chest region, [0,0,1,0]\n", " # is the abdomen region, and [0,0,0,1] is the lower body region.\n", " data[\"top_region_index\"] = [0, 1, 0, 0] # chest region\n", " data[\"bottom_region_index\"] = [0, 0, 1, 0] # abdomen region\n", " logger.info(f\"data: {data}.\")\n", "\n", " # Create the .json filename\n", " json_filename = gz_file + \".json\"\n", "\n", " # Write the dictionary to the .json file\n", " with open(json_filename, \"w\") as json_file:\n", " json.dump(data, json_file, indent=4)\n", " logger.info(f\"Save json file to {json_filename}\")\n", "\n", "\n", "folder_path = env_config_out[\"embedding_base_dir\"]\n", "gz_files = list_gz_files(folder_path)\n", "create_json_files(gz_files)\n", "\n", "logger.info(\"Completed creating .json files for all embedding files.\")" ] }, { "cell_type": "markdown", "id": "e81a9e48", "metadata": {}, "source": [ "## Step 2: Train the Model\n", "\n", "After all latent features have been created, we will initiate the multi-GPU script to train the latent diffusion model.\n", "\n", "The image generation process utilizes the [DDPM scheduler](https://arxiv.org/pdf/2006.11239) with 1,000 iterative steps. The diffusion model is optimized using L1 loss and a decayed learning rate scheduler. The batch size for this process is set to 1.\n", "\n", "Please be aware that using the H100 GPU may occasionally result in random segmentation faults. To avoid this issue, you can disable AMP by setting the `--no_amp` flag." ] }, { "cell_type": "code", "execution_count": 10, "id": "ade6389d", "metadata": {}, "outputs": [ { "name": "stderr", "output_type": "stream", "text": [ "[2025-03-14 16:14:32.570][ INFO](notebook) - Training the model...\n" ] }, { "name": "stdout", "output_type": "stream", "text": [ "\n", "[2025-03-14 16:14:39.869][ INFO](training) - Using cuda:0 of 1\n", "[2025-03-14 16:14:39.869][ INFO](training) - [config] ckpt_folder -> ./temp_work_dir/./models.\n", "[2025-03-14 16:14:39.869][ INFO](training) - [config] data_root -> ./temp_work_dir/./embeddings.\n", "[2025-03-14 16:14:39.869][ INFO](training) - [config] data_list -> ./temp_work_dir/sim_datalist.json.\n", "[2025-03-14 16:14:39.869][ INFO](training) - [config] lr -> 0.0001.\n", "[2025-03-14 16:14:39.869][ INFO](training) - [config] num_epochs -> 2.\n", "[2025-03-14 16:14:39.869][ INFO](training) - [config] num_train_timesteps -> 1000.\n", "[2025-03-14 16:14:41.316][ INFO](training) - Training from scratch.\n", "[2025-03-14 16:14:41.337][ INFO](training) - num_files_train: 2\n", "[2025-03-14 16:14:41.634][ INFO](training) - Scaling factor set to 1.159693956375122.\n", "[2025-03-14 16:14:41.634][ INFO](training) - scale_factor -> 1.159693956375122.\n", "[2025-03-14 16:14:41.637][ INFO](training) - torch.set_float32_matmul_precision -> highest.\n", "[2025-03-14 16:14:41.637][ INFO](training) - Epoch 1, lr 0.0001.\n", "[2025-03-14 16:14:42.627][ INFO](training) - [2025-03-14 16:14:42] epoch 1, iter 1/2, loss: 1.1344, lr: 0.000100000000.\n", "[2025-03-14 16:14:42.739][ INFO](training) - [2025-03-14 16:14:42] epoch 1, iter 2/2, loss: 1.1275, lr: 0.000056250000.\n", "[2025-03-14 16:14:42.783][ INFO](training) - epoch 1 average loss: 1.1310.\n", "[2025-03-14 16:14:44.540][ INFO](training) - Epoch 2, lr 2.5e-05.\n", "[2025-03-14 16:14:44.981][ INFO](training) - [2025-03-14 16:14:44] epoch 2, iter 1/2, loss: 1.1254, lr: 0.000025000000.\n", "[2025-03-14 16:14:45.106][ INFO](training) - [2025-03-14 16:14:45] epoch 2, iter 2/2, loss: 1.1201, lr: 0.000006250000.\n", "[2025-03-14 16:14:45.177][ INFO](training) - epoch 2 average loss: 1.1227.\n", "\n" ] } ], "source": [ "logger.info(\"Training the model...\")\n", "\n", "# Define the arguments for torchrun\n", "module = \"scripts.diff_model_train\"\n", "module_args = [\n", " \"--env_config\",\n", " env_config_filepath,\n", " \"--model_config\",\n", " model_config_filepath,\n", " \"--model_def\",\n", " model_def_filepath,\n", " \"--num_gpus\",\n", " str(num_gpus),\n", "]\n", "\n", "run_torchrun(module, module_args, num_gpus=num_gpus)" ] }, { "cell_type": "markdown", "id": "4bdf7b17", "metadata": {}, "source": [ "## Step 3: Infer using the Trained Model\n", "\n", "Upon completing the training of the latent diffusion model, we can employ the multi-GPU script to perform inference. By integrating the diffusion model with the autoencoder's decoder, this process will generate 3D images with specified top/bottom body regions, spacing, and dimensions." ] }, { "cell_type": "code", "execution_count": 11, "id": "1626526d", "metadata": {}, "outputs": [ { "name": "stderr", "output_type": "stream", "text": [ "[2025-03-14 16:14:49.136][ INFO](notebook) - Running inference...\n", "[2025-03-14 16:15:02.647][ INFO](notebook) - Completed all steps.\n" ] }, { "name": "stdout", "output_type": "stream", "text": [ "\n", "[2025-03-14 16:14:56.275][ INFO](inference) - Using cuda:0 of 1 with random seed: 59473\n", "[2025-03-14 16:14:56.275][ INFO](inference) - [config] ckpt_filepath -> ./temp_work_dir/./models/diff_unet_ckpt.pt.\n", "[2025-03-14 16:14:56.275][ INFO](inference) - [config] random_seed -> 59473.\n", "[2025-03-14 16:14:56.275][ INFO](inference) - [config] output_prefix -> unet_3d.\n", "[2025-03-14 16:14:56.275][ INFO](inference) - [config] output_size -> (256, 256, 128).\n", "[2025-03-14 16:14:56.275][ INFO](inference) - [config] out_spacing -> (1.0, 1.0, 0.75).\n", "[2025-03-14 16:14:56.275][ INFO](root) - `controllable_anatomy_size` is not provided.\n", "[2025-03-14 16:14:58.525][ INFO](inference) - checkpoints ./temp_work_dir/./models/diff_unet_ckpt.pt loaded.\n", "[2025-03-14 16:14:58.527][ INFO](inference) - scale_factor -> 1.159693956375122.\n", "[2025-03-14 16:14:58.528][ INFO](inference) - num_downsample_level -> 4, divisor -> 4.\n", "[2025-03-14 16:14:58.536][ INFO](inference) - noise: cuda:0, torch.float32, \n", "\n", " 0%| | 0/10 [00:00" ] }, "metadata": {}, "output_type": "display_data" } ], "source": [ "display(Image(filename=\"./figures/maisi_diff-train-examples.png\"))" ] } ], "metadata": { "kernelspec": { "display_name": "monai", "language": "python", "name": "python3" }, "language_info": { "codemirror_mode": { "name": "ipython", "version": 3 }, "file_extension": ".py", "mimetype": "text/x-python", "name": "python", "nbconvert_exporter": "python", "pygments_lexer": "ipython3", "version": "3.10.18" } }, "nbformat": 4, "nbformat_minor": 5 }