About the Stage 1 training frozen parameters

[AHupuJR]

Thank you very much for open-sourcing the codes. Good work

Just want to ask for the stage 1 traning, in the paper there is "we first only pre-train the alignment module, while fixing the VAE encoder and the diffusion model". But in fact, in the codes, in stage 1 training, the whole unet (alignment module + denoising unet) is trainable. And codes to train only the alignment module are commented. So should we uncomment this when training?

FaithDiff/train_SDXL_stage_1.py

Lines 627 to 636 in 682021a

	unet.train() # If you GPU memory is limited, you can set it to unet.requires_grad_(False)
	'''
	unet.requires_grad_(False)
	for param in unet.condition_embedding.parameters():
	param.requires_grad = True
	for param in unet.information_transformer_layes.parameters():
	param.requires_grad = True
	for param in unet.spatial_ch_projs.parameters():
	param.requires_grad = True
	'''