some model tweaks to work on small dataset with FULL minimagen model …

.gitignore

@@ -0,0 +1,7 @@
+temp/*
+training_*/
+generated_images_*/
+.idea/*
+minimagen/__pycache__/*
+main.py
+venv/*

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"
+    ]
+  }
+}

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"
+    ]
+  }
+}

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"):

minimagen/Unet.py

@@ -747,4 +747,38 @@ class SuperTest(Unet):
         memory_efficient=True
     )
     def __init__(self, *args, **kwargs):
-        super().__init__(*args, **{**Super.defaults, **kwargs})
+        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})

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