FluxPipeline produces noise when .enable_vae_slicing is used, and FluxImage2ImagePipeline does not support .enable_vae_slicing.

[Meatfucker]

Describe the bug

When using the flux pipeline, if vae slicing is enabled, it produces noise instead of images, and in the image2image pipeline it is not usable at all.

Reproduction

from diffusers import FlowMatchEulerDiscreteScheduler, AutoencoderKL, FluxTransformer2DModel, FluxPipeline, utils, FluxImg2ImgPipeline
from transformers import CLIPTextModel, CLIPTokenizer,T5EncoderModel, T5TokenizerFast
from optimum.quanto import freeze, qfloat8, quantize
import torch
import gc


async def generate_flux(prompt,
                        width,
                        height,
                        steps,
                        batch_size,
                        image=None,
                        strength=None,
                        model_name=None,
                        lora_name=None):
    width = width if width is not None else 1024
    height = height if height is not None else 1024
    steps = steps if steps is not None else 30
    batch_size = batch_size if batch_size is not None else 4
    strength = strength if strength is not None else 0.7
    if model_name is None:
        model_name = "black-forest-labs/FLUX.1-dev"
        revision = "refs/pr/3"
    dtype = torch.bfloat16
    scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained(model_name, subfolder="scheduler", revision=revision)
    text_encoder = CLIPTextModel.from_pretrained("openai/clip-vit-large-patch14", torch_dtype=dtype)
    tokenizer = CLIPTokenizer.from_pretrained("openai/clip-vit-large-patch14", torch_dtype=dtype)
    text_encoder_2 = T5EncoderModel.from_pretrained(model_name, subfolder="text_encoder_2", torch_dtype=dtype,
                                                    revision=revision)
    tokenizer_2 = T5TokenizerFast.from_pretrained(model_name, subfolder="tokenizer_2", torch_dtype=dtype,
                                                  revision=revision)
    vae = AutoencoderKL.from_pretrained(model_name, subfolder="vae", torch_dtype=dtype, revision=revision)
    transformer = FluxTransformer2DModel.from_pretrained(model_name, subfolder="transformer", torch_dtype=dtype, revision=revision)

    if image is not None:
        generator = FluxImg2ImgPipeline(scheduler=scheduler,
                                        text_encoder=text_encoder,
                                        tokenizer=tokenizer,
                                        text_encoder_2=text_encoder_2,
                                        tokenizer_2=tokenizer_2,
                                        vae=vae,
                                        transformer=transformer)
        generator.enable_vae_slicing()
    else:
        generator = FluxPipeline(scheduler=scheduler,
                                 text_encoder=text_encoder,
                                 tokenizer=tokenizer,
                                 text_encoder_2=text_encoder_2,
                                 tokenizer_2=tokenizer_2,
                                 vae=vae,
                                 transformer=transformer)
        generator.enable_vae_slicing()
    if lora_name is not None:
        try:
            generator.load_lora_weights(f"loras/flux/{lora_name}", weight_name=lora_name)
        except Exception as e:
            print(f"FLUX LORA ERROR: {e}")
    quantize(transformer, weights=qfloat8)
    freeze(transformer)
    quantize(text_encoder_2, weights=qfloat8)
    freeze(text_encoder_2)
    generator.to("cuda")
    generator.set_progress_bar_config(disable=True)
    if image is not None:
        images = generator(prompt=prompt,
                           image=image,
                           width=width, height=height,
                           num_inference_steps=steps,
                           strength=strength,
                           num_images_per_prompt=batch_size).images
    else:
        images = generator(prompt=prompt,
                           width=width, height=height,
                           num_inference_steps=steps,
                           num_images_per_prompt=batch_size).images
    generator.to("cpu")

    del generator, scheduler, text_encoder, text_encoder_2, tokenizer, tokenizer_2, vae, transformer
    torch.cuda.empty_cache()
    gc.collect()
    return images

