|
--- |
|
base_model: stabilityai/stable-diffusion-xl-base-1.0 |
|
library_name: diffusers |
|
license: openrail++ |
|
tags: |
|
- text-to-image |
|
- text-to-image |
|
- diffusers-training |
|
- diffusers |
|
- lora |
|
- template:sd-lora |
|
- stable-diffusion-xl |
|
- stable-diffusion-xl-diffusers |
|
- science |
|
- materiomics |
|
- bio-inspired |
|
- materials science |
|
instance_prompt: <leaf microstructure> |
|
widget: [] |
|
--- |
|
|
|
# SDXL Fine-tuned with Leaf Images |
|
|
|
DreamBooth is an advanced technique designed for fine-tuning text-to-image diffusion models to generate personalized images of specific subjects. By leveraging a few reference images (around 5 or so), DreamBooth integrates unique visual features of the subject into the model's output domain. |
|
|
|
This is achieved by binding a unique identifier "\<..IDENTIFIER..\>", such as \<leaf microstructure\> in this work, to the subject. An optional class-specific prior preservation loss can be used to maintain high fidelity and contextual diversity. The result is a model capable of synthesizing novel, photorealistic images of the subject in various scenes, poses, and lighting conditions, guided by text prompts. In this project, DreamBooth has been applied to render images with specific biological patterns, making it ideal for applications in materials science and engineering where accurate representation of biological material microstructures is crucial. |
|
|
|
For example, an original prompt might be: "a vase with intricate patterns, high quality." With the fine-tuned model, using the unique identifier, the prompt becomes: "a vase that resembles a \<leaf microstructure\>, high quality." This allows the model to generate images that specifically incorporate the desired biological pattern. |
|
|
|
## Model description |
|
|
|
These are LoRA adaption weights for the SDXL-base-1.0 model (```stabilityai/stable-diffusion-xl-base-1.0```). |
|
|
|
## Trigger keywords |
|
|
|
The following images were used during fine-tuning using the keyword \<leaf microstructure\>: |
|
|
|
 |
