Optimizing FLUX.1-dev: Achieving Triple Performance Gains in Image Generation

Introduction: The Shift to Self-Hosted AI

The generative AI landscape is evolving as developers increasingly choose to run their own models instead of relying solely on third-party APIs. With high-quality open-source models and efficient frameworks, users can enjoy customized generative capabilities across various modalities including images, audio, and text. However, self-hosting introduces challenges regarding performance, latency, and the need for competitive quality against proprietary solutions.

In this post, we focus on enhancing the performance of the FLUX.1-dev model from Black Forest Labs. By implementing strategic optimizations, we’ve succeeded in making it operate at speeds comparable to proprietary services, achieving a threefold increase in inference speed along with reductions in cold boot latency. For the full code implementation, refer to our GitHub repository.

Key Optimization Findings

To be competitive with existing APIs offering FLUX.1-dev image generation, we identified a target response time of under three seconds. We realized a performance boost through two main avenues:

• Utilizing standard optimizations that leverage compiler efficiency and hardware capability delivered a 1.5x speed increase.
• Implementing an approximate caching technique called First Block Caching led to an additional 2x reduction in latency.

Establishing the Performance Baseline

Before implementing optimizations, it is critical to measure the existing performance accurately. We employed the standard Hugging Face diffusers library to implement a FluxPipeline, initialized with 16-bit precision settings:

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

Our initial assessments demonstrated a mean time of approximately 6.75 seconds to generate a 1024x1024 image. This baseline latency serves as a reference point for our subsequent performance improvements.

Achieving a 1.5x Speed Enhancement with Standard Optimizations

We began with conventional optimizations, including:

• Using PyTorch’s optimizing compiler
• Merging the computations for the query, key, and value in the attention mechanism of the Transformer
• Switching to a ‘channels last’ memory layout

Optimizing the Compute Graph with the Torch Compiler

Inherent in PyTorch is a dynamic computation graph that is executed on the CPU before triggering the actual calculations on the GPU. Transforming this virtual graph into a more efficient one is crucial for performance. We took advantage of the Torch compiler, which compiles dynamic graphs into lower-level representations, optimizing them for faster execution. A configuration snippet for our setup is shown below

self.pipe.transformer = torch.compile(
    self.pipe.transformer, mode="max-autotune-no-cudagraphs", dynamic=True
)

Fusing Query, Key, and Value Computations

In the Transformer architecture, the attention mechanism typically involves three separate matrix multiplications that can slow processing. By concatenating the weight matrices for query, key, and value, we can compute them in a single operation, unlocking greater parallelism:

self.pipe.transformer.fuse_qkv_projections()

Improving Data Locality with Channels Last Memory Layout

We further optimized memory access patterns by employing a channels-last memory layout (instead of the default channels-first). Doing so enhances data locality, enabling faster access during convolutions, typical in neural networks:

self.pipe.transformer.to(memory_format=torch.channels_last)

Performance Gains from Standard Optimizations

The cumulative effect of these strategies yielded a performance improvement of approximately 1.5x. The following video illustrates how the image evolves across denoising steps for both the baseline and optimized methods:

Achieving a Further 2x Speed Improvement via Approximate Caching

While the standard optimizations were essential, we pushed for additional improvements through approximate caching techniques. Generative models, specifically diffusion models, operate through iterative steps that can sometimes demonstrate minimal changes.

By employing First Block Caching, we enable the model to skip certain steps based on whether significant changes are expected. This technique enhances performance while maintaining acceptable output quality:

apply_cache_on_pipe(
    self.pipe,
    residual_diff_threshold=0.12,
)

This strategic threshold allows us to maximize performance gains while ensuring the quality of generated images remains at a satisfactory level.

Reducing Cold Start Latency by 30x

While optimizing inference, we observed that our boot times significantly increased, caused predominantly by the Torch compiler’s profiling activity. To address this, we utilized caching to store and retrieve artifacts, which helped mitigate the large initialization times.

Employing Memory Snapshots in Modal allows us to reduce the time taken to load files and execute initializations dramatically:

@modal.enter(snap=True)
def load(self):
    self.pipe = FluxPipeline.from_pretrained(
        "black-forest-labs/FLUX.1-dev",
        torch_dtype=torch.bfloat16,
        use_safetensors=True,
    ).to("cpu")

Serving AI Models Efficiently with Modal

Through these optimizations, we've successfully reduced the serving latency of FLUX.1-dev, achieving performance levels that can compete with proprietary providers. The methodologies explored here enable broader use cases in self-hosted generative APIs, allowing organizations to harness the power of AI efficiently while controlling costs.

To discover more about the potential of our serverless platform for generative inference, check out the Modal website and see how we are facilitating robust solutions for a variety of clientele, from media companies to tech startups.