Can I train Imagen for super-resolution only?

[hmza09]

I want to train unets for super resolution. I have images with size 213x213 and I want to get super resolution images of 640x640. Can someone guide me if the script I'm using is correct or not and what should I alter to get desirable results.
`
import torch
from imagen_pytorch import Unet, Imagen, SRUnet256, ImagenTrainer

unet1 = Unet(
dim = 32,
dim_mults = (1, 2, 4),
num_resnet_blocks = 3,
layer_attns = (False, True, True),
layer_cross_attns = False,
use_linear_attn = True
)

unet2 = SRUnet256(
dim = 32,
dim_mults = (1, 2, 4),
num_resnet_blocks = (2, 4, 8),
layer_attns = (False, False, True),
layer_cross_attns = False
)

imagen = Imagen(
condition_on_text = False,
unets = (unet1, unet2),
image_sizes = (64, 256),
timesteps = 1000
)
dataset = Dataset('Image', image_size = 256)
trainer = ImagenTrainer(
imagen = imagen,
split_valid_from_train = True
).cuda()
trainer.add_train_dataset(dataset, batch_size = 16)
l_arr = []
for i in range(20000):
loss = trainer.train_step(unet_number = 1, max_batch_size = 4)
trainer.update(unet_number = 1)
l_arr.append(loss)
print(f'loss: {loss}')

if not (i % 500):
    valid_loss = trainer.valid_step(unet_number = 1, max_batch_size = 4)
    print(f'valid loss: {valid_loss}')

if not (i % 1000) and trainer.is_main: # is_main makes sure this can run in distributed
    images = trainer.sample(batch_size = 1, return_pil_images = True) 
    images[0].save(f'./sample-{i // 100}.png')
    trainer.save(f'checkpoint_{i // 100}.pt')

`

[lucidrains]

@hmza09 hey yes, i will add this feature at some point!

[lucidrains]

@hmza09 let me know if ae5a1e1 works!