|
|
|
You should use \<leaf microstructure\> to trigger this feature during image generation. |
|
|
|
[](https://colab.research.google.com/#fileId=https://huggingface.co/lamm-mit/SDXL-leaf-inspired/blob/main/SDXL_leaf_inspired_inference.ipynb) |
|
|
|
## How to use |
|
|
|
Defining some helper functions: |
|
|
|
```python |
|
from diffusers import DiffusionPipeline |
|
import torch |
|
import os |
|
from datetime import datetime |
|
from PIL import Image |
|
|
|
def generate_filename(base_name, extension=".png"): |
|
timestamp = datetime.now().strftime("%Y%m%d_%H%M%S") |
|
return f"{base_name}_{timestamp}{extension}" |
|
|
|
def save_image(image, directory, base_name="image_grid"): |
|
|
|
filename = generate_filename(base_name) |
|
file_path = os.path.join(directory, filename) |
|
image.save(file_path) |
|
print(f"Image saved as {file_path}") |
|
|
|
def image_grid(imgs, rows, cols, save=True, save_dir='generated_images', base_name="image_grid", |
|
save_individual_files=False): |
|
|
|
if not os.path.exists(save_dir): |
|
os.makedirs(save_dir) |
|
|
|
assert len(imgs) == rows * cols |
|
|
|
w, h = imgs[0].size |
|
grid = Image.new('RGB', size=(cols * w, rows * h)) |
|
grid_w, grid_h = grid.size |
|
|
|
for i, img in enumerate(imgs): |
|
grid.paste(img, box=(i % cols * w, i // cols * h)) |
|
if save_individual_files: |
|
save_image(img, save_dir, base_name=base_name+f'_{i}-of-{len(imgs)}_') |
|
|
|
if save and save_dir: |
|
save_image(grid, save_dir, base_name) |
|
|
|
return grid |
|
``` |
|
|
|
### Text-to-image |
|
|
|
Model loading: |
|
|
|
```python |
|
|
|
import torch |
|
from diffusers import DiffusionPipeline, AutoencoderKL |
|
|
|
repo_id='lamm-mit/SDXL-leaf-inspired' |
|
|
|
vae = AutoencoderKL.from_pretrained("madebyollin/sdxl-vae-fp16-fix", torch_dtype=torch.float16) |
|
base = DiffusionPipeline.from_pretrained( |
|
"stabilityai/stable-diffusion-xl-base-1.0", |
|
vae=vae, |
|
torch_dtype=torch.float16, |
|
variant="fp16", |
|
use_safetensors=True |
|
) |
|
base.load_lora_weights(repo_id) |
|
_ = base.to("cuda") |
|
|
|
refiner = DiffusionPipeline.from_pretrained( |
|
"stabilityai/stable-diffusion-xl-refiner-1.0", |
|
text_encoder_2=base.text_encoder_2, |
|
vae=base.vae, |
|
torch_dtype=torch.float16, |
|
use_safetensors=True, |
|
variant="fp16", |
|
) |
|
refiner.to("cuda") |
|
``` |
|
|
|
Image generation: |
|
|
|
```python |
|
prompt = "a vase that resembles a <leaf microstructure>, high quality" |
|
|
|
num_samples = 4 |
|
num_rows = 4 |
|
guidance_scale = 15 |
|
|
|
all_images = [] |
|
|
|
for _ in range(num_rows): |
|
# Define how many steps and what % of steps to be run on each experts (80/20) |
|
n_steps = 25 |
|
high_noise_frac = 0.8 |
|
|
|
# run both experts |
|
image = base( |
|
prompt=prompt, |
|
num_inference_steps=n_steps, guidance_scale=guidance_scale, |
|
denoising_end=high_noise_frac,num_images_per_prompt=num_samples, |
|
output_type="latent", |
|
).images |
|
image = refiner( |
|
prompt=prompt, |
|
num_inference_steps=n_steps, guidance_scale=guidance_scale, |
|
denoising_start=high_noise_frac,num_images_per_prompt=num_samples, |
|
image=image, |
|
).images |
|
|
|
all_images.extend(image) |
|
|
|
grid = image_grid(all_images, num_rows, num_samples, |
|
save_individual_files=True, |
|
) |
|
grid |
|
``` |
|
|
|
|
|
 |
|
|
|
## Fine-tuning script |
|
|
|
Download this script: [SDXL DreamBooth-LoRA_Fine-Tune.ipynb](https://huggingface.co/lamm-mit/SDXL-leaf-inspired/resolve/main/SDXL_DreamBooth_LoRA_Fine-Tune.ipynb) |
|
|
|
You need to create a local folder ```leaf_concept_dir_SDXL``` and add the leaf images (provided in this repository, see subfolder), like so: |
|
|
|
```raw |
|
mkdir leaf_concept_dir_SDXL |
|
cd leaf_concept_dir_SDXL |
|
wget https://huggingface.co/lamm-mit/SDXL-leaf-inspired/resolve/main/leaf_concept_dir_SDXL/0.jpeg |
|
wget https://huggingface.co/lamm-mit/SDXL-leaf-inspired/resolve/main/leaf_concept_dir_SDXL/1.jpeg |
|
wget https://huggingface.co/lamm-mit/SDXL-leaf-inspired/resolve/main/leaf_concept_dir_SDXL/2.jpeg |
|
wget https://huggingface.co/lamm-mit/SDXL-leaf-inspired/resolve/main/leaf_concept_dir_SDXL/3.jpeg |
|
wget https://huggingface.co/lamm-mit/SDXL-leaf-inspired/resolve/main/leaf_concept_dir_SDXL/87.jpg |
|
wget https://huggingface.co/lamm-mit/SDXL-leaf-inspired/resolve/main/leaf_concept_dir_SDXL/87.jpg |
|
wget https://huggingface.co/lamm-mit/SDXL-leaf-inspired/resolve/main/leaf_concept_dir_SDXL/88.jpg |
|
wget https://huggingface.co/lamm-mit/SDXL-leaf-inspired/resolve/main/leaf_concept_dir_SDXL/90.jpg |
|
wget https://huggingface.co/lamm-mit/SDXL-leaf-inspired/resolve/main/leaf_concept_dir_SDXL/91.jpg |
|
wget https://huggingface.co/lamm-mit/SDXL-leaf-inspired/resolve/main/leaf_concept_dir_SDXL/94.jpg |
|
cd .. |
|
``` |
|
|
|
The code will automatically download the training script. |
|
|
|
The training script can handle custom prompts associated with each image, which are generated using BLIP. |
|
|
|
For instance, for the images used here, they are: |
|
|
|
```raw |
|
{"file_name": "0.jpeg", "prompt": "<leaf microstructure>, a close up of a green plant with a lot of small holes"} |
|
{"file_name": "1.jpeg", "prompt": "<leaf microstructure>, a close up of a leaf with a small insect on it"} |
|
{"file_name": "2.jpeg", "prompt": "<leaf microstructure>, a close up of a plant with a lot of green leaves"} |
|
{"file_name": "3.jpeg", "prompt": "<leaf microstructure>, a close up of a leaf with a yellow substance in it"} |
|
{"file_name": "87.jpg", "prompt": "<leaf microstructure>, a close up of a green plant with a yellow light"} |
|
{"file_name": "88.jpg", "prompt": "<leaf microstructure>, a close up of a green plant with a white center"} |
|
{"file_name": "90.jpg", "prompt": "<leaf microstructure>, arafed leaf with a white line on the center"} |
|
{"file_name": "91.jpg", "prompt": "<leaf microstructure>, arafed image of a green leaf with a white spot"} |
|
{"file_name": "92.jpg", "prompt": "<leaf microstructure>, a close up of a leaf with a yellow light shining through it"} |
|
{"file_name": "94.jpg", "prompt": "<leaf microstructure>, arafed image of a green plant with a yellow cross"} |
|
``` |
|
|
|
Training then proceeds as: |
|
|
|
```python |
|
HF_username = 'lamm-mit' |
|
|
|
pretrained_model_name_or_path="stabilityai/stable-diffusion-xl-base-1.0" |
|
pretrained_vae_model_name_or_path="madebyollin/sdxl-vae-fp16-fix" |
|
|
|
instance_prompt ="<leaf microstructure>" |
|
instance_data_dir = "./leaf_concept_dir_SDXL/" |
|
|
|
val_prompt = "a vase that resembles a <leaf microstructure>, high quality" |
|
val_epochs = 100 |
|
|
|
instance_output_dir="leaf_LoRA_SDXL_V10" #for checkpointing |
|
``` |
|
|
|
Dataset generatio with custom per-image captions |
|
```python |
|
import requests |
|
from transformers import AutoProcessor, BlipForConditionalGeneration |
|
import torch |
|
import glob |
|
from PIL import Image |
|
import json |
|
|
|
device = "cuda" if torch.cuda.is_available() else "cpu" |
|
|
|
# load the processor and the captioning model |
|
blip_processor = AutoProcessor.from_pretrained("Salesforce/blip-image-captioning-large") |
|
blip_model = BlipForConditionalGeneration.from_pretrained("Salesforce/blip-image-captioning-large",torch_dtype=torch.float16).to(device) |
|
|
|
# captioning utility |
|
def caption_images(input_image): |
|
inputs = blip_processor(images=input_image, return_tensors="pt").to(device, torch.float16) |
|
pixel_values = inputs.pixel_values |
|
|
|
generated_ids = blip_model.generate(pixel_values=pixel_values, max_length=50) |
|
generated_caption = blip_processor.batch_decode(generated_ids, skip_special_tokens=True)[0] |
|
return generated_caption |
|
|
|
caption_prefix = f"{instance_prompt}, " |
|
with open(f'{instance_data_dir}metadata.jsonl', 'w') as outfile: |
|
for img in imgs_and_paths: |
|
caption = caption_prefix + caption_images(img[1]).split("\n")[0] |
|
entry = {"file_name":img[0].split("/")[-1], "prompt": caption} |
|
json.dump(entry, outfile) |
|
outfile.write('\n') |
|
``` |
|
This produces a JSON file in the ```instance_data_dir``` directory: |
|
|
|
```raw |
|
{"file_name": "0.jpeg", "prompt": "<leaf microstructure>, a close up of a green plant with a lot of small holes"} |
|
{"file_name": "1.jpeg", "prompt": "<leaf microstructure>, a close up of a leaf with a small insect on it"} |
|
{"file_name": "2.jpeg", "prompt": "<leaf microstructure>, a close up of a plant with a lot of green leaves"} |
|
{"file_name": "3.jpeg", "prompt": "<leaf microstructure>, a close up of a leaf with a yellow substance in it"} |
|
{"file_name": "87.jpg", "prompt": "<leaf microstructure>, a close up of a green plant with a yellow light"} |
|
{"file_name": "88.jpg", "prompt": "<leaf microstructure>, a close up of a green plant with a white center"} |
|
{"file_name": "90.jpg", "prompt": "<leaf microstructure>, arafed leaf with a white line on the center"} |
|
{"file_name": "91.jpg", "prompt": "<leaf microstructure>, arafed image of a green leaf with a white spot"} |
|
{"file_name": "92.jpg", "prompt": "<leaf microstructure>, a close up of a leaf with a yellow light shining through it"} |
|
{"file_name": "94.jpg", "prompt": "<leaf microstructure>, arafed image of a green plant with a yellow cross"} |
|
``` |
|
|
|
```python |
|
!accelerate launch train_dreambooth_lora_sdxl.py \ |
|
--pretrained_model_name_or_path="{pretrained_model_name_or_path}" \ |
|
--pretrained_vae_model_name_or_path="{pretrained_vae_model_name_or_path}"\ |
|
--dataset_name="{instance_data_dir}" \ |
|
--output_dir="{instance_output_dir}" \ |
|
--caption_column="prompt"\ |
|
--mixed_precision="fp16" \ |
|
--instance_prompt="{instance_prompt}" \ |
|
--validation_prompt="{val_prompt}" \ |
|
--validation_epochs="{val_epochs}" \ |
|
--resolution=1024 \ |
|
--train_batch_size=1 \ |
|
--gradient_accumulation_steps=3 \ |
|
--gradient_checkpointing \ |
|
--learning_rate=1e-4 \ |
|
--snr_gamma=5.0 \ |
|
--lr_scheduler="constant" \ |
|
--lr_warmup_steps=0 \ |
|
--mixed_precision="fp16" \ |
|
--use_8bit_adam \ |
|
--max_train_steps=500 \ |
|
--checkpointing_steps=500 \ |
|
--seed="0" |
|
``` |
|
|
|
### With prior preservation |
|
|
|
Set `--with_prior_preservation` flag to include prior preservation. In this case you must specify `--class_data_dir` (directory with class images) and `--class_prompt` (class prompt). You should also set `--num_class_images` to specify how many class preservation images you want to use. Either place them in the directory (specified via `--class_data_dir`) or the code with auto-generate them based off the base model. You can also provide a few yourself and let the code generate the remaining ones. |
|
|
|
An example is provided below, commented out. The code that will run here will NOT use prior preservation. |
|
|
|
Some other useful parameters that can be set include: |
|
|
|
--rank: LoRA adapter rank (LoRA alpha will be set identical to rank) |
|
--use_dora: Set if you want to use DORA |
|
|
|
Type ```python train_dreambooth_lora_sdxl.py``` to get a full list of parameters |
|
|
|
```python |
|
instance_data_dir = 'local_instance_data_dir' |
|
class_prompt = 'a prompt that describes the images in the directory local_instance_data_dir' |
|
num_class_images = 10 #how many images you want in this class |
|
|
|
!\accelerate launch train_dreambooth_lora_sdxl.py \ |
|
--pretrained_model_name_or_path="{pretrained_model_name_or_path}" \ |
|
--pretrained_vae_model_name_or_path="{pretrained_vae_model_name_or_path}"\ |
|
--dataset_name="{instance_data_dir}" \ |
|
--class_prompt="{class_prompt}" \ |
|
--num_class_images={num_class_images} \ |
|
--with_prior_preservation \ |
|
--class_data_dir="{class_data_dir}" \ |
|
--output_dir="{instance_output_dir}" \ |
|
--caption_column="prompt"\ |
|
--mixed_precision="fp16" \ |
|
--instance_prompt="{instance_prompt}" \ |
|
--validation_prompt="{val_prompt}" \ |
|
--validation_epochs={val_epochs} \ |
|
--resolution=1024 \ |
|
--train_batch_size=1 \ |
|
--gradient_accumulation_steps=4 \ |
|
--gradient_checkpointing \ |
|
--learning_rate=1e-4 \ |
|
--snr_gamma=5.0 \ |
|
--lr_scheduler="constant" \ |
|
--lr_warmup_steps=0 \ |
|
--mixed_precision="fp16" \ |
|
--use_8bit_adam \ |
|
--max_train_steps=500 \ |
|
--checkpointing_steps=500 \ |
|
--seed="0" |
|
``` |
