[WIP][Community Pipeline] InstaFlow! One-Step Stable Diffusion with Rectified Flow

[ayushtues]

Following #6037

Model/Pipeline/Scheduler description

InstaFlow is an ultra-fast, one-step image generator that achieves image quality close to Stable Diffusion, significantly reducing the demand of computational resources. This efficiency is made possible through a recent Rectified Flow technique, which trains probability flows with straight trajectories, hence inherently requiring only a single step for fast inference.

InstaFlow has several advantages:

Ultra-Fast Inference: InstaFlow models are one-step generators, which directly map noises to images and avoid multi-step sampling of diffusion models. On our machine with A100 GPU, the inference time is around 0.1 second, saving ~90% of the inference time compared to the original Stable Diffusion.
High-Quality: InstaFlow generates images with intricate details like Stable Diffusion, and have similar FID on MS COCO 2014 as state-of-the-art text-to-image GANs, like StyleGAN-T.
Simple and Efficient Training: The training process of InstaFlow merely involves supervised training. Leveraging pre-trained Stable Diffusion, it only takes 199 A100 GPU days to get InstaFlow-0.9B.

Open source status

• The model implementation is available.
• The model weights are available (Only relevant if addition is not a scheduler).

Provide useful links for the implementation

Code & Weights - https://github.com/gnobitab/InstaFlow
HF Space - https://huggingface.co/spaces/XCLiu/InstaFlow

TO-DO

• Pipeline code
• Documentation
• Tests

CC

@sayakpaul @gnobitab

[HuggingFaceDocBuilderDev]

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint.

[sayakpaul]

Let's fix the styling problems.

[sayakpaul]

Also, have you tried the pipeline yourself? Have you experienced any significant speedup? Could you maybe report some results (possibly with the generated images)?

[ayushtues]

Will try adding the LoRa results they show, since they have a better quality

[sayakpaul]

Gentle ping @ayushtues.

[ayushtues]

Will share an update soon, was away on a trip for a while

[ayushtues]

While I am not very familiar with benchmarking timings, at 32 bit precision, on a V100, I was able to generate images at around 0.03 secs (27its/sec in the tqdm log). The OG paper uses A100, which should be faster, not sure why I am getting faster inference than the 0.1 figure they had

from examples.community.instaflow_one_step import InstaFlowPipeline
import torch


pipe = InstaFlowPipeline.from_pretrained("XCLIU/instaflow_0_9B_from_sd_1_5", torch_dtype=torch.float32, cache_dir="./cache", requires_safety_checker=False)
pipe.do_lora()
pipe.to("cuda")  ### if GPU is not available, comment this line
prompt = "A hyper-realistic photo of a cute cat."

image_list = []
for i in range(10):
    images = pipe(prompt=prompt,
                num_inference_steps=1,
                guidance_scale=0.0).images
    image_list.append(images[0])
for i in range(10):
    image_list[i].save(f"./image{i}.png")

[ayushtues]

Here are some samples

[sayakpaul]

Nice very cool. Feel free to reformat and then we can merge!

[ayushtues]

Fixed formatting, and added some documentation!

[sayakpaul]

Can't we do load_lora_weights() here? Let's try use that instead.

[sayakpaul]

One comment.

[ayushtues]

