|
--- |
|
base_model: stabilityai/stable-diffusion-2 |
|
library_name: diffusers |
|
license: openrail++ |
|
tags: |
|
- text-to-image |
|
- text-to-image |
|
- diffusers-training |
|
- diffusers |
|
- stable-diffusion-2 |
|
- stable-diffusion-2-diffusers |
|
- science |
|
- materiomics |
|
- bio-inspired |
|
- materials science |
|
- text-to-3D |
|
- text-to-STL |
|
- text-t-mesh |
|
- additive manufacturing |
|
- 3D |
|
- 3D printing |
|
instance_prompt: <leaf microstructure> |
|
widget: [] |
|
--- |
|
|
|
# Stable Diffusion 2.x Fine-tuned with Leaf Images: Text-to-image and text-to-3D |
|
|
|
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 fine-tuned weights for the ```stabilityai/stable-diffusion-2``` model. This is a full fine-tune of the model using DreamBooth. |
|
|
|
|
|
## Trigger keywords |
|
|
|
The following image were used during fine-tuning using the keyword \<leaf microstructure\>: |
|
|
|
[](https://colab.research.google.com/#fileId=https://huggingface.co/lamm-mit/SD2x-leaf-inspired/blob/main/SD2x_leaf_inspired_inference.ipynb) |
|
|
|
Please use \<leaf microstructure\> to trigger the image generation. |
|
|
|
## 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 StableDiffusionPipeline, DPMSolverMultistepScheduler |
|
|
|
repo_id='lamm-mit/SD2x-leaf-inspired' |
|
|
|
pipe = StableDiffusionPipeline.from_pretrained(repo_id, |
|
scheduler = DPMSolverMultistepScheduler.from_pretrained(repo_id, subfolder="scheduler"), |
|
torch_dtype=torch.float16, |
|
).to("cuda") |
|
|
|
``` |
|
|
|
Image generation: |
|
|
|
```python |
|
prompt = "a vase that resembles a <leaf microstructure>, high quality" |
|
num_samples = 4 |
|
num_rows = 4 |
|
|
|
all_images = [] |
|
for _ in range(num_rows): |
|
images = pipe(prompt, num_images_per_prompt=num_samples, num_inference_steps=50, guidance_scale=15).images |
|
all_images.extend(images) |
|
|
|
grid = image_grid(all_images, num_rows, num_samples) |
|
grid |
|
``` |
|
 |
|
|
|
### Image-to-Image |
|
|
|
The model can be used also for image-to-image tasks. For instance, we can first generate a draft image and then further modify it. |
|
|
|
Create draft image: |
|
``` |
|
prompt = "a vase that resembles a <leaf microstructure>, high quality" |
|
num_samples = 4 |
|
num_rows = 1 |
|
|
|
all_images = [] |
|
for _ in range(num_rows): |
|
images = pipe(prompt, num_images_per_prompt=num_samples, num_inference_steps=50, guidance_scale=15).images |
|
all_images.extend(images) |
|
|
|
grid = image_grid(all_images, num_rows, num_samples, save_individual_files=True) |
|
grid |
|
``` |
|
|
|
 |
|
|
|
Now we use one of the images (second from left) and modify it using the image-to-image pipeline. You can get the image as follows (if you run the generate code yourself, the generated images will be in the subdirectory ```generated_images```): |
|
|
|
``` |
|
wget https://huggingface.co/lamm-mit/SD2x-leaf-inspired/resolve/main/image_grid_1-of-4__20240722_144702.png |
|
``` |
|
|
|
 |
|
|
|
Now, generate: |
|
``` |
|
fname='image_grid_1-of-4__20240722_144702.png' |
|
init_image = Image.open(fname).convert("RGB") |
|
init_image = init_image.resize((768, 768)) |
|
|
|
prompt = "A vase made out of a spongy material, high quality photograph, full frame." |
|
num_samples = 4 |
|
num_rows = 1 |
|
|
|
all_images = [] |
|
for _ in range(num_rows): |
|
images = img2imgpipe(prompt, image=init_image, |
|
num_images_per_prompt=num_samples, strength=0.8, num_inference_steps=75, guidance_scale=25).images |
|
all_images.extend(images) |
|
|
|
grid = image_grid(images, num_rows, num_samples, save_individual_files=True) |
|
grid |
|
``` |
|
|
|
 |
|
|
|
|
|
We can further edit the image by introducing another feature. We start from this image |
|
``` |
|
wget https://huggingface.co/lamm-mit/SD2x-leaf-inspired/resolve/main/image_grid_2-of-4__20240722_150458.png |
|
``` |
|
|
|
 |
|
``` |
|
fname='image_grid_2-of-4__20240722_150458.png' |
|
init_image = Image.open(fname).convert("RGB") |
|
init_image = init_image.resize((768, 768)) |
|
|
|
prompt = "A nicely connected white spider web." |
|
num_samples = 4 |
|
num_rows = 1 |
|
|
|
all_images = [] |
|
for _ in range(num_rows): |
|
images = img2imgpipe(prompt, image=init_image, |
|
num_images_per_prompt=num_samples, strength=0.8, num_inference_steps=10, guidance_scale=20).images |
|
all_images.extend(images) |
|
|
|
grid = image_grid(images, num_rows, num_samples, save_individual_files=True) |
|
grid |
|
``` |
|
 |
|
|
|
A detailed view of one of them: |
|
 |
|
|
|
## Text-to-3D |
|
|
|
Download this notebook: [Convert-text-to-3D.ipynb](https://huggingface.co/lamm-mit/SD2x-leaf-inspired/resolve/main/Convert-text-to-3D.ipynb) |
|
|
|
This notebook includes the code to convert text to 3D, using an algorithm that involves a fine-tuned Stable Diffusion model and InstantMesh. |
|
|
|
``` |
|
repo_id_load= 'lamm-mit/SD2x-leaf-inspired' |
|
input_image=text_to_image_SD2x (base_model=repo_id_load, n_steps=75, guidance_scale=15, |
|
prompt = "Small chair that resembles a <leaf microstructure>.", negative_prompt="" ) |
|
display (input_image) |
|
processed_image, mv_images, mv_show_images= generate_multiviews (input_image, seed=None) |
|
display (mv_show_images) |
|
output_video, output_model_obj = make_multi_views_into_3D (mv_images, target_dir='output') |
|
Video(output_video, embed=True) |
|
``` |
|
|
|
 |
|
|
|
 |
|
|
|
### 3D printed samples |
|
|
|
In this example, we generated a 3D model and created a physical sample using additive manufacturing. |
|
|
|
``` |
|
repo_id_load= 'lamm-mit/SD2x-leaf-inspired' |
|
input_image=text_to_image_SD2x (base_model=repo_id_load, n_steps=50, guidance_scale=15, |
|
prompt = "a conch shell on black background that resembles a <leaf microstructure>, high quality", negative_prompt="" ) |
|
display (input_image) |
|
processed_image, mv_images, mv_show_images= generate_multiviews (input_image, seed=None) |
|
display (mv_show_images) |
|
output_video, output_model_obj = make_multi_views_into_3D (mv_images, target_dir='output') |
|
Video(output_video, embed=True) |
|
``` |
|
|
|
 |
|
|
|
|
|
 |
|
|
|
3D printing: Slicing using Cura and resulting physical sample with gyroid infill: |
|
|
|
 |
|
|
|
## Fine-tuning script |
|
|
|
Download this script: [SD2x DreamBooth-Fine-Tune.ipynb](https://huggingface.co/lamm-mit/SD2x-leaf-inspired/resolve/main/SD2x_DreamBooth_Fine-Tune.ipynb) |
|
|
|
You need to create a local folder ```leaf_concept_dir``` and add the leaf images (provided in this repository, see subfolder), like so: |
|
|
|
```python |
|
save_path='leaf_concept_dir' |
|
urls = [ |
|
"https://www.dropbox.com/scl/fi/4s09djm4nqxmq6vhvv9si/13_.jpg?rlkey=3m2f90pjofljmlqg5uc722i6y&dl=1", |
|
"https://www.dropbox.com/scl/fi/w4jsrf0qmrcro37nxutbx/25_.jpg?rlkey=e52gnoqaar33kwrd01h1mwcnk&dl=1", |
|
"https://www.dropbox.com/scl/fi/x0xgavduor4cbxz0sdcd2/33_.jpg?rlkey=5htaicapahhn66wnsr23v1nxz&dl=1", |
|
"https://www.dropbox.com/scl/fi/2grt40acypah9h9ok607q/72_.jpg?rlkey=bl6vfv0rcas2ygsz6o3behlst&dl=1", |
|
"https://www.dropbox.com/scl/fi/ecaf9agzdj2cawspmyt5i/117_.jpg?rlkey=oqxyk9i1wtu1wtkqadd6ylyjj&dl=1", |
|
"https://www.dropbox.com/scl/fi/gw3p73r99fleozr6ckfa3/126_.jpg?rlkey=6n7kqaklczshht1ntyqunh2lt&dl=1", |
|
## You can add additional images here |
|
] |
|
images = list(filter(None,[download_image(url) for url in urls])) |
|
|
|
if not os.path.exists(save_path): |
|
os.mkdir(save_path) |
|
|
|
[image.save(f"{save_path}/{i}.jpeg") for i, image in enumerate(images)] |
|
image_grid(images, 1, len(images)) |
|
``` |
|
|
|
The training script is included in the Jupyter notebook. |
|
|
|
## More examples |
|
|
|
```python |
|
prompt = "a conch shell on black background that resembles a <leaf microstructure>, high quality" |
|
num_samples = 4 |
|
num_rows = 4 |
|
all_images = [] |
|
for _ in range(num_rows): |
|
images = pipe(prompt, num_images_per_prompt=num_samples, num_inference_steps=50, guidance_scale=15).images |
|
all_images.extend(images) |
|
grid = image_grid(all_images, num_rows, num_samples) |
|
grid |
|
``` |
|
|
|
 |
|
|
|
|
|
 |
|
|
|
|
|
 |
|
|
|
|
|
 |
|
|
|
|
|
 |
|
|
|
|
|
 |
|
|
