Flux-1.dev generation with custom caching schedules (Pro)

This tutorial demonstrates how to use the pruna package to optimize your Flux model for faster inference using your own caching schedules. Any execution times given below are measured on an H100 GPU.

1. Loading the Flux Model

First, load your Flux model.

import torch
from diffusers import FluxPipeline

pipe = FluxPipeline.from_pretrained("black-forest-labs/FLUX.1-dev", torch_dtype=torch.bfloat16).to("cuda")

Let’s generate a sample using the original FLUX.1-dev model.

prompt = (
    "An anime illustration of Sydney Opera House sitting next to Eiffel tower, under a blue night sky of "
    "roiling energy, exploding yellow stars, and radiating swirls of blue."
)

pipe(prompt, num_inference_steps=28).images[0]

2. Defining a Custom Schedule

First, we define a custom 28 steps periodic schedule.

my_schedule = [1, 1, 0, 0, 0, 1, 0, 0, 0, 1, 0, 0, 0, 1, 0, 0, 0, 1, 0, 0, 0, 1, 0, 0, 0, 1, 0, 1]

3. Smashing the Model

Now, you can smash the model using your custom schedule.

Don’t forget to replace the token by the one provided by PrunaAI.

from pruna_pro import SmashConfig, smash

# Initialize the SmashConfig
smash_config = SmashConfig({
    "custom_schedule": {
        "schedule_list": my_schedule
    },
    "torch_compile": {
        "target": "module_list"
    }
})

pipe = smash(
    model=pipe,
    token="<your_pruna_token>",
    smash_config=smash_config,
)

4. Running the Model

After the model has been compiled, we run inference for a few iterations as warm-up. You can remove torch_compile from the compiler argument in the smash_config if you prefer instant speed-up without warm-up iterations.

prompt = (
    "An anime illustration of Sydney Opera House sitting next to Eiffel tower, under a blue night sky of "
    "roiling energy, exploding yellow stars, and radiating swirls of blue."
)

for _ in range(5):
    pipe(prompt, num_inference_steps=28).images[0]

Run the model to generate images with accelerated inference.

pipe(prompt, num_inference_steps=28).images[0]

Wrap Up

Congratulations! You have successfully smashed the Flux-1.dev model with a custom schedule. Enjoy the speed-up!

For further control over the quality of caching, see how to change the caching mode in the FLUX smashing tutorial.