Fix ValueError in maisi_train_controlnet_tutorial.ipynb

generation/maisi/maisi_inference_tutorial.ipynb

@@ -364,25 +364,25 @@
     "device = torch.device(\"cuda\")\n",
     "\n",
     "autoencoder = define_instance(args, \"autoencoder_def\").to(device)\n",
-    "checkpoint_autoencoder = torch.load(args.trained_autoencoder_path)\n",
+    "checkpoint_autoencoder = torch.load(args.trained_autoencoder_path, weights_only=True)\n",
     "autoencoder.load_state_dict(checkpoint_autoencoder)\n",
     "\n",
     "diffusion_unet = define_instance(args, \"diffusion_unet_def\").to(device)\n",
-    "checkpoint_diffusion_unet = torch.load(args.trained_diffusion_path)\n",
+    "checkpoint_diffusion_unet = torch.load(args.trained_diffusion_path, weights_only=False)\n",
     "diffusion_unet.load_state_dict(checkpoint_diffusion_unet[\"unet_state_dict\"], strict=True)\n",
     "scale_factor = checkpoint_diffusion_unet[\"scale_factor\"].to(device)\n",
     "\n",
     "controlnet = define_instance(args, \"controlnet_def\").to(device)\n",
-    "checkpoint_controlnet = torch.load(args.trained_controlnet_path)\n",
+    "checkpoint_controlnet = torch.load(args.trained_controlnet_path, weights_only=False)\n",
     "monai.networks.utils.copy_model_state(controlnet, diffusion_unet.state_dict())\n",
     "controlnet.load_state_dict(checkpoint_controlnet[\"controlnet_state_dict\"], strict=True)\n",
     "\n",
     "mask_generation_autoencoder = define_instance(args, \"mask_generation_autoencoder_def\").to(device)\n",
-    "checkpoint_mask_generation_autoencoder = torch.load(args.trained_mask_generation_autoencoder_path)\n",
+    "checkpoint_mask_generation_autoencoder = torch.load(args.trained_mask_generation_autoencoder_path, weights_only=True)\n",
     "mask_generation_autoencoder.load_state_dict(checkpoint_mask_generation_autoencoder)\n",
     "\n",
     "mask_generation_diffusion_unet = define_instance(args, \"mask_generation_diffusion_def\").to(device)\n",
-    "checkpoint_mask_generation_diffusion_unet = torch.load(args.trained_mask_generation_diffusion_path)\n",
+    "checkpoint_mask_generation_diffusion_unet = torch.load(args.trained_mask_generation_diffusion_path, weights_only=True)\n",
     "mask_generation_diffusion_unet.load_state_dict(checkpoint_mask_generation_diffusion_unet[\"unet_state_dict\"])\n",
     "mask_generation_scale_factor = checkpoint_mask_generation_diffusion_unet[\"scale_factor\"]\n",
     "\n",

generation/maisi/scripts/diff_model_create_training_data.py

@@ -20,6 +20,7 @@
 import nibabel as nib
 import numpy as np
 import torch
+import torch.distributed as dist
 
 import monai
 from monai.transforms import Compose
