Question on how to train my own stable diffusion image variation model

[chokevin8]

Hello, I'm trying to come up with an image variation model, where the input would be an unstained pathology image and the model would output a stained version of the pathology image. Both input and output are png images of same size. The unstained and stained image transformation would look something like this. Ideally, the diffusion model would be trained on true dataset of stained version of pathology image so that it can receive unstained version of pathology image as an input and output the "virtually" stained version of pathology image. I was wondering if it was possible to train a custom model by using the stable diffusion model found here on Huggingface with my own dataset of stained version of the pathology image.

I've gone ahead and tried implementing this fine-tuning process using the fine-tuning code from your example while using the model from here but I encountered two errors: "TypeError: Gradient accumulation supports only int and dict types" and the error "NameError: name 'trainer' is not defined" which are both related to the pytorch lightning trainer. I have the correct version of pytorch-lightning installed (1.4.2) just as required in requirements.txt. Maybe this is different since this is the fine tuning code from the pokemon example? This is the line where the error occurs:

main.py lines 865-870:
if 'accumulate_grad_batches' in lightning_config.trainer:
accumulate_grad_batches = lightning_config.trainer.accumulate_grad_batches
else:
accumulate_grad_batches = 1
rank_zero_print(f"accumulate_grad_batches = {accumulate_grad_batches}")
lightning_config.trainer.accumulate_grad_batches = accumulate_grad_batches

Lastly, where would be the .yaml config file be found for this specific image variation model? Is it this one? It doesn't say anywhere which config file to use. Thank you so much for any help, really appreciate it.

[pokameng]

@chokevin8
how to solve the bug: KeyError: 'model.diffusion_model.input_blocks.0.0.weight'
Thanks!