Ronen Kamenetsky, Sara Dorfman, Daniel Garibi, Roni Paiss, Or Patashnik, Daniel Cohen-Or

Large-scale text-to-image diffusion models have become the backbone of modern image editing, yet text prompts alone do not offer adequate control over the editing process. Two properties are especially desirable: disentanglement, where changing one attribute does not unintentionally alter others, and continuous control, where the strength of an edit can be smoothly adjusted. We introduce a method for disentangled and continuous editing through token-level manipulation of text embeddings. The edits are applied by manipulating the embeddings along carefully chosen directions, which control the strength of the target attribute. To identify such directions, we employ a Sparse Autoencoder (SAE), whose sparse latent space exposes semantically isolated dimensions. Our method operates directly on text embeddings without modifying the diffusion process, making it model agnostic and broadly applicable to various image synthesis backbones. Experiments show that it enables intuitive and efficient manipulations with continuous control across diverse attributes and domains.

The official implementation of the paper SAEdit: Token-level control for continuous image editing via Sparse AutoEncoder

Our code builds on the requirement of the diffusers library. To set up the environment, please run:

conda env create -f environment.yaml
conda activate saeedit_env

or install requirements:

pip install -r requirements.txt

To run our method, insert the SAEdit block as a callback function as can be seen bellow, and in the example subfolder.

import torch
import yaml
from diffusers import FluxPipeline
from src.saedit import SAEditCallback
from src.models.matryoshka_sae import GlobalBatchTopKMatryoshkaSAE

def main():
    device = "cuda" if torch.cuda.is_available() else "cpu"
    prompt = "a portrait of a man riding a donkey in the snow"
    sentence_pair_paths = "configs/variations/smiling_man.yaml"
    tokens_to_edit = ['man']

    # Load model
    model = FluxPipeline.from_pretrained(
        "black-forest-labs/FLUX.1-dev",
        torch_dtype=torch.bfloat16
    )
    model.enable_model_cpu_offload()
    
    # Load SAE model
    sae = GlobalBatchTopKMatryoshkaSAE.from_pretrained(
                                        "Ronenk94/T5_matryoshka_sae",
                                         device="cuda")    
    # Load variation configuration
    with open(sentence_pair_paths, "r") as f:
        sentence_pairs = yaml.safe_load(f)
    
    # Create SAEdit callback
    callback = SAEditCallback(
        pipeline=model,
        sae=sae,
        prompt=prompt,
        source_tokens_to_edit=tokens_to_edit,
        factor=0.8, # set 0 to retrieve original image
        sentence_pairs=sentence_pairs["sentence_pairs"],
        max_sequence_length=256
    )
    
    # Generate image
    output = model(
        prompt=prompt,
        guidance_scale=3.5,
        height=1024,
        width=1024,
        num_inference_steps=40,
        max_sequence_length=256,
        generator=torch.Generator(device=device).manual_seed(42),
        callback_on_step_end=callback,
        callback_on_step_end_tensor_inputs=['prompt_embeds']
    )
    
    # Save result
    output.images[0].save("output.png")
    print("Image saved to output.png")

if __name__ == "__main__":
    main()

We have made the weights of our model public in hugging face.

To train a new SAE model two steps are required:

First, generate a Zarr file containing T5 embeddings from your training dataset (should be ran with torchrun):

torchrun src/train/generate_data.py --text_encoder_path <The path of the text encoder e.g T5EncoderModel> --text_tokenizer_path <The path of the corresponding tokenizer> --output_path <the path of the output> --data_output_length <number of rows in the resulting dataset>

This will create a Zarr file containing all T5 embeddings from your training dataset.

Train the SAE model using the generated embeddings:

python src/train/train.py --cfg src/train/configs/train_matryoshka_topk_300_dict_65k.yml

Configuration Setup:

• Update the dataset_path in your chosen config file (.yml) to point to the Zarr file created in Step 1
• Adjust other hyperparameters (batch size, learning rate, etc.) as needed
• The trained model will be saved as a snapshot (.pt file) that can be loaded and used with SAEdit

This code builds on the code from the diffusers library. In addition for defining and training the SAE code was borrowed from Matryoshka_sae

If you use this code for your research, please cite the following work:

@misc{kamenetsky2025saedittokenlevelcontrolcontinuous,
      title={SAEdit: Token-level control for continuous image editing via Sparse AutoEncoder}, 
      author={Ronen Kamenetsky and Sara Dorfman and Daniel Garibi and Roni Paiss and Or Patashnik and Daniel Cohen-Or},
      year={2025},
      eprint={2510.05081},
      archivePrefix={arXiv},
      primaryClass={cs.GR},
      url={https://arxiv.org/abs/2510.05081}, 
}