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

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
• Table of Contents
• Table of Contents
• Model Details
  • Model Description
• Uses
  • Direct Use
  • Downstream Use [Optional]
  • Out-of-Scope Use
• Bias, Risks, and Limitations
  • Recommendations
• Training Details
  • Training Data
  • Training Procedure
    • Preprocessing
    • Speeds, Sizes, Times
• Evaluation
  • Testing Data, Factors & Metrics
    • Testing Data
    • Factors
    • Metrics
  • Results
• Model Examination
• Environmental Impact
• Technical Specifications [optional]
  • Model Architecture and Objective
  • Compute Infrastructure
    • Hardware
    • Software
• Citation
• Glossary [optional]
• More Information [optional]
• Model Card Authors [optional]
• Model Card Contact
• 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:

• 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

• most common aspect ratios in the dataset (before preprocessing)

• predefined aspect ratios

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

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.

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")