FluxSpace: Disentangled Semantic Editing in Rectified Flow Transformers

Official implementation of FluxSpace, a domain-agnostic image editing method for rectified flow transformers like FLUX. This work enables precise, attribute-specific modifications without affecting unrelated aspects of the image.

Abstract

Rectified flow models have emerged as a dominant approach in image generation, showcasing impressive capabilities in high-quality image synthesis. However, despite their effectiveness in visual generation, rectified flow models often struggle with disentangled editing of images. This limitation prevents the ability to perform precise, attribute-specific modifications without affecting unrelated aspects of the image. In this paper, we introduce FluxSpace, a domain-agnostic image editing method leveraging a representation space with the ability to control the semantics of images generated by rectified flow transformers, such as Flux. By leveraging the representations learned by the transformer blocks within the rectified flow models, we propose a set of semantically interpretable representations that enable a wide range of image editing tasks, from fine-grained image editing to artistic creation. This work offers a scalable and effective image editing approach, along with its disentanglement capabilities.

Install

# Using conda (recommended)
conda create -n fluxspace python=3.12 -y
conda activate fluxspace

# Install the requirements
pip install -r requirements.txt

Quick Start

import torch
from flux_semantic_pipeline import GenerationPipelineSemantic

# Load pipeline
pipe = GenerationPipelineSemantic.from_pretrained(
    "black-forest-labs/FLUX.1-dev", 
    torch_dtype=torch.bfloat16
).to("cuda")
pipe.register_transformer_blocks()

# Generating the original image
generator = torch.Generator().manual_seed(0)
image = pipe(
    prompt="A portrait photo of a man",
    edit_prompt="",
    guidance_scale=3.5,
    num_inference_steps=30,
    edit_start_iter=30,
    edit_global_scale=0.0,
    edit_content_scale=0,
    attention_threshold=0.0,
    generator=generator
).images[0]

image.save("original_image.png")

# Generate with semantic editing
generator = torch.Generator().manual_seed(0)
image = pipe(
    prompt="A portrait photo of a man",
    edit_prompt="eyeglasses",
    guidance_scale=3.5,
    num_inference_steps=30,
    edit_start_iter=3,
    edit_global_scale=0.8,
    edit_content_scale=5.0,
    attention_threshold=0.5,
    generator=generator
).images[0]

image.save("edited_image.png")

Examples

For comprehensive examples and interactive demos, see demo.ipynb.

Citation

If you use this work in your research, please cite our paper:

@article{dalva2024fluxspace,
  title={FluxSpace: Disentangled Semantic Editing in Rectified Flow Transformers},
  author={Dalva, Yusuf and Venkatesh, Kavana and Yanardag, Pinar},
  journal={arXiv preprint arXiv:2412.09611},
  year={2024}
}

@inproceedings{dalva2025fluxspace,
  title={FluxSpace: Disentangled Semantic Editing in Rectified Flow Models},
  author={Dalva, Yusuf and Venkatesh, Kavana and Yanardag, Pinar},
  booktitle={Proceedings of the Computer Vision and Pattern Recognition Conference},
  pages={13083--13092},
  year={2025}
}