Question on how to train a 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 images look something like this. 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 custom dataset (stained version of the pathology image). Ideally the model would take an unstained image as an input and output the stained version, and the model would learn the diffusion process for this image variation.

I'm currently looking at the code https://github.com/huggingface/diffusers/blob/3816c9ad9fa51c1ad53bd0302962f2329aaa91b8/src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_image_variation.py#L26 from this repo, would this be a good starting point? I'm not sure if it's feasible, as this question has already been asked but was left unanswered in this thread. Any help would be appreciated!

[chokevin8]

Also, since it would be kind of a "style transfer" between the unstained and stained images, I would want the structure and textual elements of the image to be preserved during the transfer- I was wondering if adding an auxillary loss function (ex. structural similarity loss) would do the job? If yes, where would this loss function be added?

[patrickvonplaten]

Maybe @justinpinkney could help here :-)

[chokevin8]

@patrickvonplaten Thank you for the reference, @justinpinkney I would appreciate any help with this. Also, where would be the .yaml config file be found for this specific image variation model? Is it this one? thank you so much!

[pokameng]

@justinpinkney I would appreciate any help with this

[pokameng]

@patrickvonplaten I've been confused for a long time and no one seems to have replied to this question

[chokevin8]

I've been able to run the code, but I'm not sure what's happening, this is what I get as an error message: @justinpinkney I'd appreciate any guidance, thanks a lot in advance.

LatentDiffusion: Running in eps-prediction mode DiffusionWrapper has 859.52 M params. Keeping EMAs of 688. making attention of type 'vanilla' with 512 in_channels Working with z of shape (1, 4, 32, 32) = 4096 dimensions. making attention of type 'vanilla' with 512 in_channels Attempting to load state from C:\Users\Kevin.cache\huggingface\hub\models--lambdalabs--stable-diffusion-image-conditioned\snapshots\c24656f0a8ef81625d4fc4ec097dd97971139ad9\sd-clip-vit-l14-img-embed_ema_only.ckpt Found nested key 'state_dict' in checkpoint, loading this instead

Global seed set to 23 +--------------------- Traceback (most recent call last) ---------------------+	C:\Users\Kevin\PycharmProjects\wsi_analysis\kevin\stable-diffusion\main.py:		683 in
680 old_state = old_state["state_dict"]
681
682 #Check if we need to port weights from 4ch input to 8ch
> 683 in_filters_load = old_state["model.diffusionmodel.input
684 new_state = model.state_dict()
685 in_filters_current = new_state["model.diffusion_model.inp
686 in_shape = in_filters_current.shape

+-----------------------------------------------------------------------------+ KeyError: 'model.diffusion_model.input_blocks.0.0.weight'

During handling of the above exception, another exception occurred: +--------------------- Traceback (most recent call last) ---------------------+	C:\Users\Kevin\PycharmProjects\wsi_analysis\kevin\stable-diffusion\main.py:		939 in
936 dst = os.path.join(dst, "debug_runs", name)
937 os.makedirs(os.path.split(dst)[0], exist_ok=True)
938 os.rename(logdir, dst)
> 939 if trainer.global_rank == 0:
940 rank_zero_print(trainer.profiler.summary())
941

+-----------------------------------------------------------------------------+ NameError: name 'trainer' is not defined

[chokevin8]

I've now successfully handled the issue with the "Found nested key 'state_dict' in checkpoint, loading this instead" and the , "KeyError: 'model.diffusion_model.input_blocks.0.0.weight" error. I've reinstalled the packages and their correct versions in requirements.txt. But now I'm getting two errors: A new error "TypeError: Gradient accumulation supports only int and dict types" and still the existing error "NameError: name 'trainer' is not defined" which is 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.

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

Seems like for the both errors, there may be some bug with either the main.py or the trainer module of the pytorch-lightning.

Would appreciate any help- it's starting to frustrate me with the lack of documentation on this topic. Thanks a lot!

[justinpinkney]

@chokevin8 is the issue above when running my stable-diffusion code or diffusers? If it's for my code better to create an issue over there (I'll try and not be too slow to respond)

[chokevin8]

@justinpinkney It's an issue with your code I think, since I was running the inference code you provided with the config file updated as you previously mentioned in this thread. Since I can't transfer this issue, do you mind transferring this thread to your repo then? Once you transfer them, I can provide more details about the specific errors I'm having with the snippets of code included. Thanks a lot!

[chokevin8]

@justinpinkney Just a heads up for this in case you may have forgotten :)

[justinpinkney]

@chokevin8 can't move this issue I'm afraid, feel free to create a new one in my repo though

[chokevin8]

@justinpinkney Just resummarized my issue and opened a new issue on your repo! Closing this one, this repo is moved to here.

[pokameng]

@chokevin8 hello, i am also face the same problem: KeyError: 'model.diffusion_model.input_blocks.0.0.weight' when i load the finetune model. Can you tell me how to resolve this poblem?\ Thanks!!!