---
license: openrail++
tags:
- stable-diffusion
- text-to-image
pinned: true
---
# Model Card for flex-diffusion-2-1
<!-- Provide a quick summary of what the model is/does. [Optional] -->
stable-diffusion-2-1 (stabilityai/stable-diffusion-2-1) finetuned with different aspect ratios.
## TLDR:
### There are 2 models in this repo:
- One based on stable-diffusion-2-1 (stabilityai/stable-diffusion-2-1) finetuned for 6k steps.
- One based on stable-diffusion-2-base (stabilityai/stable-diffusion-2-base) finetuned for 6k steps, on the same dataset.
For usage, see - [How to Get Started with the Model](#how-to-get-started-with-the-model)
### It aims to solve the following issues:
1. Generated images looks like they are cropped from a larger image.
2. Generating non-square images creates weird results, due to the model being trained on square images.
Examples:
| resolution | model | stable diffusion | flex diffusion |
|:---------------:|:-------:|:----------------------------:|:-----------------------------:|
| 576x1024 (9:16) | v2-1 |  |  |
| 576x1024 (9:16) | v2-base |  |  |
| 1024x576 (16:9) | v2-1 |  |  |
| 1024x576 (16:9) | v2-base |  |  |
### Limitations:
1. It's trained on a small dataset, so it's improvements may be limited.
2. For each aspect ratio, it's trained on only a fixed resolution. So it may not be able to generate images of different resolutions.
For 1:1 aspect ratio, it's fine-tuned at 512x512, although flex-diffusion-2-1 was last finetuned at 768x768.
### Potential improvements:
1. Train on a larger dataset.
2. Train on different resolutions even for the same aspect ratio.
3. Train on specific aspect ratios, instead of a range of aspect ratios.
# Table of Contents
- [Model Card for flex-diffusion-2-1](#model-card-for--model_id-)
- [Table of Contents](#table-of-contents)
- [Table of Contents](#table-of-contents-1)
- [Model Details](#model-details)
- [Model Description](#model-description)
- [Uses](#uses)
- [Direct Use](#direct-use)
- [Downstream Use [Optional]](#downstream-use-optional)
- [Out-of-Scope Use](#out-of-scope-use)
- [Bias, Risks, and Limitations](#bias-risks-and-limitations)
- [Recommendations](#recommendations)
- [Training Details](#training-details)
- [Training Data](#training-data)
- [Training Procedure](#training-procedure)
- [Preprocessing](#preprocessing)
- [Speeds, Sizes, Times](#speeds-sizes-times)
- [Evaluation](#evaluation)
- [Testing Data, Factors & Metrics](#testing-data-factors--metrics)
- [Testing Data](#testing-data)
- [Factors](#factors)
- [Metrics](#metrics)
- [Results](#results)
- [Model Examination](#model-examination)
- [Environmental Impact](#environmental-impact)
- [Technical Specifications [optional]](#technical-specifications-optional)
- [Model Architecture and Objective](#model-architecture-and-objective)
- [Compute Infrastructure](#compute-infrastructure)
- [Hardware](#hardware)
- [Software](#software)
- [Citation](#citation)
- [Glossary [optional]](#glossary-optional)
- [More Information [optional]](#more-information-optional)
- [Model Card Authors [optional]](#model-card-authors-optional)
- [Model Card Contact](#model-card-contact)
- [How to Get Started with the Model](#how-to-get-started-with-the-model)
# Model Details
## Model Description
<!-- Provide a longer summary of what this model is/does. -->
stable-diffusion-2-1 (stabilityai/stable-diffusion-2-1) finetuned for dynamic aspect ratios.
finetuned resolutions:
| | width | height | aspect ratio |
|---:|--------:|---------:|:---------------|
| 0 | 512 | 1024 | 1:2 |
| 1 | 576 | 1024 | 9:16 |
| 2 | 576 | 960 | 3:5 |
| 3 | 640 | 1024 | 5:8 |
| 4 | 512 | 768 | 2:3 |
| 5 | 640 | 896 | 5:7 |
| 6 | 576 | 768 | 3:4 |
| 7 | 512 | 640 | 4:5 |
| 8 | 640 | 768 | 5:6 |
| 9 | 640 | 704 | 10:11 |
| 10 | 512 | 512 | 1:1 |
| 11 | 704 | 640 | 11:10 |
| 12 | 768 | 640 | 6:5 |
| 13 | 640 | 512 | 5:4 |
| 14 | 768 | 576 | 4:3 |
| 15 | 896 | 640 | 7:5 |
| 16 | 768 | 512 | 3:2 |
| 17 | 1024 | 640 | 8:5 |
| 18 | 960 | 576 | 5:3 |
| 19 | 1024 | 576 | 16:9 |
| 20 | 1024 | 512 | 2:1 |
- **Developed by:** Jonathan Chang
- **Model type:** Diffusion-based text-to-image generation model
- **Language(s)**: English
- **License:** creativeml-openrail-m
- **Parent Model:** https://huggingface.co/stabilityai/stable-diffusion-2-1
- **Resources for more information:** More information needed
# Uses
- see https://huggingface.co/stabilityai/stable-diffusion-2-1
# Training Details
## Training Data
- LAION aesthetic dataset, subset of it with 6+ rating
- https://laion.ai/blog/laion-aesthetics/
- https://huggingface.co/datasets/ChristophSchuhmann/improved_aesthetics_6plus
- I only used a small portion of that, see [Preprocessing](#preprocessing)
- most common aspect ratios in the dataset (before preprocessing)
| | aspect_ratio | counts |
|---:|:---------------|---------:|
| 0 | 1:1 | 154727 |
| 1 | 3:2 | 119615 |
| 2 | 2:3 | 61197 |
| 3 | 4:3 | 52276 |
| 4 | 16:9 | 38862 |
| 5 | 400:267 | 21893 |
| 6 | 3:4 | 16893 |
| 7 | 8:5 | 16258 |
| 8 | 4:5 | 15684 |
| 9 | 6:5 | 12228 |
| 10 | 1000:667 | 12097 |
| 11 | 2:1 | 11006 |
| 12 | 800:533 | 10259 |
| 13 | 5:4 | 9753 |
| 14 | 500:333 | 9700 |
| 15 | 250:167 | 9114 |
| 16 | 5:3 | 8460 |
| 17 | 200:133 | 7832 |
| 18 | 1024:683 | 7176 |
| 19 | 11:10 | 6470 |
- predefined aspect ratios
| | width | height | aspect ratio |
|---:|--------:|---------:|:---------------|
| 0 | 512 | 1024 | 1:2 |
| 1 | 576 | 1024 | 9:16 |
| 2 | 576 | 960 | 3:5 |
| 3 | 640 | 1024 | 5:8 |
| 4 | 512 | 768 | 2:3 |
| 5 | 640 | 896 | 5:7 |
| 6 | 576 | 768 | 3:4 |
| 7 | 512 | 640 | 4:5 |
| 8 | 640 | 768 | 5:6 |
| 9 | 640 | 704 | 10:11 |
| 10 | 512 | 512 | 1:1 |
| 11 | 704 | 640 | 11:10 |
| 12 | 768 | 640 | 6:5 |
| 13 | 640 | 512 | 5:4 |
| 14 | 768 | 576 | 4:3 |
| 15 | 896 | 640 | 7:5 |
| 16 | 768 | 512 | 3:2 |
| 17 | 1024 | 640 | 8:5 |
| 18 | 960 | 576 | 5:3 |
| 19 | 1024 | 576 | 16:9 |
| 20 | 1024 | 512 | 2:1 |
## Training Procedure
<!-- This relates heavily to the Technical Specifications. Content here should link to that section when it is relevant to the training procedure. -->
### Preprocessing
1. download files with url & caption from https://huggingface.co/datasets/ChristophSchuhmann/improved_aesthetics_6plus
- I only used the first file `train-00000-of-00007-29aec9150af50f9f.parquet`
2. use img2dataset to convert to webdataset
- https://github.com/rom1504/img2dataset
- I put train-00000-of-00007-29aec9150af50f9f.parquet in a folder called `first-file`
- the output folder is `/mnt/aesthetics6plus`, change this to your own folder
```bash
echo INPUT_FOLDER=first-file
echo OUTPUT_FOLDER=/mnt/aesthetics6plus
img2dataset --url_list $INPUT_FOLDER --input_format "parquet"\
--url_col "URL" --caption_col "TEXT" --output_format webdataset\
--output_folder $OUTPUT_FOLDER --processes_count 3 --thread_count 6 --image_size 1024 --resize_only_if_bigger --resize_mode=keep_ratio_largest \
--save_additional_columns '["WIDTH","HEIGHT","punsafe","similarity"]' --enable_wandb True
```
3. The data-loading code will do preprocessing on the fly, so no need to do anything else. But it's not optimized for speed, the GPU utilization fluctuates between 80% and 100%. And it's not written for multi-GPU training, so use it with caution. The code will do the following:
- use webdataset to load the data
- calculate the aspect ratio of each image
- find the closest aspect ratio & it's associated resolution from the predefined resolutions: `argmin(abs(aspect_ratio - predefined_aspect_ratios))`. E.g. if the aspect ratio is 1:3, the closest resolution is 1:2. and it's associated resolution is 512x1024.
- keeping the aspect ratio, resize the image such that it's larger or equal to the associated resolution on each side. E.g. resize to 512x(512*3) = 512x1536
- random crop the image to the associated resolution. E.g. crop to 512x1024
- if more than 10% of the image is lost in the cropping, discard this example.
- batch examples by aspect ratio, so all examples in a batch have the same aspect ratio
### Speeds, Sizes, Times
- Dataset size: 100k image-caption pairs, before filtering.
- I didn't wait for the whole dataset to be downloaded, I copied the first 10 tar files and their index files to a new folder called `aesthetics6plus-small`, with 100k image-caption pairs in total. The full dataset is a lot bigger.
- Hardware: 1 RTX3090 GPUs
- Optimizer: 8bit Adam
- Batch size: 32
- actual batch size: 2
- gradient_accumulation_steps: 16
- effective batch size: 32
- Learning rate: warmup to 2e-6 for 500 steps and then kept constant
- Learning rate: 2e-6
- Training steps: 6k
- Epoch size (approximate): 32 * 6k / 100k = 1.92 (not accounting for the filtering)
- Each example is seen 1.92 times on average.
- Training time: approximately 1 day
## Results
More information needed
# Model Card Authors
Jonathan Chang
# How to Get Started with the Model
Use the code below to get started with the model.
```python
from diffusers import StableDiffusionPipeline, DPMSolverMultistepScheduler, UNet2DConditionModel
def use_DPM_solver(pipe):
pipe.scheduler = DPMSolverMultistepScheduler.from_config(pipe.scheduler.config)
return pipe
pipe = StableDiffusionPipeline.from_pretrained(
"stabilityai/stable-diffusion-2-1",
unet = UNet2DConditionModel.from_pretrained("ttj/flex-diffusion-2-1", subfolder="2-1/unet", torch_dtype=torch.float16),
torch_dtype=torch.float16,
)
# for v2-base, use the following line instead
#pipe = StableDiffusionPipeline.from_pretrained(
# "stabilityai/stable-diffusion-2-base",
# unet = UNet2DConditionModel.from_pretrained("ttj/flex-diffusion-2-1", subfolder="2-base/unet", torch_dtype=torch.float16),
# torch_dtype=torch.float16)
pipe = use_DPM_solver(pipe).to("cuda")
pipe = pipe.to("cuda")
prompt = "a professional photograph of an astronaut riding a horse"
image = pipe(prompt).images[0]
image.save("astronaut_rides_horse.png")
```