train rflow

Signed-off-by: Can-Zhao <canz@nvidia.com>

Author: Can-Zhao

generation/maisi/scripts/diff_model_infer.py

@@ -91,12 +91,11 @@ def prepare_tensors(args: argparse.Namespace, device: torch.device) -> tuple:
     top_region_index_tensor = np.array(args.diffusion_unet_inference["top_region_index"]).astype(float) * 1e2
     bottom_region_index_tensor = np.array(args.diffusion_unet_inference["bottom_region_index"]).astype(float) * 1e2
     spacing_tensor = np.array(args.diffusion_unet_inference["spacing"]).astype(float) * 1e2
-    modality_tensor = np.array([args.diffusion_unet_inference["modality"]]).astype(int)
 
     top_region_index_tensor = torch.from_numpy(top_region_index_tensor[np.newaxis, :]).half().to(device)
     bottom_region_index_tensor = torch.from_numpy(bottom_region_index_tensor[np.newaxis, :]).half().to(device)
     spacing_tensor = torch.from_numpy(spacing_tensor[np.newaxis, :]).half().to(device)
-    modality_tensor = torch.from_numpy(modality_tensor[np.newaxis, :]).long().to(device)
+    modality_tensor = args.diffusion_unet_inference["modality"]*torch.ones((len(spacing_tensor)),dtype=torch.long).to(device)
 
     return top_region_index_tensor, bottom_region_index_tensor, spacing_tensor, modality_tensor
 

generation/maisi/scripts/infer_controlnet.py

@@ -161,6 +161,7 @@ def main():
             top_region_index_tensor = None
             bottom_region_index_tensor = None
         spacing_tensor = batch["spacing"].to(device)
+        modality_tensor = args.controlnet_infer["modality"]*torch.ones((len(labels),),dtype=torch.long).to(device)
         out_spacing = tuple((batch["spacing"].squeeze().numpy() / 100).tolist())
         # get target dimension
         dim = batch["dim"]
@@ -180,6 +181,7 @@ def main():
             top_region_index_tensor=top_region_index_tensor,
             bottom_region_index_tensor=bottom_region_index_tensor,
             spacing_tensor=spacing_tensor,
+            modality_tensor=modality_tensor,
             latent_shape=latent_shape,
             output_size=output_size,
             noise_factor=1.0,

generation/maisi/scripts/utils.py

@@ -306,10 +306,10 @@ def prepare_maisi_controlnet_json_dataloader(
         LoadImaged(keys=["image", "label"], image_only=True, ensure_channel_first=True),
         Orientationd(keys=["label"], axcodes="RAS"),
         EnsureTyped(keys=["label"], dtype=torch.uint8, track_meta=True),
-        Lambdad(keys="top_region_index", func=lambda x: torch.FloatTensor(x)),
-        Lambdad(keys="bottom_region_index", func=lambda x: torch.FloatTensor(x)),
+        Lambdad(keys="top_region_index", func=lambda x: torch.FloatTensor(x),allow_missing_keys=True),
+        Lambdad(keys="bottom_region_index", func=lambda x: torch.FloatTensor(x),allow_missing_keys=True),
         Lambdad(keys="spacing", func=lambda x: torch.FloatTensor(x)),
-        Lambdad(keys=["top_region_index", "bottom_region_index", "spacing"], func=lambda x: x * 1e2),
+        Lambdad(keys=["top_region_index", "bottom_region_index", "spacing"], func=lambda x: x * 1e2,allow_missing_keys=True),
     ]
     train_transforms, val_transforms = Compose(common_transform), Compose(common_transform)