diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..fad67d0 --- /dev/null +++ b/.gitignore @@ -0,0 +1,7 @@ +temp/* +training_*/ +generated_images_*/ +.idea/* +minimagen/__pycache__/* +main.py +venv/* \ No newline at end of file diff --git a/dataset/anime_captions.json b/dataset/anime_captions.json new file mode 100644 index 0000000..47191a6 --- /dev/null +++ b/dataset/anime_captions.json @@ -0,0 +1,54 @@ +{ + "train": { + "image_url": [ + "https://m.media-amazon.com/images/I/51g-1QzlBaL._AC_SY580_.jpg", + "https://m.media-amazon.com/images/I/51VOIAeQLVL._AC_UF894,1000_QL80_.jpg", + "https://m.media-amazon.com/images/I/41gsqNDTjOL._AC_SY580_.jpg", + "https://m.media-amazon.com/images/I/61ioYU4QfQL._AC_UF894,1000_QL80_.jpg", + "https://m.media-amazon.com/images/I/71L+1etqrJL._AC_UF894,1000_QL80_.jpg", + "https://m.media-amazon.com/images/I/61OJ3XqcMuL._AC_UF894,1000_QL80_.jpg", + "https://m.media-amazon.com/images/I/61Yu1RLNhBL._AC_UF894,1000_QL80_.jpg", + "https://m.media-amazon.com/images/I/51W0Hr+L+ML._AC_SY580_.jpg", + "https://m.media-amazon.com/images/I/51nenSklFgL._AC_SY780_.jpg", + "https://m.media-amazon.com/images/I/91pBAMSzZbL._AC_UF894,1000_QL80_.jpg", + "https://i.ebayimg.com/images/g/QDQAAOSwjvNhsXYp/s-l1600.png", + "https://m.media-amazon.com/images/I/81UaaTlfBVL._AC_UF894,1000_QL80_.jpg", + "https://m.media-amazon.com/images/I/511KfynF4jL._AC_UF894,1000_QL80_.jpg", + "https://m.media-amazon.com/images/I/61V5-czH4rL._AC_UF894,1000_QL80_.jpg", + "https://m.media-amazon.com/images/I/713cslmT3aL.jpg", + "https://m.media-amazon.com/images/I/71BPDF6V1FL._AC_UF894,1000_QL80_.jpg" + ], + "caption": [ + "I don't want to rule anything. Being King of the Pirates means being more free than anyone.", + "There's no such thing as a painless lesson. They just don't exist. Sacrifices are necessary; you can't gain anything without losing something first.", + "Quit it with your stupid reasoning! I meant to say that I'll stand any pain for my friends!! I'm not gonna give up on them", + "There’s nothing special about being born. Not a thing. Most of the universe is just death, nothing more", + "No, see, I don't think of it like I'm saving the world. The fact is, it's because I'm usually trying to challenge the strongest warriors I can find.", + "I love the kind of woman that can kick my ass.", + "Fine, Kakarot, you are the mightiest Saiyan, I've admitted that much. At least for now. But don't you dare think that this is over", + "I am the Pretty Guardian who Fights for Love and for Justice. I am Sailor Moon! And now in the name of the moon, I will punish you!", + "Nobody cares whether I live or die.", + "Just because you die doesn’t mean that the people you killed will come back to life. instead using this sword to save just one more soul is repentance in true sense.", + "It's not a sense of justice. Figuring out difficult cases is my hobby. If you measured good and evil deeds by current laws, I would be responsible for many crimes.", + "If it's for the sake of avenging those who were killed, to make sure there are no more victims, I won't hesitate to bring my blade down upon any demon's neck", + "If I can't even protect my captain's dream, then whatever ambition I have is nothing but talk! Luffy must be the man who becomes the Pirate King!", + "In the ninja world, those who break the rules are trash, that's true, but those who abandon their comrades are worse than trash.", + "If we catch Kira, he is evil. If he wins and rules the world, then he is justice.", + "I never wanted to take out the Uchiha clan. But it was the only way to save my little brother." + ] + }, + "validation": { + "image_url": [ + "https://m.media-amazon.com/images/I/51g-1QzlBaL._AC_SY580_.jpg", + "https://m.media-amazon.com/images/I/51VOIAeQLVL._AC_UF894,1000_QL80_.jpg", + "https://m.media-amazon.com/images/I/41gsqNDTjOL._AC_SY580_.jpg", + "https://m.media-amazon.com/images/I/61ioYU4QfQL._AC_UF894,1000_QL80_.jpg" + ], + "caption": [ + "I don't want to rule anything. Being King of the Pirates means being more free than anyone.", + "There's no such thing as a painless lesson. They just don't exist. Sacrifices are necessary; you can't gain anything without losing something first.", + "Quit it with your stupid reasoning! I meant to say that I'll stand any pain for my friends!! I'm not gonna give up on them", + "There’s nothing special about being born. Not a thing. Most of the universe is just death, nothing more" + ] + } +} \ No newline at end of file diff --git a/dataset/movie_captions.json b/dataset/movie_captions.json new file mode 100644 index 0000000..6d1e081 --- /dev/null +++ b/dataset/movie_captions.json @@ -0,0 +1,78 @@ +{ + "train": { + "image_url": [ + "https://image.tmdb.org/t/p/w500/1hRoyzDtpgMU7Dz4JF22RANzQO7.jpg", + "https://m.media-amazon.com/images/M/MV5BMjAxMzY3NjcxNF5BMl5BanBnXkFtZTcwNTI5OTM0Mw@@._V1_.jpg", + "https://m.media-amazon.com/images/I/41kTVLeW1CL._AC_.jpg", + "https://m.media-amazon.com/images/I/81SIVdnkUmL.jpg", + "https://m.media-amazon.com/images/I/91tuQ9AVNxL.jpg", + "https://m.media-amazon.com/images/I/81kz06oSUeL.jpg", + "https://m.media-amazon.com/images/I/912vL7muRwL.jpg", + "https://m.media-amazon.com/images/I/51V0Mw23gmL._AC_.jpg", + "https://m.media-amazon.com/images/I/81Luju2cHuL.jpg", + "https://m.media-amazon.com/images/I/41XiDnayGtL.jpg", + "https://m.media-amazon.com/images/M/MV5BN2EwM2I5OWMtMGQyMi00Zjg1LWJkNTctZTdjYTA4OGUwZjMyXkEyXkFqcGdeQXVyMTMxODk2OTU@._V1_FMjpg_UX1000_.jpg", + "https://m.media-amazon.com/images/M/MV5BMTk4ODQzNDY3Ml5BMl5BanBnXkFtZTcwODA0NTM4Nw@@._V1_.jpg", + "https://m.media-amazon.com/images/I/51unGrb-AAL._AC_.jpg", + "https://cdn.shopify.com/s/files/1/0057/3728/3618/products/3998dd3fa7e628e415e9805b960bec61_480x.progressive.jpg?v=1573592743", + "https://m.media-amazon.com/images/M/MV5BMjE5MzcyNjk1M15BMl5BanBnXkFtZTcwMjQ4MjcxOQ@@._V1_.jpg", + "https://m.media-amazon.com/images/I/81EBp0vOZZL.jpg" + ], + "caption": [ + "Your mind is the scene of the crime.\n", + "Why So Serious?\n", + "Enter the World of Pandora.\n", + "Some assembly required.\n", + "Witness the beginning of a happy ending\n", + "Mankind was born on Earth. It was never meant to die here.\n", + "Life, liberty and the pursuit of vengeance.\n", + "All heroes start somewhere.\n", + "Mischief. Mayhem. Soap.\n", + "May The Odds Be Ever In Your Favor.\n", + "What a Lovely Day.\n", + "The Legend Ends\n", + "Welcome to the Real World.\n", + "Heroes aren't born. They're built.\n", + "Unleash the power behind the armor.\n", + "One ring to rule them all\n" + ] + }, + "validation": { + "image_url": [ + "https://image.tmdb.org/t/p/w500/1hRoyzDtpgMU7Dz4JF22RANzQO7.jpg", + "https://m.media-amazon.com/images/M/MV5BMjAxMzY3NjcxNF5BMl5BanBnXkFtZTcwNTI5OTM0Mw@@._V1_.jpg", + "https://m.media-amazon.com/images/I/41kTVLeW1CL._AC_.jpg", + "https://m.media-amazon.com/images/I/81SIVdnkUmL.jpg", + "https://m.media-amazon.com/images/I/91tuQ9AVNxL.jpg", + "https://m.media-amazon.com/images/I/81kz06oSUeL.jpg", + "https://m.media-amazon.com/images/I/912vL7muRwL.jpg", + "https://m.media-amazon.com/images/I/51V0Mw23gmL._AC_.jpg", + "https://m.media-amazon.com/images/I/81Luju2cHuL.jpg", + "https://m.media-amazon.com/images/I/41XiDnayGtL.jpg", + "https://m.media-amazon.com/images/M/MV5BN2EwM2I5OWMtMGQyMi00Zjg1LWJkNTctZTdjYTA4OGUwZjMyXkEyXkFqcGdeQXVyMTMxODk2OTU@._V1_FMjpg_UX1000_.jpg", + "https://m.media-amazon.com/images/M/MV5BMTk4ODQzNDY3Ml5BMl5BanBnXkFtZTcwODA0NTM4Nw@@._V1_.jpg", + "https://m.media-amazon.com/images/I/51unGrb-AAL._AC_.jpg", + "https://cdn.shopify.com/s/files/1/0057/3728/3618/products/3998dd3fa7e628e415e9805b960bec61_480x.progressive.jpg?v=1573592743", + "https://m.media-amazon.com/images/M/MV5BMjE5MzcyNjk1M15BMl5BanBnXkFtZTcwMjQ4MjcxOQ@@._V1_.jpg", + "https://m.media-amazon.com/images/I/81EBp0vOZZL.jpg" + ], + "caption": [ + "Your mind is the scene of the crime.\n", + "Why So Serious?\n", + "Enter the World of Pandora.\n", + "Some assembly required.\n", + "Witness the beginning of a happy ending\n", + "Mankind was born on Earth. It was never meant to die here.\n", + "Life, liberty and the pursuit of vengeance.\n", + "All heroes start somewhere.\n", + "Mischief. Mayhem. Soap.\n", + "May The Odds Be Ever In Your Favor.\n", + "What a Lovely Day.\n", + "The Legend Ends\n", + "Welcome to the Real World.\n", + "Heroes aren't born. They're built.\n", + "Unleash the power behind the armor.\n", + "One ring to rule them all\n" + ] + } +} \ No newline at end of file diff --git a/inference.py b/inference.py index 14f762e..2986b8f 100644 --- a/inference.py +++ b/inference.py @@ -1,3 +1,4 @@ +import json from argparse import ArgumentParser from minimagen.generate import load_minimagen, sample_and_save @@ -8,11 +9,18 @@ parser.add_argument("-d", "--TRAINING_DIRECTORY", dest="TRAINING_DIRECTORY", help="Training directory to use for inference", type=str) args = parser.parse_args() +# Override the training directory +# args.TRAINING_DIRECTORY = "training_20221218_035553" + minimagen = load_minimagen(args.TRAINING_DIRECTORY) if args.CAPTIONS is None: print("\nNo caption supplied - using the default of \"a happy dog\".\n") captions = ['a happy dog'] + captions = ['a very typical bus station'] + with open("dataset/anime_captions.json", "r") as f: + dset = json.load(f) + captions = dset["train"]["caption"] elif not args.CAPTIONS.endswith(".txt"): captions = [args.CAPTIONS] elif args.CAPTIONS.endswith(".txt"): diff --git a/minimagen/Unet.py b/minimagen/Unet.py index 96afdf3..a438a98 100644 --- a/minimagen/Unet.py +++ b/minimagen/Unet.py @@ -747,4 +747,38 @@ class SuperTest(Unet): memory_efficient=True ) def __init__(self, *args, **kwargs): - super().__init__(*args, **{**Super.defaults, **kwargs}) \ No newline at end of file + super().__init__(*args, **{**Super.defaults, **kwargs}) + + +class BaseMovie(Unet): + """ + Base image generation U-Net for movie poster generation. + """ + + defaults = dict( + dim=256, + dim_mults=(1, 2, 4), + num_resnet_blocks=2, + layer_attns=(False, True, True), + layer_cross_attns=(False, True, True), + memory_efficient=False + ) + + def __init__(self, *args, **kwargs): + super().__init__(*args, **{**Base.defaults, **kwargs}) + + +class SuperMovie(Unet): + """ + Super-Resolution U-Net for movie poster generation. + """ + defaults = dict( + dim=128, + dim_mults=(1, 2, 4), + num_resnet_blocks=(2, 4, 8), + layer_attns=(False, False, True), + layer_cross_attns=(False, False, True), + memory_efficient=True + ) + def __init__(self, *args, **kwargs): + super().__init__(*args, **{**Super.defaults, **kwargs}) diff --git a/minimagen/training.py b/minimagen/training.py index 43452a8..cca9f37 100644 --- a/minimagen/training.py +++ b/minimagen/training.py @@ -208,6 +208,11 @@ def get_minimagen_parser(): parser.add_argument("-test", "--TESTING", dest="TESTING", help="Whether to test with smaller dataset", action='store_true') parser.set_defaults(TESTING=False) + + parser.add_argument("-movie", "--MOVIE", dest="MOVIE", help="Whether to use movie dataset", + action='store_true') + parser.set_defaults(MOVIE=False) + return parser @@ -279,8 +284,56 @@ def ConceptualCaptions(args, smalldata=False, testset=False): :return: test_dataset if :code:`testset` else (train_dataset, valid_dataset) """ dset = load_dataset("conceptual_captions") + if smalldata: + # num = 16 + num = 1000 + vi = dset['validation']['image_url'][:num] + vc = dset['validation']['caption'][:num] + ti = dset['train']['image_url'][:num] + tc = dset['train']['caption'][:num] + dset = datasets.Dataset = {'train': { + 'image_url': ti, + 'caption': tc, + }, 'num_rows': num, + 'validation': { + 'image_url': vi, + 'caption': vc, }, 'num_rows': num} + + if testset: + # Torch test dataset + test_dataset = MinimagenDataset(dset, max_length=args.MAX_NUM_WORDS, train=False, encoder_name=args.T5_NAME, + side_length=args.IMG_SIDE_LEN) + return test_dataset + else: + # Torch train/valid dataset + dataset_train_valid = MinimagenDataset(dset, max_length=args.MAX_NUM_WORDS, encoder_name=args.T5_NAME, + train=True, + side_length=args.IMG_SIDE_LEN) + + # Split into train/valid + train_size = int(args.TRAIN_VALID_FRAC * len(dataset_train_valid)) + valid_size = len(dataset_train_valid) - train_size + train_dataset, valid_dataset = torch.utils.data.random_split(dataset_train_valid, [train_size, valid_size]) + if args.VALID_NUM is not None: + valid_dataset.indices = valid_dataset.indices[:args.VALID_NUM + 1] + return train_dataset, valid_dataset + + +def MovieCaptions(args, smalldata=False, testset=False): + """ + Load `movie captions dataset: ~/MinImagen/dataset/movie_captions.json`_ + + :param args: Arguments Namespace/dictionary parsed from :func:`~.minimagen.training.get_minimagen_parser` + :param smalldata: Whether to return a small subset of the data (for testing code) + :param testset: Whether to return the testing set (vs training/valid) + :return: test_dataset if :code:`testset` else (train_dataset, valid_dataset) + """ + # with open("dataset/movie_captions.json", "r") as f: + with open("dataset/anime_captions.json", "r") as f: + dset = json.load(f) if smalldata: num = 16 + # num = 100 vi = dset['validation']['image_url'][:num] vc = dset['validation']['caption'][:num] ti = dset['train']['image_url'][:num] @@ -378,6 +431,7 @@ def train(): # Every 10% of the way through epoch, save states in case of training failure if batch_num % args.CHCKPT_NUM == 0: + # if int(batch_num % (total_num_batches / 10)) == 0: with training_dir(): with open('training_progess.txt', 'a') as f: f.write(f'{"-" * 10}Checkpoint created at batch number {batch_num}{"-" * 10}\n') @@ -389,7 +443,9 @@ def train(): torch.save(imagen.unets[idx].state_dict(), model_path) # Write and batch average training loss so far - avg_loss = [i / batch_num for i in running_train_loss] + # avg_loss = [i / batch_num for i in running_train_loss] + # Above line of code seems to be incorrectly written by the original devs; below seems more correct + avg_loss = [i / len(batch["image"]) for i in running_train_loss] with training_dir(): with open('training_progess.txt', 'a') as f: f.write(f'U-Nets Avg Train Losses Epoch {epoch + 1} Batch {batch_num}: ' @@ -402,7 +458,8 @@ def train(): imagen.train(False) print(f'\n{"-" * 10}Validation...{"-" * 10}') - for vbatch in tqdm(valid_dataloader): + # for vbatch in tqdm(valid_dataloader): + for vbatch in valid_dataloader: if not vbatch: continue @@ -444,6 +501,8 @@ def train(): imagen.train(True) running_train_loss = [0. for i in range(len(unets))] + # total_num_batches = math.ceil(len(train_dataloader) / train_dataloader.batch_size) + total_num_batches = len(train_dataloader) print(f'\n{"-" * 10}Training...{"-" * 10}') for batch_num, batch in tqdm(enumerate(train_dataloader)): try: @@ -544,9 +603,11 @@ def load_testing_parameters(args): BATCH_SIZE=2, MAX_NUM_WORDS=32, IMG_SIDE_LEN=128, - EPOCHS=2, + # EPOCHS=2, + EPOCHS=5, T5_NAME='t5_small', - TRAIN_VALID_FRAC=0.5, + # TRAIN_VALID_FRAC=0.5, + TRAIN_VALID_FRAC=0.9, TIMESTEPS=25, # Do not make less than 20 OPTIM_LR=0.0001 ) @@ -556,6 +617,30 @@ def load_testing_parameters(args): return args +def load_semi_workload_parameters(args): + """ + Load command line arguments that are conducive to testing training movie generation scripts. + + :param args: Arguments Namespace returned from parsing :func:`~.minimagen.training.get_minimagen_parser`. + """ + d = dict( + BATCH_SIZE=2, + MAX_NUM_WORDS=64, + IMG_SIDE_LEN=128, + # EPOCHS=2, + EPOCHS=200, + T5_NAME='t5_base', + # TRAIN_VALID_FRAC=0.5, + TRAIN_VALID_FRAC=0.9, + TIMESTEPS=1000, # Do not make less than 20 + OPTIM_LR=0.0001 + ) + + # Replace relevant values in arg dict + args.__dict__ = {**args.__dict__, **d} + return args + + def create_directory(dir_path): """ Creates a training directory at the given path if it does not exist already and returns a context manager that diff --git a/train.py b/train.py index 7a9b8d4..608ac95 100644 --- a/train.py +++ b/train.py @@ -6,12 +6,12 @@ from minimagen.Imagen import Imagen -from minimagen.Unet import Unet, Base, Super, BaseTest, SuperTest +from minimagen.Unet import Unet, Base, Super, BaseTest, SuperTest, BaseMovie, SuperMovie from minimagen.generate import load_minimagen, load_params from minimagen.t5 import get_encoded_dim from minimagen.training import get_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 + load_restart_training_parameters, load_testing_parameters, MovieCaptions, load_semi_workload_parameters # Get device device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu") @@ -39,12 +39,20 @@ args = load_restart_training_parameters(args, justparams=True) # If testing, lower parameter values to lower computational load and also to lower amount of data being used. +args.TESTING = True if args.TESTING: - args = load_testing_parameters(args) - train_dataset, valid_dataset = ConceptualCaptions(args, smalldata=True) + # args = load_testing_parameters(args) + args = load_semi_workload_parameters(args) + + train_dataset, valid_dataset = MovieCaptions(args, smalldata=True) + # train_dataset, valid_dataset = ConceptualCaptions(args, smalldata=True) else: + # train_dataset, valid_dataset = MovieCaptions(args, smalldata=False) train_dataset, valid_dataset = ConceptualCaptions(args, smalldata=False) +# print current state of args +print(f"args: {args}") + # Create dataloaders dl_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) @@ -62,7 +70,8 @@ # If not loading a training from a checkpoint if args.TESTING: # If testing, use tiny MinImagen for low computational load - unets_params = [get_default_args(BaseTest), get_default_args(SuperTest)] + # unets_params = [get_default_args(BaseTest), get_default_args(SuperTest)] + unets_params = [get_default_args(BaseMovie), get_default_args(SuperMovie)] # Else if not loading Unet/Imagen settings from a config (parameters) folder, use defaults elif not args.PARAMETERS: