{"id":13488772,"url":"https://github.com/TencentARC/BrushNet","last_synced_at":"2025-03-28T01:37:52.930Z","repository":{"id":227271378,"uuid":"769852109","full_name":"TencentARC/BrushNet","owner":"TencentARC","description":"[ECCV 2024] The official implementation of paper \"BrushNet: A Plug-and-Play Image Inpainting Model with Decomposed Dual-Branch Diffusion\"","archived":false,"fork":false,"pushed_at":"2024-07-17T11:17:08.000Z","size":39036,"stargazers_count":1391,"open_issues_count":44,"forks_count":115,"subscribers_count":41,"default_branch":"main","last_synced_at":"2024-10-23T01:29:16.790Z","etag":null,"topics":["diffusion","diffusion-models","eccv","eccv2024","image-inpainting","text-to-image"],"latest_commit_sha":null,"homepage":"https://tencentarc.github.io/BrushNet/","language":"Python","has_issues":true,"has_wiki":null,"has_pages":null,"mirror_url":null,"source_name":null,"license":"other","status":null,"scm":"git","pull_requests_enabled":true,"icon_url":"https://github.com/TencentARC.png","metadata":{"files":{"readme":"README.md","changelog":null,"contributing":"CONTRIBUTING.md","funding":null,"license":"LICENSE","code_of_conduct":"CODE_OF_CONDUCT.md","threat_model":null,"audit":null,"citation":"CITATION.cff","codeowners":null,"security":null,"support":null,"governance":null,"roadmap":null,"authors":null,"dei":null,"publiccode":null,"codemeta":null}},"created_at":"2024-03-10T08:44:03.000Z","updated_at":"2024-10-21T15:22:52.000Z","dependencies_parsed_at":"2024-03-22T14:15:41.460Z","dependency_job_id":"0eb88b67-e1f1-429b-8d75-5667c232bb67","html_url":"https://github.com/TencentARC/BrushNet","commit_stats":null,"previous_names":["tencentarc/brushnet"],"tags_count":0,"template":false,"template_full_name":null,"repository_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/TencentARC%2FBrushNet","tags_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/TencentARC%2FBrushNet/tags","releases_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/TencentARC%2FBrushNet/releases","manifests_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/TencentARC%2FBrushNet/manifests","owner_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/owners/TencentARC","download_url":"https://codeload.github.com/TencentARC/BrushNet/tar.gz/refs/heads/main","host":{"name":"GitHub","url":"https://github.com","kind":"github","repositories_count":222061537,"owners_count":16924686,"icon_url":"https://github.com/github.png","version":null,"created_at":"2022-05-30T11:31:42.601Z","updated_at":"2022-07-04T15:15:14.044Z","host_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub","repositories_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories","repository_names_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repository_names","owners_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/owners"}},"keywords":["diffusion","diffusion-models","eccv","eccv2024","image-inpainting","text-to-image"],"created_at":"2024-07-31T18:01:21.524Z","updated_at":"2025-03-28T01:37:52.924Z","avatar_url":"https://github.com/TencentARC.png","language":"Python","readme":"# BrushNet\n\nThis repository contains the implementation of the ECCV2024 paper \"BrushNet: A Plug-and-Play Image Inpainting Model with Decomposed Dual-Branch Diffusion\"\n\nKeywords: Image Inpainting, Diffusion Models, Image Generation\n\n\u003e [Xuan Ju](https://github.com/juxuan27)\u003csup\u003e12\u003c/sup\u003e, [Xian Liu](https://alvinliu0.github.io/)\u003csup\u003e12\u003c/sup\u003e, [Xintao Wang](https://xinntao.github.io/)\u003csup\u003e1*\u003c/sup\u003e, [Yuxuan Bian](https://scholar.google.com.hk/citations?user=HzemVzoAAAAJ\u0026hl=zh-CN\u0026oi=ao)\u003csup\u003e2\u003c/sup\u003e, [Ying Shan](https://www.linkedin.com/in/YingShanProfile/)\u003csup\u003e1\u003c/sup\u003e, [Qiang Xu](https://cure-lab.github.io/)\u003csup\u003e2*\u003c/sup\u003e\u003cbr\u003e\n\u003e \u003csup\u003e1\u003c/sup\u003eARC Lab, Tencent PCG \u003csup\u003e2\u003c/sup\u003eThe Chinese University of Hong Kong \u003csup\u003e*\u003c/sup\u003eCorresponding Author\n\n\n\u003cp align=\"center\"\u003e\n  \u003ca href=\"https://tencentarc.github.io/BrushNet/\"\u003e🌐Project Page\u003c/a\u003e |\n  \u003ca href=\"https://arxiv.org/abs/2403.06976\"\u003e📜Arxiv\u003c/a\u003e |\n  \u003ca href=\"https://forms.gle/9TgMZ8tm49UYsZ9s5\"\u003e🗄️Data\u003c/a\u003e |\n  \u003ca href=\"https://drive.google.com/file/d/1IkEBWcd2Fui2WHcckap4QFPcCI0gkHBh/view\"\u003e📹Video\u003c/a\u003e |\n  \u003ca href=\"https://huggingface.co/spaces/TencentARC/BrushNet\"\u003e🤗Hugging Face Demo\u003c/a\u003e |\n\u003c/p\u003e\n\n\n\n**📖 Table of Contents**\n\n\n- [BrushNet](#brushnet)\n  - [🔥 Update Log](#-update-log)\n  - [TODO](#todo)\n  - [🛠️ Method Overview](#️-method-overview)\n  - [🚀 Getting Started](#-getting-started)\n    - [Environment Requirement 🌍](#environment-requirement-)\n    - [Data Download ⬇️](#data-download-️)\n  - [🏃🏼 Running Scripts](#-running-scripts)\n    - [Training 🤯](#training-)\n    - [Inference 📜](#inference-)\n    - [Evaluation 📏](#evaluation-)\n  - [🤝🏼 Cite Us](#-cite-us)\n  - [💖 Acknowledgement](#-acknowledgement)\n\n\n\n## 🔥 Update Log\n- [2024/12/17] 📢 📢  [BrushEdit](https://github.com/TencentARC/BrushEdit) are released, an efficient, white-box, free-form image editing tool powered by LLM-agents and an all-in-one inpainting model.\n- [2024/12/17] 📢 📢 [BrushNetX](https://huggingface.co/TencentARC/BrushEdit/tree/main/brushnetX) (Stronger BrushNet) models are released.\n\n## TODO\n\n- [x] Release trainig and inference code\n- [x] Release checkpoint (sdv1.5)\n- [x] Release checkpoint (sdxl). Sadly, I only have V100 for training this checkpoint, which can only train with a batch size of 1 with a slow speed. The current ckpt is only trained for a small step number thus perform not well. But fortunately, [yuanhang](https://github.com/yuanhangio) volunteer to help training a better version. Please stay tuned! Thank [yuanhang](https://github.com/yuanhangio) for his effort!\n- [x] Release evluation code\n- [x] Release gradio demo\n- [x] Release comfyui demo. Thank [nullquant](https://github.com/nullquant) ([ConfyUI-BrushNet](https://github.com/nullquant/ComfyUI-BrushNet)) and [kijai](https://github.com/kijai) ([ComfyUI-BrushNet-Wrapper](https://github.com/kijai/ComfyUI-BrushNet-Wrapper)) for helping!\n- [x] Release [trainig data](https://huggingface.co/datasets/random123123/BrushData). Thank [random123123](https://huggingface.co/random123123) for helping!\n- [x] We use BrushNet to participate in CVPR2024 GenAI Media Generation Challenge Workshop and get top prize! The solution is provided in [InstructionGuidedEditing](InstructionGuidedEditing)\n- [x] Release a new version of checkpoint (sdxl).\n\n## 🛠️ Method Overview\n\nBrushNet is a diffusion-based text-guided image inpainting model that can be plug-and-play into any pre-trained diffusion model. Our architectural design incorporates two key insights: (1) dividing the masked image features and noisy latent reduces the model's learning load, and (2) leveraging dense per-pixel control over the entire pre-trained model enhances its suitability for image inpainting tasks. More analysis can be found in the main paper.\n\n![](examples/brushnet/src/model.png)\n\n\n\n## 🚀 Getting Started\n\n### Environment Requirement 🌍\n\nBrushNet has been implemented and tested on Pytorch 1.12.1 with python 3.9.\n\nClone the repo:\n\n```\ngit clone https://github.com/TencentARC/BrushNet.git\n```\n\nWe recommend you first use `conda` to create virtual environment, and install `pytorch` following [official instructions](https://pytorch.org/). For example:\n\n\n```\nconda create -n diffusers python=3.9 -y\nconda activate diffusers\npython -m pip install --upgrade pip\npip install torch==1.12.1+cu116 torchvision==0.13.1+cu116 torchaudio==0.12.1 --extra-index-url https://download.pytorch.org/whl/cu116\n```\n\nThen, you can install diffusers (implemented in this repo) with:\n\n```\npip install -e .\n```\n\nAfter that, you can install required packages thourgh:\n\n```\ncd examples/brushnet/\npip install -r requirements.txt\n```\n\n### Data Download ⬇️\n\n\n**Dataset**\n\nYou can download the BrushData and BrushBench [here](https://forms.gle/9TgMZ8tm49UYsZ9s5) (as well as the EditBench we re-processed), which are used for training and testing the BrushNet. By downloading the data, you are agreeing to the terms and conditions of the license. The data structure should be like:\n\n```\n|-- data\n    |-- BrushData\n        |-- 00200.tar\n        |-- 00201.tar\n        |-- ...\n    |-- BrushDench\n        |-- images\n        |-- mapping_file.json\n    |-- EditBench\n        |-- images\n        |-- mapping_file.json\n```\n\n\nNoted: *We only provide a part of the BrushData in google drive due to the space limit. [random123123](https://huggingface.co/random123123) has helped upload a full dataset on hugging face [here](https://huggingface.co/datasets/random123123/BrushData). Thank for his help!*\n\n\n**Checkpoints**\n\nCheckpoints of BrushNet can be downloaded from [here](https://drive.google.com/drive/folders/1fqmS1CEOvXCxNWFrsSYd_jHYXxrydh1n?usp=drive_link). The ckpt folder contains \n\n- BrushNet pretrained checkpoints for Stable Diffusion v1.5 (`segmentation_mask_brushnet_ckpt` and `random_mask_brushnet_ckpt`)\n- pretrinaed Stable Diffusion v1.5 checkpoint (e.g., realisticVisionV60B1_v51VAE from [Civitai](https://civitai.com/)). You can use `scripts/convert_original_stable_diffusion_to_diffusers.py` to process other models downloaded from Civitai. \n- BrushNet pretrained checkpoints for Stable Diffusion XL (`segmentation_mask_brushnet_ckpt_sdxl_v1` and `random_mask_brushnet_ckpt_sdxl_v0`).  A better version will be shortly released by [yuanhang](https://github.com/yuanhangio). Please stay tuned!\n- pretrinaed Stable Diffusion XL checkpoint (e.g., juggernautXL_juggernautX from [Civitai](https://civitai.com/)). You can use `StableDiffusionXLPipeline.from_single_file(\"path of safetensors\").save_pretrained(\"path to save\",safe_serialization=False)` to process other models downloaded from Civitai. \n\n\n\nThe data structure should be like:\n\n\n```\n|-- data\n    |-- BrushData\n    |-- BrushDench\n    |-- EditBench\n    |-- ckpt\n        |-- realisticVisionV60B1_v51VAE\n            |-- model_index.json\n            |-- vae\n            |-- ...\n        |-- segmentation_mask_brushnet_ckpt\n        |-- segmentation_mask_brushnet_ckpt_sdxl_v0\n        |-- random_mask_brushnet_ckpt\n        |-- random_mask_brushnet_ckpt_sdxl_v0\n        |-- ...\n```\n\nThe checkpoint in `segmentation_mask_brushnet_ckpt` and `segmentation_mask_brushnet_ckpt_sdxl_v0` provide checkpoints trained on BrushData, which has segmentation prior (mask are with the same shape of objects). The `random_mask_brushnet_ckpt` and `random_mask_brushnet_ckpt_sdxl` provide a more general ckpt for random mask shape.\n\n## 🏃🏼 Running Scripts\n\n\n### Training 🤯\n\nYou can train with segmentation mask using the script:\n\n```\n# sd v1.5\naccelerate launch examples/brushnet/train_brushnet.py \\\n--pretrained_model_name_or_path runwayml/stable-diffusion-v1-5 \\\n--output_dir runs/logs/brushnet_segmentationmask \\\n--train_data_dir data/BrushData \\\n--resolution 512 \\\n--learning_rate 1e-5 \\\n--train_batch_size 2 \\\n--tracker_project_name brushnet \\\n--report_to tensorboard \\\n--resume_from_checkpoint latest \\\n--validation_steps 300\n--checkpointing_steps 10000 \n\n# sdxl\naccelerate launch examples/brushnet/train_brushnet_sdxl.py \\\n--pretrained_model_name_or_path stabilityai/stable-diffusion-xl-base-1.0 \\\n--output_dir runs/logs/brushnetsdxl_segmentationmask \\\n--train_data_dir data/BrushData \\\n--resolution 1024 \\\n--learning_rate 1e-5 \\\n--train_batch_size 1 \\\n--gradient_accumulation_steps 4 \\\n--tracker_project_name brushnet \\\n--report_to tensorboard \\\n--resume_from_checkpoint latest \\\n--validation_steps 300 \\\n--checkpointing_steps 10000 \n```\n\nTo use custom dataset, you can process your own data to the format of BrushData and revise `--train_data_dir`.\n\nYou can train with random mask using the script (by adding `--random_mask`):\n\n```\n# sd v1.5\naccelerate launch examples/brushnet/train_brushnet.py \\\n--pretrained_model_name_or_path runwayml/stable-diffusion-v1-5 \\\n--output_dir runs/logs/brushnet_randommask \\\n--train_data_dir data/BrushData \\\n--resolution 512 \\\n--learning_rate 1e-5 \\\n--train_batch_size 2 \\\n--tracker_project_name brushnet \\\n--report_to tensorboard \\\n--resume_from_checkpoint latest \\\n--validation_steps 300 \\\n--random_mask\n\n# sdxl\naccelerate launch examples/brushnet/train_brushnet_sdxl.py \\\n--pretrained_model_name_or_path stabilityai/stable-diffusion-xl-base-1.0 \\\n--output_dir runs/logs/brushnetsdxl_randommask \\\n--train_data_dir data/BrushData \\\n--resolution 1024 \\\n--learning_rate 1e-5 \\\n--train_batch_size 1 \\\n--gradient_accumulation_steps 4 \\\n--tracker_project_name brushnet \\\n--report_to tensorboard \\\n--resume_from_checkpoint latest \\\n--validation_steps 300 \\\n--checkpointing_steps 10000 \\\n--random_mask\n```\n\n\n\n### Inference 📜\n\nYou can inference with the script:\n\n```\n# sd v1.5\npython examples/brushnet/test_brushnet.py\n# sdxl\npython examples/brushnet/test_brushnet_sdxl.py\n```\n\nSince BrushNet is trained on Laion, it can only guarantee the performance on general scenarios. We recommend you train on your own data (e.g., product exhibition, virtual try-on) if you have high-quality industrial application requirements. We would also be appreciate if you would like to contribute your trained model!\n\nYou can also inference through gradio demo:\n\n```\n# sd v1.5\npython examples/brushnet/app_brushnet.py\n```\n\n\n### Evaluation 📏\n\nYou can evaluate using the script:\n\n```\npython examples/brushnet/evaluate_brushnet.py \\\n--brushnet_ckpt_path data/ckpt/segmentation_mask_brushnet_ckpt \\\n--image_save_path runs/evaluation_result/BrushBench/brushnet_segmask/inside \\\n--mapping_file data/BrushBench/mapping_file.json \\\n--base_dir data/BrushBench \\\n--mask_key inpainting_mask\n```\n\nThe `--mask_key` indicates which kind of mask to use, `inpainting_mask` for inside inpainting and `outpainting_mask` for outside inpainting. The evaluation results (images and metrics) will be saved in `--image_save_path`. \n\n\n\n*Noted that you need to ignore the nsfw detector in `src/diffusers/pipelines/brushnet/pipeline_brushnet.py#1261` to get the correct evaluation results. Moreover, we find different machine may generate different images, thus providing the results on our machine [here](https://drive.google.com/drive/folders/1dK3oIB2UvswlTtnIS1iHfx4s57MevWdZ?usp=sharing).*\n\n\n## 🤝🏼 Cite Us\n\n```\n@misc{ju2024brushnet,\n  title={BrushNet: A Plug-and-Play Image Inpainting Model with Decomposed Dual-Branch Diffusion}, \n  author={Xuan Ju and Xian Liu and Xintao Wang and Yuxuan Bian and Ying Shan and Qiang Xu},\n  year={2024},\n  eprint={2403.06976},\n  archivePrefix={arXiv},\n  primaryClass={cs.CV}\n}\n```\n\n\n## 💖 Acknowledgement\n\u003cspan id=\"acknowledgement\"\u003e\u003c/span\u003e\n\nOur code is modified based on [diffusers](https://github.com/huggingface/diffusers), thanks to all the contributors!\n\n","funding_links":[],"categories":["SD-inpaint","Repos"],"sub_categories":[],"project_url":"https://awesome.ecosyste.ms/api/v1/projects/github.com%2FTencentARC%2FBrushNet","html_url":"https://awesome.ecosyste.ms/projects/github.com%2FTencentARC%2FBrushNet","lists_url":"https://awesome.ecosyste.ms/api/v1/projects/github.com%2FTencentARC%2FBrushNet/lists"}