Package: minimagen
MinImagenA Minimal implementation of theImagentext-to-image model.SeeBuild Your Own Imagen Text-to-Image Modelfor a tutorial on how to build MinImagen.SeeHow Imagen Actually Worksfor a detailed explanation of Imagen's operating principles.Given a caption of an image, the text-to-image modelImagenwill generate an image that reflects the scene described by the caption. The model is acascading diffusion model, using aT5 text encoderto generate a caption encoding which conditions a base image generator and then a sequence of super-resolution models through which the output of the base image generator is passed.In particular, two notable contributions are the developments of:Noise Conditioning Augmentation, which noises low-resolution conditioning images in the super-resolution models, andDynamic Thresholdingwhich helps prevent image saturation at highclassifier-free guidanceweights.Table of ContentsAttribution NoteInstallationDocumentationUsage - Command Linemain.py- training and image generation in sequencetrain.py- training a MinImagen instanceinference.py- generated images using a MinImagen instanceUsage - PackageTrainingImage GenerationAdditional ResourcesSocialsAttribution NoteThis implementation is largely based on Phil Wang'sImagen implementation.InstallationTo install MinImagen, run the following command in the terminal:$pipinstallminimagenNote that MinImagen requires Python3.9 or higherDocumentationSee theMinImagen Documentationto learn more about the package.Usage - Command LineIf you have cloned this repo, you can use the provided scripts to get started with MinImagen.main.pyFor the most basic usage, simply enter the MinImagen directory and run the following in the terminal:$pythonmain.pyThis will create a small MinImagen instance and train it on a tiny amount of data, and then use this MinImagen instance to generate an image.Afterrunning the script, you will see a directory calledtraining_<TIMESTAMP>.This directory is called aTraining Directoryand is generated when training a MinImagen instance.It contains information about the configuration (parameterssubdirectory), and contains the model checkpoints (state_dictsandtmpdirectories).It also contains atraining_progress.txtfile that records training progress.You will also see a directory calledgenerated_images_<TIMESTEP>.This directory contains a folder of images generated by the model (generated_images).It also containscaptions.txtfiles, which documents the captions that were input to get the images (where the line index of a given caption corresponds to the image number in thegenerated_iamgesfolder).Finally, this directory also containsimagen_training_directory.txt, which specifies the name of the Training Directory used to load the MinImagen instance / generate images.train.pymain.pysimply runstrain.pyandinference.pyin series, the former to train the model and the latter to generate the image.To train a model, simply runtrain.pyand specify relevant command line arguments. Thepossible argumentsare:--PARAMETERSor-p, which specifies a directory that specifies the MinImagen configuration to use. It should be structured like aparameterssubdirectory within a Training Directory (example inparameters).--NUM_WORKERS"or-n, which specifies the number of workers to use for the DataLoaders.--BATCH_SIZEor-b, which specifies the batch size to use during training.--MAX_NUM_WORDSor-mw, which specifies the maximum number of words allowed in a caption.--IMG_SIDE_LENor-s, specifies the final side length of the square images the MinImagen will output.--EPOCHSor-e, which specifies the number of training epochs.--T5_NAME-t5, which specifies the name of T5 encoder to use.--TRAIN_VALID_FRACor-f, which specifies the fraction of dataset to use for training (vs. validation).--TIMESTEPSor-t, which specifies the number of timesteps in Diffusion Process.--OPTIM_LRor-lr, which specifies the learning rate for Adam optimizer.--ACCUM_ITERor-ai, which specifies the number of batches to accumulate for gradient accumulation.--CHCKPT_NUMor-cn, which specifies the interval of batches to create a temporary model checkpoint at during training.--VALID_NUMor-vn, which specifies the number of validation images to use. If None, uses full amount from train/valid split. The reason for including this is that, even with an e.g. 0.99--TRAIN_VALID_FRAC, a prohibitively large number of images could still be left for validation for very large datasets.--RESTART_DIRECTORYor-rd, training directory to load MinImagen instance from if resuming training. A new Training Directory will be created for the training, leaving the previous Training Directory from which the checkpoint is loaded unperturbed.--TESTINGor-test, which is used to run the script with a small MinImagen instance and small dataset for testing.For example, to run a small training using the provided exampleparametersfolder, run the following in the terminal:pythontrain.py--PARAMETERS./parameters--BATCH_SIZE2--TIMESTEPS25--TESTINGAfter execution, you will see a newtraining_<TIMESTAMP>Training Directorythat contains the files aslisted abovefrom the training.inference.pyTo generate images using a model from aTraining Directory, we can useinference.py. Simply runinference.pyand specify relevant command line arguments. The possible arguments are:--TRAINING_DIRECTORY"or-d, which specifies the training directory from which to load the MinImagen instance for inference.--CAPTIONSor-c, which specifies either (a) a single caption to generate an image for, or (b) a filepath to a.txtfile that contains a list of captions to generate images for, where each caption is on a new line.For example, to generate images for the example captions provided incaptions.txtusing the model generated from the above training line, simply runpythoninference.py-CAPTIONScaptions.txt--TRAINING_DIRECTORYtraining_<TIMESTAMP>whereTIMESTAMPis replaced with the appropriate value from your training.Usage - PackageTrainingA minimal training script using theminimagenpackage is shown below. Seetrain.pyfor a more built-up version of the below code.importosfromdatetimeimportdatetimeimporttorch.utils.datafromtorchimportoptimfromminimagen.ImagenimportImagenfromminimagen.UnetimportUnet,Base,Super,BaseTest,SuperTestfromminimagen.generateimportload_minimagen,load_paramsfromminimagen.t5importget_encoded_dimfromminimagen.trainingimportget_minimagen_parser,ConceptualCaptions,get_minimagen_dl_opts,\create_directory,get_model_params,get_model_size,save_training_info,get_default_args,MinimagenTrain,\load_restart_training_parameters,load_testing_parameters# Get devicedevice=torch.device("cuda:0"iftorch.cuda.is_available()else"cpu")# Command line argument parserparser=get_minimagen_parser()args=parser.parse_args()# Create training directorytimestamp=datetime.now().strftime("%Y%m%d_%H%M%S")dir_path=f"./training_{timestamp}"training_dir=create_directory(dir_path)# Replace some cmd line args to lower computational load.args=load_testing_parameters(args)# Load subset of Conceptual Captions dataset.train_dataset,valid_dataset=ConceptualCaptions(args,smalldata=True)# Create dataloadersdl_opts={**get_minimagen_dl_opts(device),'batch_size':args.BATCH_SIZE,'num_workers':args.NUM_WORKERS}train_dataloader=torch.utils.data.DataLoader(train_dataset,**dl_opts)valid_dataloader=torch.utils.data.DataLoader(valid_dataset,**dl_opts)# Use small U-Nets to lower computational load.unets_params=[get_default_args(BaseTest),get_default_args(SuperTest)]unets=[Unet(**unet_params).to(device)forunet_paramsinunets_params]# Specify MinImagen parametersimagen_params=dict(image_sizes=(int(args.IMG_SIDE_LEN/2),args.IMG_SIDE_LEN),timesteps=args.TIMESTEPS,cond_drop_prob=0.15,text_encoder_name=args.T5_NAME)# Create MinImagen from UNets with specified imagen parametersimagen=Imagen(unets=unets,**imagen_params).to(device)# Fill in unspecified arguments with defaults to record complete config (parameters) fileunets_params=[{**get_default_args(Unet),**i}foriinunets_params]imagen_params={**get_default_args(Imagen),**imagen_params}# Get the size of the Imagen model in megabytesmodel_size_MB=get_model_size(imagen)# Save all training info (config files, model size, etc.)save_training_info(args,timestamp,unets_params,imagen_params,model_size_MB,training_dir)# Create optimizeroptimizer=optim.Adam(imagen.parameters(),lr=args.OPTIM_LR)# Train the MinImagen instanceMinimagenTrain(timestamp,args,unets,imagen,train_dataloader,valid_dataloader,training_dir,optimizer,timeout=30)Image GenerationA minimal inference script using theminimagenpackage is shown below. Seeinference.pyfor a more built-up version of the below code.fromargparseimportArgumentParserfromminimagen.generateimportload_minimagen,sample_and_save# Command line argument parserparser=ArgumentParser()parser.add_argument("-d","--TRAINING_DIRECTORY",dest="TRAINING_DIRECTORY",help="Training directory to use for inference",type=str)args=parser.parse_args()# Specify the caption(s) to generate images forcaptions=['a happy dog']# Use `sample_and_save` to generate and save the iamgessample_and_save(captions,training_directory=args.TRAINING_DIRECTORY)# Alternatively, rather than specifying a Training Directory, you can input just a MinImagen instance to use for image generation.# In this case, information about the MinImagen instance used to generate the images will not be saved.minimagen=load_minimagen(args.TRAINING_DIRECTORY)sample_and_save(captions,minimagen=minimagen)To see more of what MinImagen has to offer, or to get additional details on the scripts above, check out theMinImagen DocumentationAdditional ResourcesFor a step-by-step guide on how to build the version of Imagen in this repository, seeBuild Your Own Imagen Text-to-Image Model.For an deep-dive into how Imagen works, seeHow Imagen Actually Works.For additional learning resources for Machine Learning and Deep Learning, check out ourBlog.Check out ourYouTube channelfor videos and other learning resources on Machine Learning and Deep Learning.Follow our newsletterto stay up to date on our recent content.Read about Diffusion Models in ourIntroduction to Diffusion Models for Machine Learningguide.Read more about Imagen here.SocialsFollow us onTwitterfor more Deep Learning content.