Opening article
A demonstration script utilizing the new Flux2Pipeline and Flux2Transformer2DModel from Hugging Face Diffusers encounters an out-of-memory error during execution. The code loads the diffusers/FLUX.2-dev-bnb-4bit repository utilizing PyTorch and Mistral3ForConditionalGeneration components. The process fails during the forward pass of the text encoder.
import torch from transformers import Mistral3ForConditionalGeneration
from diffusers import Flux2Pipeline, Flux2Transformer2DModel
repo_id = "diffusers/FLUX.2-dev-bnb-4bit" device = "cuda:0" torch_dtype = torch.bfloat16
transformer = Flux2Transformer2DModel.from_pretrained( repo_id, subfolder="transformer", torch_dtype=torch_dtype, device_map="cpu" ) text_encoder = Mistral3ForConditionalGeneration.from_pretrained( repo_id, subfolder="text_encoder", dtype=torch_dtype, device_map="cpu" )
pipe = Flux2Pipeline.from_pretrained( repo_id, transformer=transformer, text_encoder=text_encoder, torch_dtype=torch_dtype ) pipe.enable_model_cpu_offload()
Flax classes are deprecated and will be removed in Diffusers v1.0.0. We recommend migrating to PyTorch classes or pinning your version of Diffusers. Flax classes are deprecated and will be removed in Diffusers v1.0.0. We recommend migrating to PyTorch classes or pinning your version of Diffusers. /usr/local/lib/python3.12/dist-packages/huggingface_hub/utils/_validators.py:206: UserWarning: The local_dir_use_symlinks argument is deprecated and ignored in hf_hub_download . Downloading to a local directory does not use symlinks anymore. warnings.warn( Download complete: 0.00/0.00 [00:00<?, ?B/s]Fetching 2 files: 100% 2/2 [00:00<00:00, 98.80it/s]Loading checkpoint shards: 100% 2/2 [00:01<00:00, 2.03it/s]Download complete: 0.00/0.00 [00:00<?, ?B/s]Fetching 4 files: 100% 4/4 [00:00<00:00, 167.57it/s]Loading weights: 100% 585/585 [00:02<00:00, 310.91it/s, Materializing param=model.vision_tower.transformer.layers.23.ffn_norm.weight]The tied weights mapping and config for this model specifies to tie model.language_model.embed_tokens.weight to lm_head.weight, but both are present in the checkpoints, so we will NOT tie them. You should update the config with tie_word_embeddings=False to silence this warning Loading pipeline components...: 100% 5/5 [00:03<00:00, 1.14it/s]--------------------------------------------------------------------------- OutOfMemoryError Traceback (most recent call last) /tmp/ipykernel_18947/863729753.py in <cell line: 0>() 22 prompt = "Realistic macro photograph of a hermit crab using a soda can as its shell, partially emerging from the can, captured with sharp detail and natural colors, on a sunlit beach with soft shadows and a shallow depth of field, with blurred ocean waves in the background. The can has the text BFL Diffusers on it and it has a color gradient that start with #FF5733 at the top and transitions to #33FF57 at the bottom." 23 ---> 24 image = pipe( 25 prompt=prompt, 26 generator=torch.Generator(device=device).manual_seed(42),
prompt = "Realistic macro photograph of a hermit crab using a soda can as its shell, partially emerging from the can, captured with sharp detail and natural colors, on a sunlit beach with soft shadows and a shallow depth of field, with blurred ocean waves in the background. The can has the text BFL Diffusers on it and it has a color gradient that start with #FF5733 at the top and transitions to #33FF57 at the bottom."
image = pipe( prompt=prompt, generator=torch.Generator(device=device).manual_seed(42), num_inference_steps=50, # 28 is a good trade-off guidance_scale=4, ).images[0]
36 frames/usr/local/lib/python3.12/dist-packages/torch/utils/_contextlib.py in decorate_context(*args, **kwargs) 122 # pyrefly: ignore [bad-context-manager] 123 with ctx_factory(): --> 124 return func(*args, **kwargs) 125 126 return decorate_context
Hardware resource constraints such as out-of-memory errors remain a practical hurdle when running large multimodal pipelines combining transformer-based text encoders and diffusion models.