Hi @sayakpaul the current code (

diffusers/examples/community/instaflow_one_step.py

Line 433 in 8793472

def do_lora(self, lora_path="Lykon/dreamshaper-7", base_sd="runwayml/stable-diffusion-v1-5", alpha=1.0):

) does the below, using DreamShaper for LoRA (https://huggingface.co/Lykon/dreamshaper-7/tree/main)

• unet = instaflow unet + (dreamshaper unet - SD unet)
• VAE = dreamshaper VAE
• Encoder = dreamshaper Encoder

Is there a way to do this using load_lora_weights? Or would we need to create a LoRA compatible checkpoint manually for this model

[sayakpaul]

If the underlying LoRA format is new, we cannot assume load_lora_weights() will work right off the bat.

[patrickvonplaten]

Good to merge for me whenever!

[ayushtues]

@sayakpaul since the text encoder and VAE are just copies of the DreamShaper model and the Unet is following a different strategy, I am not sure how to do this using just load_lora_weights

[sayakpaul]

The UNet doesn't have different pre-trained weights right?

So, if you subclass the UNet with UNet2DConditionLoadersMixin (example) and the pipeline with LoraLoaderMixin (example; for SDXL, it should SDXLLoraLoaderMixin), that should cut it until and unless I am missing out on something utterly obvious.

[ayushtues]

Small update here, I digged a bit into it, and the VAE & Encoder are the same b/w DreamShaper and Instaflow, probably just the standard SD 1.5 ones, so just need to figure out LoRA for the Unet.

Will create the LoRA weights for the Unet, and upload them to a seperate location and load from there using load_lora_weights

[ayushtues]

@sayakpaul is there a breaking change to save_attn_procs recently, was trying it out for saving the Unet and wasn't working.

Doesn't seem to be working for a simple example also,

import torch
from diffusers import DiffusionPipeline

pipeline = DiffusionPipeline.from_pretrained(
    "CompVis/stable-diffusion-v1-4",
    torch_dtype=torch.float32,
)
pipeline.unet.save_attn_procs("./", weight_name="pytorch_custom_diffusion_weights.bin")

Gives the error

TypeError                                 Traceback (most recent call last)
[<ipython-input-2-532e5d8c7935>](https://localhost:8080/#) in <cell line: 8>()
      6     torch_dtype=torch.float32,
      7 )
----> 8 pipeline.unet.save_attn_procs("./", weight_name="pytorch_custom_diffusion_weights.bin")

5 frames
[/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py](https://localhost:8080/#) in add_module(self, name, module)
    608         """
    609         if not isinstance(module, Module) and module is not None:
--> 610             raise TypeError(f"{torch.typename(module)} is not a Module subclass")
    611         elif not isinstance(name, str):
    612             raise TypeError(f"module name should be a string. Got {torch.typename(name)}")

TypeError: diffusers.models.attention_processor.AttnProcessor2_0 is not a Module subclass

Tried both the stable version of diffusers and building from source, colab here - https://colab.research.google.com/drive/1v9i09Mi9L93Cts0zywEj10t1rJVCuJmj?usp=sharing

[sayakpaul]

save_attn_procs() is deprecated.

[ayushtues]

Is there any new functionality for similar usecases, namely saving weights which we can load using load_attn_procs?

[sayakpaul]

You can use load_lora_weights(). The weights should be serialized via LoraLoaderMixin.save_lora_weights(). Here's a concrete example:

diffusers/examples/dreambooth/train_dreambooth_lora_sdxl.py

Line 1673 in dadd55f

StableDiffusionXLPipeline.save_lora_weights(

[ayushtues]

While I am yet to figure out how to use the Dreamshaper model using LORA (it doesn't seem to be using a trivial lora format), the pipeline does seem to be able to use other LORAs using simply load_lora_weights. An example with https://huggingface.co/artificialguybr/logo-redmond-1-5v-logo-lora-for-liberteredmond-sd-1-5, a LORA for generating logos

import torch
from examples.community.instaflow_one_step import InstaFlowPipeline


pipe = InstaFlowPipeline.from_pretrained("XCLIU/instaflow_0_9B_from_sd_1_5", torch_dtype=torch.float32, cache_dir="E:/InstaFlow/code/cache", requires_safety_checker=False) 
pipe.load_lora_weights("artificialguybr/logo-redmond-1-5v-logo-lora-for-liberteredmond-sd-1-5")
prompt = "logo, A logo for a fitness app, dynamic running figure, energetic colors (red, orange) ),LogoRedAF ,"
generator = torch.Generator(device="cpu").manual_seed(0)
images = pipe(prompt=prompt, 
              num_inference_steps=1,
              generator=generator, 
              guidance_scale=0.0).images 
images[0].save("./image1.png")

Without LORA

WITH LORA

[sayakpaul]

Nice. Let's include this in the README. And then ship!

[ayushtues]

Made the changes, one small thing, I haven't yet tested directly loading the custom pipeline using DiffusionPipeline, my current understanding is that it requires the PR to be merged, that is the case right?

[sayakpaul]

Made the changes, one small thing, I haven't yet tested directly loading the custom pipeline using DiffusionPipeline, my current understanding is that it requires the PR to be merged, that is the case right?

Hoe have you tested then? Ccing @patrickvonplaten for better clarification here.

[patrickvonplaten]

Good to merge for me!

[ayushtues]

Hi @sayakpaul, I tested out the pipeline locally using DiffusionPipeline as well, colab here -https://colab.research.google.com/drive/1v9i09Mi9L93Cts0zywEj10t1rJVCuJmj?usp=sharing#scrollTo=2rMIAjT7uGVD

To test it locally, we need to create a new folder, with a pipeline.py file containing the code of the pipeline, in this case I made a instaflow_pipe folder containing a pipeline.py and it worked!

import torch
from diffusers import DiffusionPipeline

pipe = DiffusionPipeline.from_pretrained("XCLIU/instaflow_0_9B_from_sd_1_5", custom_pipeline="./instaflow_pipe")

After the PR is merged, from_pretrained of DiffusionPipeline picks up the code from the community example folder in the main branch of diffusers from github by mentioning the name of the pipeline file (instaflow_one_step in this case),
as can be seen here in the comments - https://github.com/huggingface/diffusers/blob/d3eb3b35bec35f3f4e56b7528577714bf489d464/src/diffusers/pipeline_utils.py#L248C1-L277C31

so then we can do

import torch
from diffusers import DiffusionPipeline

pipe = DiffusionPipeline.from_pretrained("XCLIU/instaflow_0_9B_from_sd_1_5", custom_pipeline="instaflow_one_step")

[ayushtues]

Let me know if anything else is needed, else we can merge the PR!

[sayakpaul]

Once the CI is green will merge. Thanks for contributions!