Handle corrupt images

[CerebralSeed]

This code will inform the user of a corrupt image in their set and continue training.

[lucidrains]

oh nice

not sure if selecting another image is the best way to go about it though

[CerebralSeed]

Ideally, we would just skip that index. Perhaps could send an argument out, if that image load was an error. How about this:

    def __getitem__(self, index):
        try:
            path = self.paths[index]
            img = Image.open(path)
            img = self.transform(img)
            err = False
        except:
            print("Error with image", str(path))
            img = None
            err = True
        return img, err

And then in the discriminator trainer:

            if self.unconditional:
                real_images, err = next(dl_iter)
            else:
                result, err = next(dl_iter)
                assert isinstance(result, tuple), 'dataset should return a tuple of two items for text conditioned training, (images: Tensor, texts: List[str])'
                real_images, texts = result
            if not err: 
                # requires grad for real images, for gradient penalty
    
                real_images = real_images.to(self.device)
                real_images.requires_grad_()
                ...

And something similar for the generator, inside of the generate_kwargs.

[lucidrains]

@CerebralSeed maybe we should just run through the entire dataset once at init? could be parallelized too with joblib

[CerebralSeed]

@lucidrains - been under the weather the last few days and stuck in bed unable to focus. Feeling better now. Just started going through some of the changes you've made and, wow! I've got a lot of catching up to do.

Interesting idea. How long would that take, though, with LAION 5B?

By the way, if you're open to the idea, I can go through the code and make it generalized for the 1D and 3D cases with an added argument on initialization.

Once you get a more finalized version, I'm going to work on incorporating 3D Conv branches on the 2D case as a variant. I've found that fixes the generative model issue with hands and other objects with multiple degrees of freedom on the depth dimension in UNets.

I saw you removed the discriminator learning rate multiplier recently. Did that not work out so well?