Logs

The fluxpipeline one produces no errors, the fluximage2imagepipeline says it cannot find the .enable_vae_slicing

System Info

accelerate==1.4.0
aiofiles==24.1.0
aiohappyeyeballs==2.6.1
aiohttp==3.11.18
aiosignal==1.3.2
annotated-types==0.7.0
anyio==4.8.0
attrs==25.3.0
backoff==2.2.1
beautifulsoup4==4.13.4
bitsandbytes==0.45.3
certifi==2025.1.31
cffi==1.17.1
chardet==5.2.0
charset-normalizer==3.4.1
click==8.1.8
contourpy==1.3.2
cryptography==44.0.2
cycler==0.12.1
dataclasses-json==0.6.7
diffusers==0.33.1
emoji==2.14.1
eval_type_backport==0.2.2
faiss-cpu==1.10.0
fastapi==0.115.11
filelock==3.17.0
filetype==1.2.0
fonttools==4.58.0
frozenlist==1.6.0
fsspec==2025.2.0
ftfy==6.3.1
greenlet==3.2.1
h11==0.14.0
html5lib==1.1
httpcore==1.0.7
httpx==0.28.1
httpx-sse==0.4.0
huggingface-hub==0.29.2
idna==3.10
imageio==2.37.0
imageio-ffmpeg==0.6.0
importlib_metadata==8.6.1
Jinja2==3.1.6
joblib==1.4.2
jsonpatch==1.33
jsonpointer==3.0.0
kiwisolver==1.4.8
langchain-core==0.3.56
langchain-text-splitters==0.3.8
langdetect==1.0.9
langsmith==0.3.34
loguru==0.7.3
lxml==5.4.0
MarkupSafe==3.0.2
marshmallow==3.26.1
matplotlib==3.10.3
mpmath==1.3.0
multidict==6.4.3
mypy_extensions==1.1.0
nest-asyncio==1.6.0
networkx==3.4.2
ninja==1.11.1.3
nltk==3.9.1
numpy==2.2.3
nvidia-cublas-cu12==12.4.5.8
nvidia-cuda-cupti-cu12==12.4.127
nvidia-cuda-nvrtc-cu12==12.4.127
nvidia-cuda-runtime-cu12==12.4.127
nvidia-cudnn-cu12==9.1.0.70
nvidia-cufft-cu12==11.2.1.3
nvidia-curand-cu12==10.3.5.147
nvidia-cusolver-cu12==11.6.1.9
nvidia-cusparse-cu12==12.3.1.170
nvidia-cusparselt-cu12==0.6.2
nvidia-nccl-cu12==2.21.5
nvidia-nvjitlink-cu12==12.4.127
nvidia-nvtx-cu12==12.4.127
olefile==0.47
opencv-python==4.11.0.86
optimum==1.24.0
optimum-quanto==0.2.7
orjson==3.10.16
packaging==24.2
peft==0.14.0
pillow==11.1.0
propcache==0.3.1
protobuf==6.30.0
psutil==7.0.0
pycparser==2.22
pydantic==2.11.3
pydantic-settings==2.9.1
pydantic_core==2.33.1
pyparsing==3.2.3
pypdf==5.4.0
python-dateutil==2.9.0.post0
python-dotenv==1.1.0
python-iso639==2025.2.18
python-magic==0.4.27
python-oxmsg==0.0.2
PyYAML==6.0.2
RapidFuzz==3.13.0
regex==2024.11.6
requests==2.32.3
requests-toolbelt==1.0.0
safetensors==0.5.3
scikit-learn==1.6.1
scipy==1.15.2
sentence-transformers==4.1.0
sentencepiece==0.2.0
setuptools==75.8.2
six==1.17.0
sniffio==1.3.1
soupsieve==2.7
SQLAlchemy==2.0.40
starlette==0.46.0
sympy==1.13.1
tenacity==9.1.2
threadpoolctl==3.6.0
timm==1.0.15
tokenizers==0.21.0
torch==2.6.0
torchaudio==2.6.0
torchvision==0.21.0
tqdm==4.67.1
transformers==4.50.2
triton==3.2.0
typing-inspect==0.9.0
typing-inspection==0.4.0
typing_extensions==4.12.2
unstructured==0.17.2
unstructured-client==0.34.0
urllib3==2.3.0
uvicorn==0.34.0
wcwidth==0.2.13
webencodings==0.5.1
wrapt==1.17.2
yarl==1.20.0
zipp==3.21.0
zstandard==0.23.0

