Package: mediffusion
MediffusionDiffusion models have significantly impacted the realm of image generation. In a bid to reduce the technical complexity, we aim to lower the entry barrier for the medical community. To achieve this, we have introducedmediffusion, a user-friendly diffusion package that can be effortlessly tailored to address medical problems with less than 20 lines of code. We have utilized various codebases, includingguided diffusionandLDM, enhancing their robustness for medical use cases. We plan to update this package regularly. Embracing the spirit of open science, we invite you to consider sharing  a demo notebook of your work should you choose to utilize this package.Happy Coding!Setup and InstallationStep 1: Create a Conda EnvironmentIf you haven't installed Conda yet, you can download it fromhere. After installing, create a new Conda environment by running:condacreate--namemediffusionpython=3.10Activate the environment:condaactivatemediffusionStep 2: Install PyTorchInstall PyTorch specifically for CUDA 11.8 by running:pipinstalltorchtorchvisiontorchaudio--index-urlhttps://download.pytorch.org/whl/cu118Step 3: Install The PackageYou can install the latest version from github using:pipinstallmediffusionThis will install all the necessary packages.Training1. HyperparametersBefore starting the training, it is recommended that you set up some global constants and environment variables:os.environ["CUDA_VISIBLE_DEVICES"]="0,1"os.environ['WANDB_API_KEY']="WANDB-API-KEY"TOTAL_IMAGE_SEEN=40e6BATCH_SIZE=36NUM_DEVICES=2# number of devices in CUDA_VISIBLE_DEVICESTRAIN_ITERATIONS=int(TOTAL_IMAGE_SEEN/(BATCH_SIZE*NUM_DEVICES))2. Preparing DataTo prepare the data, you need to create a dataset where each element is a dictionary. The dictionary should have the key "img" and may also contain additional keys like "cls" and "concat" depending on the type of condition. One way to do this is by using MONAI. Below is a sample code snippet:importmonaiasmntrain_data_dicts=[{"img":"./image1.dcm","cls":2},{"img":"./image2.dcm","cls":0}]valid_data_dicts=[{"img":"./image9.dcm","cls":1}]transforms=mn.transforms.Compose([mn.transforms.LoadImageD(keys="img"),mn.transforms.SelectItemsD(keys=["img","cls"]),mn.transforms.ScaleIntensityD(keys=["img"],minv=-1,maxv=1),mn.transforms.ToTensorD(keys=["img","cls"],dtype=torch.float,track_meta=False),])train_ds=mn.data.Dataset(data=train_data_dicts,transform=transforms)valid_ds=mn.data.Dataset(data=valid_data_dicts,transform=transforms)train_sampler=torch.utils.data.RandomSampler(train_ds,replacement=True,num_samples=TOTAL_IMAGE_SEEN)At the end of this step, you should havetrain_ds,val_dsandtrain_sampler.3. Configuring ModelConfiguration Fields ExplanationBelow is a table that provides descriptions for each element in the configuration file:SectionFieldDescriptiondiffusiontimestepsThe number of timesteps in the diffusion processschedule_nameThe name of the schedule (e.g., "cosine")enforce_zero_terminal_snrWhether to enforce zero terminal SNR (True/False)schedule_paramsParameters related to the diffusion schedule-- beta_startStarting value for beta in the schedule-- beta_endEnding value for beta in the schedule-- cosine_sParameter for cosine scheduletimestep_respacingCan be a list of respacings. For example, with 200 steps, [10,20] means in the first 100, get 10 samples and in the next 100, get 20 samples.mean_typeType of mean model (e.g., "VELOCITY")var_typeType of variance model (e.g., "LEARNED_RANGE")loss_typeThe type of loss to use (e.g., "MSE")optimizerlrLearning ratetypeThe type of optimizer to usevalidationclassifier_cond_scaleClassifier guidance scale for validation logging.protocolInference protocol for logging validation resultslog_originalWhether to log the original validation data (True/False)log_concatWhether to log the concatenated images (True/False)log_cls_indicesWhether to log the entire cls vector (default value of -1), or specefic indices from the cls vector (user should provide a list of desired cls indices)modelinput_sizeThe input size of the model. Can be an integer for square and cube images or a list of integers for specific axes, like [64, 64, 32]dimsNumber of dimensions, 2 or 3 for 2D and 3D imagesattention_resolutionsList of resolutions for attention layerschannel_multList of multipliers for each layer's channelsdropoutDropout ratein_channelsNumber of input channels (image channels + concat channels)out_channelsNumber of output channels (image channels or image channels * 2 if learning the variance)model_channelsNumber of convolution channels in the modelnum_head_channelsNumber of attention head channelsnum_headsNumber of attention headsnum_heads_upsampleNumber of attention head after upsamplingnum_res_blocksList of the number of residual blocks for each layerresblock_updownWhether to use residual blocks for down/up sampling (True/False)use_checkpointWhether to use checkpointing (True/False)use_new_attention_orderWhether to use the new attention ordering (True/False)use_scale_shift_normWhether to use scale-shift normalization (True/False)scale_skip_connectionWhether to scaleskip connections (True/False)num_classesNumber of classes for conditioningconcat_channelsNumber of concatenatong channels for conditioning (for super-resolution or inpainting)guidance_drop_probDrop probability for the classifier free guidance scale trainingFor sample configurations, please checkout thesample_configsdirectory.Note: If a field is left out of the config file, the default value is infered based on this file:mediffusion/default_config/default.yaml.Instantiating ModelYou can instantiate the model using the configuration file and dataset as follows:frommediffusionimportDiffusionModulemodel=DiffusionModule("./config.yaml",train_ds=train_ds,val_ds=valid_ds,dl_workers=2,train_sampler=train_sampler,batch_size=32,# train batch sizeval_batch_size=16# validation batch size (recommended size is half of batch_size))4. Setting Up TrainerYou can set up the trainer using theTrainerclass:frommediffusionimportTrainertrainer=Trainer(max_steps=TRAIN_ITERATIONS,val_check_interval=5000,root_directory="./outputs",# where to save the weights and logsprecision="16-mixed",# mixed precision trainingdevices=-1,# use all the devices in CUDA_VISIBLE_DEVICESnodes=1,wandb_project="Your_Project_Name",logger_instance="Your_Logger_Instance",)5. Training the ModelFinally, to train your model, you simply call:trainer.fit(model)Prediction1. Loading the ModelFirst, import theDiffusionModuleclass and load the pre-trained model checkpoint. The model is then moved to the CUDA device and set to inference mode. Additionally, you may choose to enable half-precision for better performance:frommediffusionimportDiffusionModulemodel=DiffusionModule("./config.yaml")model.load_ckpt("./outputs/pl/last.ckpt",ema=True)model.cuda().half()model.eval()2. Preparing InputPrepare the noise and model keyword arguments. Here,"cls"specifies the class condition and is set to 0:importtorchnoise=torch.randn(1,1,256,256)model_kwargs={"cls":torch.tensor([0]).cuda().half()}Note: You can use other keys likeconcatand/orcls_embed. To find out more, look at thetutorialsdirectory.3. Making PredictionsTo make a prediction, use thepredictmethod from theDiffusionModuleclass:img=model.predict(noise,model_kwargs=model_kwargs,classifier_cond_scale=4,inference_protocol="DDIM100")noise: The input noise tensormodel_kwargs: A dictionary containing additional model configurations (e.g., class conditions)classifier_cond_scale: The scale used for the classifier free guidance condition during inferenceinference_protocol: The inference protocol to be used (e.g.,"DDIM100")Theimgis the generated output based on the model's inference (C:H:W(:D)). To save the image, you need to transpose it first, due to the different axis conventions.Note: The model currently supports the following solvers:DDPM,DDIM,IDDIM(for inverse diffusion), andPLMS. As an example,"PLMS100"means using thePLMSsolver for100steps.TutorialsFor more hands-on tutorials on how to effectively use this package, please check thetutorialsfolder in the GitHub repository. These tutorials provide step-by-step instructions, Colab notebooks, and explanations to help you get started with the software.File NameDescriptionNotebook Link01_2d_ddpmGetting started with training a simple 2D class-conditioned DDPM.📓02_2d_inpaintingImage inpainting with 2D diffusion model (repaint method)📓TO-DOThe following features and improvements are currently on our development roadmap:Cross-attentionDPM-SolverVAE for LDMWe are actively working on these features and they will be available in future releases.Issues and ContributionsIssuesIf you encounter any issues while using this package, we encourage you to open an issue in the GitHub repository. Your feedback helps us to improve the software and resolve any bugs or limitations.ContributionsContributions to the codebase are always welcome. If you have a feature request, bugfix, or any other contribution, feel free to submit a pull request.Development OpportunitiesIf you're interested in actively participating in the development of this package, please send us a Direct Message (DM). We're always open to collaboration and would be delighted to have you on board.CitationIf you find this work useful, please consider citing the parent project:@article{KHOSRAVI2023107832,title={Few-shot biomedical image segmentation using diffusion models: Beyond image generation},journal={Computer Methods and Programs in Biomedicine},volume={242},pages={107832},year={2023},issn={0169-2607},doi={https://doi.org/10.1016/j.cmpb.2023.107832},url={https://www.sciencedirect.com/science/article/pii/S0169260723004984},author={Bardia Khosravi and Pouria Rouzrokh and John P. Mickley and Shahriar Faghani and Kellen Mulford and Linjun Yang and A. Noelle Larson and Benjamin M. Howe and Bradley J. Erickson and Michael J. Taunton and Cody C. Wyles},}