ogkalu2 / sketch-guided-stable-diffusion

Unofficial Implementation of the Google Paper - https://sketch-guided-diffusion.github.io/

Sketch-Guided-Text-To-Image-Diffusion

Unofficial Implementation of the Google Paper - https://sketch-guided-diffusion.github.io/

Pretrained LGP

I have Uploaded a pretrained LGP on Stable Diffusion v1.5 here - https://huggingface.co/ogkalu/Stable-Diffusion-Latent-Guidance-Predictor. Trained on 4995 Imagenet Images for 16 epochs.

Training the Latent Guidance Predictor

First, you will need a dataset of a few thousamd images and their corresponding edge maps. You can download only as many Imagenet Images as needed here - https://github.com/mf1024/ImageNet-Datasets-Downloader (Note: Set Multi-Processing Workers to 1 to avoid downloading more Images than specified). Next, use pidinet to generate the edge maps. I suggest using my fork to easily generate the edge maps for a large corpus of images. https://github.com/ogkalu2/pidinet-for-imagenet

The Only requirement to use my train_LGP.py file to train the LGP is that the dataset be in this format -

• Edge Maps and Image Dataset in separate folders
• The Image Dataset Folder should consist of subfolders, their names corresponding to the labels/captions of each image
• If the Image name is example-name.jpg then the Edge Map name should be example-name.png
• These requirements are only to load the data as https://github.com/mf1024/ImageNet-Datasets-Downloader and https://github.com/ogkalu2/pidinet-for-imagenet provide them

# Train
# The unet and vae arguments can be changed to custom models available on hugging face or locally with the diffusers files in the appropriate subfolders. It is however uneccesary to re-train the LGP for most custom models.

!python train_LGP.py --dataset_dir /path/to/image/dataset --edge_maps_dir /path/to/edge_map/dataset --batch_size 15 --LGP_path /path/to/save/trained/LGP --epochs 16 --lr 0.0001 --device cuda --vae runwayml/stable-diffusion-v1-5 --unet runwayml/stable-diffusion-v1-5

Evaluating the Latent Guidance Predictor

This is to see the edge maps generated by a pretrained LGP from the extracted Unet activations. This is not Stable Diffusion Inference

!python Evaluate_LGP.py --caption "woman" --noise_strength 0.3 --image_path /path/to/image/or/sketch/ --LGP_path /path/to/pretrained/LGP --vae runwayml/stable-diffusion-v1-5 --unet runwayml/stable-diffusion-v1-5 --device cuda

Pretrained LGP

I have Uploaded a pretrained LGP on Stable Diffusion v1.5 here - https://huggingface.co/ogkalu/Stable-Diffusion-Latent-Guidance-Predictor

Citations

@article{voynov2022sketch, title={Sketch-Guided Text-to-Image Diffusion Models}, author={Voynov, Andrey and Abernan, Kfir and Cohen-Or, Daniel}, booktitle={arXiv preprint arXiv:2211.13752}, year={2022} }

@misc{baranchuk2021labelefficient, title={Label-Efficient Semantic Segmentation with Diffusion Models}, author={Dmitry Baranchuk and Ivan Rubachev and Andrey Voynov and Valentin Khrulkov and Artem Babenko}, year={2021}, eprint={2112.03126}, archivePrefix={arXiv}, primaryClass={cs.CV} }

sketch-guided-stable-diffusion's Issues

issue about text_encoder

Hello, I encountered bugs when I run the code. "text_encoder = CLIPTextModel.from_pretrained("openai/clip-vit-large-patch14").to(device)"
something like that:
..size mismatch for text_model.encoder.layers.11.mlp.fc2.bias: copying a param with shape torch.Size([768]) from checkpoint, the shape in current model is torch.Size([512]).
size mismatch for text_model.encoder.layers.11.layer_norm2.weight: copying a param with shape torch.Size([768]) from checkpoint, the shape in current model is torch.Size([512]).
size mismatch for text_model.encoder.layers.11.layer_norm2.bias: copying a param with shape torch.Size([768]) from checkpoint, the shape in current model is torch.Size([512]).
size mismatch for text_model.final_layer_norm.weight: copying a param with shape torch.Size([768]) from checkpoint, the shape in current model is torch.Size([512]).
size mismatch for text_model.final_layer_norm.bias: copying a param with shape torch.Size([768]) from checkpoint, the shape in current model is torch.Size([512]).

Could you help me fix it?

You seem to only share the pt model, I would like to consult where to put this model to work

I downloaded your pt model and put it in the embeddings folder, but it doesn't work, is it my operation wrong?

LGP Inference

Thanks for posting an amazing implementation.

I managed to train the LGP model, but having trouble in inferance, where we are supposed to guide reverse process of diffusion model using the gradient of the LGP.

Could you please provide the complete inferance code ? Or at least some crude outlines ?