Who can help?

No response

[asomoza]

let's use this issue for the enable_vae_slicing generating noise, for the img2img, it should be an easy fix.

I can reproduce this issue, using it without quantization works and also it works with bnb, so this is a quanto only problem.

A more concise code to reproduce the issue with quanto:

import torch
from optimum.quanto import freeze, qfloat8, quantize
from transformers import T5EncoderModel

from diffusers import FluxPipeline, FluxTransformer2DModel


repo_id = "black-forest-labs/FLUX.1-dev"
dtype = torch.bfloat16

text_encoder_2 = T5EncoderModel.from_pretrained(repo_id, subfolder="text_encoder_2", torch_dtype=dtype)
transformer = FluxTransformer2DModel.from_pretrained(repo_id, subfolder="transformer", torch_dtype=dtype)

pipe = FluxPipeline.from_pretrained(repo_id, text_encoder_2=text_encoder_2, transformer=transformer, torch_dtype=dtype)
pipe.enable_vae_slicing()

quantize(transformer, weights=qfloat8)
freeze(transformer)
quantize(text_encoder_2, weights=qfloat8)
freeze(text_encoder_2)

pipe.enable_model_cpu_offload()

image = pipe(prompt="a capybara", num_inference_steps=28).images[0]
image.save("flux.png")

code to test bitsandbytes:

import torch
from transformers import BitsAndBytesConfig as BitsAndBytesConfig
from transformers import T5EncoderModel

from diffusers import BitsAndBytesConfig as DiffusersBitsAndBytesConfig
from diffusers import FluxPipeline, FluxTransformer2DModel


repo_id = "black-forest-labs/FLUX.1-dev"
dtype = torch.float16

quant_config = BitsAndBytesConfig(load_in_8bit=True)

text_encoder_8bit = T5EncoderModel.from_pretrained(
    "black-forest-labs/FLUX.1-dev",
    subfolder="text_encoder_2",
    quantization_config=quant_config,
    torch_dtype=torch.float16,
)

quant_config = DiffusersBitsAndBytesConfig(load_in_8bit=True)
transformer_8bit = FluxTransformer2DModel.from_pretrained(
    "black-forest-labs/FLUX.1-dev",
    subfolder="transformer",
    quantization_config=quant_config,
    torch_dtype=torch.float16,
)

pipe = FluxPipeline.from_pretrained(
    repo_id, text_encoder_2=text_encoder_8bit, transformer=transformer_8bit, torch_dtype=dtype
)

pipe.enable_vae_slicing()
pipe.enable_model_cpu_offload()

image = pipe(prompt="a capybara", num_inference_steps=28).images[0]
image.save("flux_bnb.png")

quanto	bnb

cc: @sayakpaul

[sayakpaul]

Hmm it seems like float8 issue really and that too on some specific hardware (as also found out by @DN6, cc: @SunMarc for awareness). Does it happen with int8 too?

[asomoza]

no, tested with int8 and it works, so it's really just an issue with float8, so in the meantime, we can just say that people that have this error should switch to int8 or/and the quantization backend.

[SunMarc]

cc @DerekLiu35

[DerekLiu35]

One way to get your code to work for FluxPipeline is to downgrade optimum-quanto to 0.2.4, didn't test with FluxImg2ImgPipeline.