From 93fe3dc99398e1f92cc870ec976a1a80cc934c4a Mon Sep 17 00:00:00 2001 From: Manish Singh Saini Date: Sun, 18 Dec 2022 02:15:18 -0500 Subject: [PATCH 1/6] some model tweaks to work on small dataset with FULL minimagen model from paper --- .gitignore | 5 +++++ inference.py | 1 + minimagen/training.py | 21 +++++++++++++++------ train.py | 4 +++- 4 files changed, 24 insertions(+), 7 deletions(-) create mode 100644 .gitignore diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..db72c77 --- /dev/null +++ b/.gitignore @@ -0,0 +1,5 @@ +temp/* +training_*/ +generated_images_*/ +.idea/* +minimagen/__pycache__/* diff --git a/inference.py b/inference.py index 14f762e..73faddb 100644 --- a/inference.py +++ b/inference.py @@ -13,6 +13,7 @@ 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'] elif not args.CAPTIONS.endswith(".txt"): captions = [args.CAPTIONS] elif args.CAPTIONS.endswith(".txt"): diff --git a/minimagen/training.py b/minimagen/training.py index 43452a8..a96ef7f 100644 --- a/minimagen/training.py +++ b/minimagen/training.py @@ -280,7 +280,8 @@ def ConceptualCaptions(args, smalldata=False, testset=False): """ dset = load_dataset("conceptual_captions") if smalldata: - num = 16 + # num = 16 + num = 100 vi = dset['validation']['image_url'][:num] vc = dset['validation']['caption'][:num] ti = dset['train']['image_url'][:num] @@ -378,6 +379,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 +391,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 +406,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 +449,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: @@ -541,12 +548,14 @@ def load_testing_parameters(args): :param args: Arguments Namespace returned from parsing :func:`~.minimagen.training.get_minimagen_parser`. """ d = dict( - BATCH_SIZE=2, + BATCH_SIZE=8, MAX_NUM_WORDS=32, IMG_SIDE_LEN=128, - EPOCHS=2, + # EPOCHS=2, + EPOCHS=50, 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 ) diff --git a/train.py b/train.py index 7a9b8d4..0248705 100644 --- a/train.py +++ b/train.py @@ -39,11 +39,13 @@ 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) else: - train_dataset, valid_dataset = ConceptualCaptions(args, smalldata=False) + # train_dataset, valid_dataset = ConceptualCaptions(args, smalldata=False) + train_dataset, valid_dataset = ConceptualCaptions(args, smalldata=True) # Create dataloaders dl_opts = {**get_minimagen_dl_opts(device), 'batch_size': args.BATCH_SIZE, 'num_workers': args.NUM_WORKERS} From c60e6d4fb52829ae953adf749fb73a0b642f9601 Mon Sep 17 00:00:00 2001 From: Manish Singh Saini Date: Sun, 18 Dec 2022 18:10:37 -0500 Subject: [PATCH 2/6] minor edits --- .gitignore | 1 + minimagen/training.py | 2 +- train.py | 3 +++ 3 files changed, 5 insertions(+), 1 deletion(-) diff --git a/.gitignore b/.gitignore index db72c77..e074c92 100644 --- a/.gitignore +++ b/.gitignore @@ -3,3 +3,4 @@ training_*/ generated_images_*/ .idea/* minimagen/__pycache__/* +main.py diff --git a/minimagen/training.py b/minimagen/training.py index a96ef7f..8358607 100644 --- a/minimagen/training.py +++ b/minimagen/training.py @@ -552,7 +552,7 @@ def load_testing_parameters(args): MAX_NUM_WORDS=32, IMG_SIDE_LEN=128, # EPOCHS=2, - EPOCHS=50, + EPOCHS=350, T5_NAME='t5_small', # TRAIN_VALID_FRAC=0.5, TRAIN_VALID_FRAC=0.9, diff --git a/train.py b/train.py index 0248705..eff88b7 100644 --- a/train.py +++ b/train.py @@ -47,6 +47,9 @@ # train_dataset, valid_dataset = ConceptualCaptions(args, smalldata=False) train_dataset, valid_dataset = ConceptualCaptions(args, smalldata=True) +# 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) From 23cf88151e2815afaa9ec47da994c866ebe2f065 Mon Sep 17 00:00:00 2001 From: Manish Singh Saini Date: Sun, 18 Dec 2022 18:11:56 -0500 Subject: [PATCH 3/6] minor edits --- .gitignore | 1 + 1 file changed, 1 insertion(+) diff --git a/.gitignore b/.gitignore index e074c92..fad67d0 100644 --- a/.gitignore +++ b/.gitignore @@ -4,3 +4,4 @@ generated_images_*/ .idea/* minimagen/__pycache__/* main.py +venv/* \ No newline at end of file From 14cc1c1c9d0a7d5630d509598b7109289e509ad8 Mon Sep 17 00:00:00 2001 From: Manish Singh Saini Date: Sun, 18 Dec 2022 23:47:50 -0500 Subject: [PATCH 4/6] added custom Unet Base and Super classes for training movie dataset --- dataset/movie_captions.json | 78 ++++++++++++++++++++++++++++++++++++ inference.py | 3 ++ minimagen/Unet.py | 36 ++++++++++++++++- minimagen/training.py | 79 ++++++++++++++++++++++++++++++++++++- train.py | 18 +++++---- 5 files changed, 204 insertions(+), 10 deletions(-) create mode 100644 dataset/movie_captions.json 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 73faddb..1b3cb51 100644 --- a/inference.py +++ b/inference.py @@ -8,6 +8,9 @@ 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: diff --git a/minimagen/Unet.py b/minimagen/Unet.py index 96afdf3..3cc7f69 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=128, + 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 8358607..1184ed3 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 @@ -314,6 +319,52 @@ def ConceptualCaptions(args, smalldata=False, testset=False): 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: + 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] + 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 get_minimagen_dl_opts(device): """Returns dictionary of default MinImagen dataloader options""" return {'batch_size': 4, @@ -548,11 +599,11 @@ def load_testing_parameters(args): :param args: Arguments Namespace returned from parsing :func:`~.minimagen.training.get_minimagen_parser`. """ d = dict( - BATCH_SIZE=8, + BATCH_SIZE=2, MAX_NUM_WORDS=32, IMG_SIDE_LEN=128, # EPOCHS=2, - EPOCHS=350, + EPOCHS=5, T5_NAME='t5_small', # TRAIN_VALID_FRAC=0.5, TRAIN_VALID_FRAC=0.9, @@ -565,6 +616,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=8, + MAX_NUM_WORDS=64, + IMG_SIDE_LEN=128, + # EPOCHS=2, + EPOCHS=50, + 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 eff88b7..fefc5f3 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,13 +39,16 @@ 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 +args.TESTING = True if args.TESTING: - args = load_testing_parameters(args) + # 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 = ConceptualCaptions(args, smalldata=False) - train_dataset, valid_dataset = ConceptualCaptions(args, smalldata=True) + # 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}") @@ -67,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: From 68e304d2ded029340e9429c1160d9666e1a8bb0e Mon Sep 17 00:00:00 2001 From: Manish Singh Saini Date: Tue, 20 Dec 2022 04:23:47 -0500 Subject: [PATCH 5/6] changde some config paramters; added anime_captions.json for testing --- dataset/anime_captions.json | 54 +++++++++++++++++++++++++++++++++++++ minimagen/Unet.py | 2 +- minimagen/training.py | 9 ++++--- train.py | 4 +-- 4 files changed, 62 insertions(+), 7 deletions(-) create mode 100644 dataset/anime_captions.json 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/minimagen/Unet.py b/minimagen/Unet.py index 3cc7f69..a438a98 100644 --- a/minimagen/Unet.py +++ b/minimagen/Unet.py @@ -756,7 +756,7 @@ class BaseMovie(Unet): """ defaults = dict( - dim=128, + dim=256, dim_mults=(1, 2, 4), num_resnet_blocks=2, layer_attns=(False, True, True), diff --git a/minimagen/training.py b/minimagen/training.py index 1184ed3..cca9f37 100644 --- a/minimagen/training.py +++ b/minimagen/training.py @@ -286,7 +286,7 @@ def ConceptualCaptions(args, smalldata=False, testset=False): dset = load_dataset("conceptual_captions") if smalldata: # num = 16 - num = 100 + num = 1000 vi = dset['validation']['image_url'][:num] vc = dset['validation']['caption'][:num] ti = dset['train']['image_url'][:num] @@ -328,7 +328,8 @@ def MovieCaptions(args, smalldata=False, testset=False): :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/movie_captions.json", "r") as f: + with open("dataset/anime_captions.json", "r") as f: dset = json.load(f) if smalldata: num = 16 @@ -623,11 +624,11 @@ def load_semi_workload_parameters(args): :param args: Arguments Namespace returned from parsing :func:`~.minimagen.training.get_minimagen_parser`. """ d = dict( - BATCH_SIZE=8, + BATCH_SIZE=2, MAX_NUM_WORDS=64, IMG_SIDE_LEN=128, # EPOCHS=2, - EPOCHS=50, + EPOCHS=200, T5_NAME='t5_base', # TRAIN_VALID_FRAC=0.5, TRAIN_VALID_FRAC=0.9, diff --git a/train.py b/train.py index fefc5f3..608ac95 100644 --- a/train.py +++ b/train.py @@ -44,8 +44,8 @@ # 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) + 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) From 9c6d1437ba8b7a8eb257bd2c56ffc1ce7aeb039c Mon Sep 17 00:00:00 2001 From: Manish Singh Saini Date: Tue, 20 Dec 2022 14:54:11 -0500 Subject: [PATCH 6/6] added custom captions in inference.py --- inference.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/inference.py b/inference.py index 1b3cb51..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 @@ -17,6 +18,9 @@ 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"):