Add ability to train decoder using embedding-image pairs

[Veldrovive]

I am implementing a single node training script for the decoder and it seems @lucidrains has implemented a wrapper script for this purpose that is already feature-full. Currently, the forward pass is implemented as follows:

DALLE2-pytorch/dalle2_pytorch/train.py

Lines 189 to 199 in 1d5dc08

	def forward(
	self,
	x,
	*,
	unet_number,
	divisor = 1,
	**kwargs
	):
	with autocast(enabled = self.amp):
	loss = self.decoder(x, unet_number = unet_number, **kwargs)
	return self.scale(loss / divisor, unet_number = unet_number)

This lacks the ability to substitute our own image embeddings in the case where we have precomputed embedding-image pairs. The functionality is already mostly supported by the Decoder network where image_embed can be passed to the forward method so this could be implemented by simply adding the image_embed parameter as a pass though to decoder.forward. However, it would also be convenient to make the clip model optional in the Decoder constructor. I already started on this a week ago in this branch by adding the ability to set clip_image_size and channels separately from a clip model.

There are only a few small changes that would be necessary to implement this feature so I could put together a pull request to do this.

[lucidrains]

@Veldrovive Hi Aidan! Indeed that is the case, and I can get this finished in the next half hour, been meaning to get around to it!

[lucidrains]

@Veldrovive here you go https://github.com/lucidrains/DALLE2-pytorch/releases/tag/0.0.106

[Veldrovive]

Great! For the DecoderTrainer are you thinking to just use kwargs for image_embed and not put a specific named parameter for it?

[lucidrains]

@Veldrovive yup, for wrapper i usually just forward kwargs to whatever is being wrapped (instead of using some fancy forwarding module)