OneDiffusion
This repository contains the OneDiffusion model presented in the paper One Diffusion to Generate Them All.
VAE model: SD3 VAE Text encoder: T5-xl
Installation
conda create -n onediffusion_env python=3.8 &&
conda activate onediffusion_env &&
pip install torch==2.3.1 torchvision==0.18.1 torchaudio==2.3.1 --index-url https://download.pytorch.org/whl/cu118 &&
pip install "git+https://github.com/facebookresearch/pytorch3d.git" &&
pip install -r requirements.txt
Quick start
Check inference.py for more detailed. For text-to-image, you can use below code snipe.
import torch
from onediffusion.diffusion.pipelines.onediffusion import OneDiffusionPipeline
device = torch.device('cuda:0')
pipeline = OneDiffusionPipeline.from_pretrained("lehduong/OneDiffusion").to(device=device, dtype=torch.bfloat16)
NEGATIVE_PROMPT = "monochrome, greyscale, low-res, bad anatomy, bad hands, text, error, missing fingers, extra digit, fewer digits, cropped, worst quality, low quality, normal quality, jpeg artifacts, signature, watermark, username, blurry, artist name, poorly drawn, bad anatomy, wrong anatomy, extra limb, missing limb, floating limbs, disconnected limbs, mutation, mutated, ugly, disgusting, blurry, amputation"
output = pipeline(
prompt="[[text2image]] A bipedal black cat wearing a huge oversized witch hat, a wizards robe, casting a spell,in an enchanted forest. The scene is filled with fireflies and moss on surrounding rocks and trees",
negative_prompt=NEGATIVE_PROMPT,
num_inference_steps=50,
guidance_scale=4,
height=1024,
width=1024,
)
output.images[0].save('text2image_output.jpg')
You can run the gradio demo with:
python gradio_demo.py --captioner molmo # [molmo, llava, disable]
The demo provides guidance and helps format the prompt properly for each task.
-
By default, it loads the quantized Molmo for captioning source images.
which significantly increases memory usage. You generally need a GPU with at least $40$ GB of memory to run the demo.You generally need a GPU with at least $21$ GB of memory to run the demo. -
Opting to use LLaVA can reduce this requirement to $\approx 27$ GB, though the resulting captions may be less accurate in some cases.
-
You can also manually provide the caption for each input image and run with
disablemode. In this mode, the returned caption is an empty string, but you should still press theGenerate Captionbutton so that the code formats the input text properly. The memory requirement for this mode is $\approx 12$ GB.
Note that the above required memory can change if you use higher resolution or more input images.
Qualitative Results
1. Text-to-Image
2. ID customization
3. Multiview generation
Single image to multiview:
Text to multiview:
4. Condition-to-Image and vice versa
5. Subject-driven generation
We finetuned the model on Subject-200K dataset (along with all other tasks) for additional 40k steps. The model is now capable of subject-driven generation.
6. Text-guide image editing
We finetuned the model on OmniEdit dataset for additional 30K steps.
7. Zero-shot Task combinations
We found that the model can handle multiple tasks in a zero-shot setting by combining condition images and task tokens without any fine-tuning, as shown in the examples below. However, its performance on these combined tasks might not be robust, and the model’s behavior may change if the order of task tokens or captions is altered. For example, when using both image inpainting and ID customization together, the target prompt and the caption of the masked image must be identical. If you plan to use such combinations, we recommend fine-tuning the model on these tasks to achieve better performance and simpler usage.
License
The model is trained on several non-commercially licensed datasets (e.g., DL3DV, Unsplash), thus, model weights are released under a CC BY-NC license as described in LICENSE.
Citation
@misc{le2024diffusiongenerate,
title={One Diffusion to Generate Them All},
author={Duong H. Le and Tuan Pham and Sangho Lee and Christopher Clark and Aniruddha Kembhavi and Stephan Mandt and Ranjay Krishna and Jiasen Lu},
year={2024},
eprint={2411.16318},
archivePrefix={arXiv},
primaryClass={cs.CV},
url={https://arxiv.org/abs/2411.16318},
}