@@ -146,7 +147,7 @@ def process_file(
         out_path.parent.mkdir(parents=True, exist_ok=True)
         logger.info(f"out_filename: {out_filename}")
 
-        with torch.cuda.amp.autocast():
+        with torch.amp.autocast("cuda"):
             pt_nda = torch.from_numpy(nda_image).float().to(device).unsqueeze(0).unsqueeze(0)
             z = autoencoder.encode_stage_2_inputs(pt_nda)
             logger.info(f"z: {z.size()}, {z.dtype}")
@@ -175,7 +176,7 @@ def diff_model_create_training_data(env_config_path: str, model_config_path: str
 
     autoencoder = define_instance(args, "autoencoder_def").to(device)
     try:
-        checkpoint_autoencoder = torch.load(args.trained_autoencoder_path)
+        checkpoint_autoencoder = torch.load(args.trained_autoencoder_path, weights_only=True)
         autoencoder.load_state_dict(checkpoint_autoencoder)
     except Exception:
         logger.error("The trained_autoencoder_path does not exist!")
@@ -202,6 +203,9 @@ def diff_model_create_training_data(env_config_path: str, model_config_path: str
 
         process_file(filepath, args, autoencoder, device, plain_transforms, new_transforms, logger)
 
+    if dist.is_initialized():
+        dist.destroy_process_group()
+
 
 if __name__ == "__main__":
     parser = argparse.ArgumentParser(description="Diffusion Model Training Data Creation")

generation/maisi/scripts/diff_model_infer.py

@@ -20,6 +20,7 @@
 import nibabel as nib
 import numpy as np
 import torch
+import torch.distributed as dist
 from tqdm import tqdm
 
 from monai.inferers import sliding_window_inference
@@ -59,13 +60,13 @@ def load_models(args: argparse.Namespace, device: torch.device, logger: logging.
     """
     autoencoder = define_instance(args, "autoencoder_def").to(device)
     try:
-        checkpoint_autoencoder = torch.load(args.trained_autoencoder_path)
+        checkpoint_autoencoder = torch.load(args.trained_autoencoder_path, weights_only=True)
         autoencoder.load_state_dict(checkpoint_autoencoder)
     except Exception:
         logger.error("The trained_autoencoder_path does not exist!")
 
     unet = define_instance(args, "diffusion_unet_def").to(device)
-    checkpoint = torch.load(f"{args.model_dir}/{args.model_filename}", map_location=device)
+    checkpoint = torch.load(f"{args.model_dir}/{args.model_filename}", map_location=device, weights_only=False)
     unet.load_state_dict(checkpoint["unet_state_dict"], strict=True)
     logger.info(f"checkpoints {args.model_dir}/{args.model_filename} loaded.")
 
@@ -149,7 +150,7 @@ def run_inference(
     autoencoder.eval()
     unet.eval()
 
-    with torch.cuda.amp.autocast(enabled=True):
+    with torch.amp.autocast("cuda", enabled=True):
         for t in tqdm(noise_scheduler.timesteps, ncols=110):
             model_output = unet(
                 x=image,
@@ -271,7 +272,7 @@ def diff_model_infer(env_config_path: str, model_config_path: str, model_def_pat
     )
 
     timestamp = datetime.now().strftime("%Y%m%d%H%M%S")
-    output_path = "{0}/{1}_seed{2}_size{3:d}x{4:d}x{5:d}_spacing{6:.2f}x{7:.2f}x{8:.2f}_{9}.nii.gz".format(
+    output_path = "{0}/{1}_seed{2}_size{3:d}x{4:d}x{5:d}_spacing{6:.2f}x{7:.2f}x{8:.2f}_{9}_rank{10}.nii.gz".format(
         args.output_dir,
         output_prefix,
         random_seed,
@@ -282,9 +283,13 @@ def diff_model_infer(env_config_path: str, model_config_path: str, model_def_pat
         out_spacing[1],
         out_spacing[2],
         timestamp,
+        local_rank,
     )
     save_image(data, output_size, out_spacing, output_path, logger)
 
+    if dist.is_initialized():
+        dist.destroy_process_group()
+
 
 if __name__ == "__main__":
     parser = argparse.ArgumentParser(description="Diffusion Model Inference")

generation/maisi/scripts/diff_model_setting.py

@@ -32,7 +32,8 @@ def setup_logging(logger_name: str = "") -> logging.Logger:
         logging.Logger: Configured logger.
     """
     logger = logging.getLogger(logger_name)
-    logger.addFilter(RankFilter())
+    if dist.is_initialized():
+        logger.addFilter(RankFilter())
     logging.basicConfig(
         level=logging.INFO,
         format="[%(asctime)s.%(msecs)03d][%(levelname)5s](%(name)s) - %(message)s",
@@ -80,9 +81,13 @@ def initialize_distributed() -> tuple:
     Returns:
         tuple: local_rank, world_size, and device.
     """
-    dist.init_process_group(backend="nccl", init_method="env://")
-    local_rank = dist.get_rank()
-    world_size = dist.get_world_size()
+    if torch.cuda.is_available() and torch.cuda.device_count() > 1:
+        dist.init_process_group(backend="nccl", init_method="env://")
+        local_rank = dist.get_rank()
+        world_size = dist.get_world_size()
+    else:
+        local_rank = 0
+        world_size = 1
     device = torch.device("cuda", local_rank)
     torch.cuda.set_device(device)
     return local_rank, world_size, device

generation/maisi/scripts/diff_model_train.py

@@ -20,7 +20,7 @@
 
 import torch
 import torch.distributed as dist
-from torch.cuda.amp import GradScaler, autocast
+from torch.amp import GradScaler, autocast
 from torch.nn.parallel import DistributedDataParallel
 
 import monai
@@ -64,17 +64,22 @@ def prepare_data(
     Returns:
         ThreadDataLoader: Data loader for training.
     """
+
+    def _load_data_from_file(file_path, key):
+        with open(file_path) as f:
+            return torch.FloatTensor(json.load(f)[key])
+
     train_transforms = Compose(
         [
             monai.transforms.LoadImaged(keys=["image"]),
             monai.transforms.EnsureChannelFirstd(keys=["image"]),
             monai.transforms.Lambdad(
-                keys="top_region_index", func=lambda x: torch.FloatTensor(json.load(open(x))["top_region_index"])
+                keys="top_region_index", func=lambda x: _load_data_from_file(x, "top_region_index")
             ),
             monai.transforms.Lambdad(
-                keys="bottom_region_index", func=lambda x: torch.FloatTensor(json.load(open(x))["bottom_region_index"])
+                keys="bottom_region_index", func=lambda x: _load_data_from_file(x, "bottom_region_index")
             ),
-            monai.transforms.Lambdad(keys="spacing", func=lambda x: torch.FloatTensor(json.load(open(x))["spacing"])),
+            monai.transforms.Lambdad(keys="spacing", func=lambda x: _load_data_from_file(x, "spacing")),
             monai.transforms.Lambdad(keys="top_region_index", func=lambda x: x * 1e2),
             monai.transforms.Lambdad(keys="bottom_region_index", func=lambda x: x * 1e2),
             monai.transforms.Lambdad(keys="spacing", func=lambda x: x * 1e2),
@@ -231,7 +236,7 @@ def train_one_epoch(
 
         optimizer.zero_grad(set_to_none=True)
 
-        with autocast(enabled=True):
+        with autocast("cuda", enabled=True):
             noise = torch.randn(
                 (num_images_per_batch, 4, images.size(-3), images.size(-2), images.size(-1)), device=device
             )
@@ -365,7 +370,7 @@ def diff_model_train(env_config_path: str, model_config_path: str, model_def_pat
     ]
     lr_scheduler = create_lr_scheduler(optimizer, total_steps)
     loss_pt = torch.nn.L1Loss()
-    scaler = GradScaler()
+    scaler = GradScaler("cuda")
 
     torch.set_float32_matmul_precision("highest")
     logger.info("torch.set_float32_matmul_precision -> highest.")
@@ -403,6 +408,9 @@ def diff_model_train(env_config_path: str, model_config_path: str, model_def_pat
                 args,
             )
 
+    if dist.is_initialized():
+        dist.destroy_process_group()
+
 
 if __name__ == "__main__":
     parser = argparse.ArgumentParser(description="Diffusion Model Training")

generation/maisi/scripts/infer_controlnet.py

@@ -69,9 +69,12 @@ def main():
     logger.info(f"Number of GPUs: {torch.cuda.device_count()}")
     logger.info(f"World_size: {world_size}")
 
-    env_dict = json.load(open(args.environment_file, "r"))
-    config_dict = json.load(open(args.config_file, "r"))
-    training_config_dict = json.load(open(args.training_config, "r"))
+    with open(args.environment_file, "r") as env_file:
+        env_dict = json.load(env_file)
+    with open(args.config_file, "r") as config_file:
+        config_dict = json.load(config_file)
+    with open(args.training_config, "r") as training_config_file:
+        training_config_dict = json.load(training_config_file)
 
     for k, v in env_dict.items():
         setattr(args, k, v)

generation/maisi/scripts/sample.py

@@ -117,7 +117,7 @@ def ldm_conditional_sample_one_mask(
     """
     recon_model = ReconModel(autoencoder=autoencoder, scale_factor=scale_factor).to(device)
 
-    with torch.no_grad(), torch.cuda.amp.autocast():
+    with torch.no_grad(), torch.amp.autocast("cuda"):
         # Generate random noise
         latents = initialize_noise_latents(latent_shape, device)
         anatomy_size = torch.FloatTensor(anatomy_size).unsqueeze(0).unsqueeze(0).half().to(device)
@@ -226,7 +226,7 @@ def ldm_conditional_sample_one_image(
 
     recon_model = ReconModel(autoencoder=autoencoder, scale_factor=scale_factor).to(device)
 
-    with torch.no_grad(), torch.cuda.amp.autocast():
+    with torch.no_grad(), torch.amp.autocast("cuda"):
         logging.info("---- Start generating latent features... ----")
         start_time = time.time()
         # generate segmentation mask

generation/maisi/scripts/train_controlnet.py

@@ -23,7 +23,7 @@
 import torch.nn.functional as F
 from monai.networks.utils import copy_model_state
 from monai.utils import RankFilter
-from torch.cuda.amp import GradScaler, autocast
+from torch.amp import GradScaler, autocast
 from torch.nn.parallel import DistributedDataParallel as DDP
 from torch.utils.tensorboard import SummaryWriter
 
@@ -71,9 +71,12 @@ def main():
     logger.info(f"Number of GPUs: {torch.cuda.device_count()}")
     logger.info(f"World_size: {world_size}")
 
-    env_dict = json.load(open(args.environment_file, "r"))
-    config_dict = json.load(open(args.config_file, "r"))
-    training_config_dict = json.load(open(args.training_config, "r"))
+    with open(args.environment_file, "r") as env_file:
+        env_dict = json.load(env_file)
+    with open(args.config_file, "r") as config_file:
+        config_dict = json.load(config_file)
+    with open(args.training_config, "r") as training_config_file:
+        training_config_dict = json.load(training_config_file)
 
     for k, v in env_dict.items():
         setattr(args, k, v)
@@ -151,7 +154,7 @@ def main():
 
     # Step 4: training
     n_epochs = args.controlnet_train["n_epochs"]
-    scaler = GradScaler()
+    scaler = GradScaler("cuda")
     total_step = 0
     best_loss = 1e4
 
@@ -174,7 +177,7 @@ def main():
 
             optimizer.zero_grad(set_to_none=True)
 
-            with autocast(enabled=True):
+            with autocast("cuda", enabled=True):
                 # generate random noise
                 noise_shape = list(inputs.shape)
                 noise = torch.randn(noise_shape, dtype=inputs.dtype).to(device)