Flux (original) (raw)

Flux is a series of text-to-image generation models based on diffusion transformers. To know more about Flux, check out the original blog post by the creators of Flux, Black Forest Labs.

Original model checkpoints for Flux can be found here. Original inference code can be found here.

Flux can be quite expensive to run on consumer hardware devices. However, you can perform a suite of optimizations to run it faster and in a more memory-friendly manner. Check out this section for more details. Additionally, Flux can benefit from quantization for memory efficiency with a trade-off in inference latency. Refer to this blog post to learn more. For an exhaustive list of resources, check out this gist.

Flux comes in two variants:

Both checkpoints have slightly difference usage which we detail below.

Timestep-distilled

import torch from diffusers import FluxPipeline

pipe = FluxPipeline.from_pretrained("black-forest-labs/FLUX.1-schnell", torch_dtype=torch.bfloat16) pipe.enable_model_cpu_offload()

prompt = "A cat holding a sign that says hello world" out = pipe( prompt=prompt, guidance_scale=0., height=768, width=1360, num_inference_steps=4, max_sequence_length=256, ).images[0] out.save("image.png")

Guidance-distilled

import torch from diffusers import FluxPipeline

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

prompt = "a tiny astronaut hatching from an egg on the moon" out = pipe( prompt=prompt, guidance_scale=3.5, height=768, width=1360, num_inference_steps=50, ).images[0] out.save("image.png")

Running FP16 inference

Flux can generate high-quality images with FP16 (i.e. to accelerate inference on Turing/Volta GPUs) but produces different outputs compared to FP32/BF16. The issue is that some activations in the text encoders have to be clipped when running in FP16, which affects the overall image. Forcing text encoders to run with FP32 inference thus removes this output difference. See here for details.

FP16 inference code:

import torch from diffusers import FluxPipeline

pipe = FluxPipeline.from_pretrained("black-forest-labs/FLUX.1-schnell", torch_dtype=torch.bfloat16)

pipe.enable_sequential_cpu_offload() pipe.vae.enable_slicing() pipe.vae.enable_tiling()

pipe.to(torch.float16)

prompt = "A cat holding a sign that says hello world" out = pipe( prompt=prompt, guidance_scale=0., height=768, width=1360, num_inference_steps=4, max_sequence_length=256, ).images[0] out.save("image.png")

Single File Loading for the FluxTransformer2DModel

The FluxTransformer2DModel supports loading checkpoints in the original format shipped by Black Forest Labs. This is also useful when trying to load finetunes or quantized versions of the models that have been published by the community.

`FP8` inference can be brittle depending on the GPU type, CUDA version, and `torch` version that you are using. It is recommended that you use the `optimum-quanto` library in order to run FP8 inference on your machine.

The following example demonstrates how to run Flux with less than 16GB of VRAM.

First install optimum-quanto

pip install optimum-quanto

Then run the following example

import torch from diffusers import FluxTransformer2DModel, FluxPipeline from transformers import T5EncoderModel, CLIPTextModel from optimum.quanto import freeze, qfloat8, quantize

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

transformer = FluxTransformer2DModel.from_single_file("https://huggingface.co/Kijai/flux-fp8/blob/main/flux1-dev-fp8.safetensors", torch_dtype=dtype) quantize(transformer, weights=qfloat8) freeze(transformer)

text_encoder_2 = T5EncoderModel.from_pretrained(bfl_repo, subfolder="text_encoder_2", torch_dtype=dtype) quantize(text_encoder_2, weights=qfloat8) freeze(text_encoder_2)

pipe = FluxPipeline.from_pretrained(bfl_repo, transformer=None, text_encoder_2=None, torch_dtype=dtype) pipe.transformer = transformer pipe.text_encoder_2 = text_encoder_2

pipe.enable_model_cpu_offload()

prompt = "A cat holding a sign that says hello world" image = pipe( prompt, guidance_scale=3.5, output_type="pil", num_inference_steps=20, generator=torch.Generator("cpu").manual_seed(0) ).images[0]

image.save("flux-fp8-dev.png")

FluxPipeline

class diffusers.FluxPipeline

< source >

( scheduler: FlowMatchEulerDiscreteScheduler vae: AutoencoderKL text_encoder: CLIPTextModel tokenizer: CLIPTokenizer text_encoder_2: T5EncoderModel tokenizer_2: T5TokenizerFast transformer: FluxTransformer2DModel )

Parameters

The Flux pipeline for text-to-image generation.

Reference: https://blackforestlabs.ai/announcing-black-forest-labs/

__call__

< source >

( prompt: Union = None prompt_2: Union = None height: Optional = None width: Optional = None num_inference_steps: int = 28 timesteps: List = None guidance_scale: float = 3.5 num_images_per_prompt: Optional = 1 generator: Union = None latents: Optional = None prompt_embeds: Optional = None pooled_prompt_embeds: Optional = None output_type: Optional = 'pil' return_dict: bool = True joint_attention_kwargs: Optional = None callback_on_step_end: Optional = None callback_on_step_end_tensor_inputs: List = ['latents'] max_sequence_length: int = 512 ) → ~pipelines.flux.FluxPipelineOutput or tuple

Parameters

Returns

~pipelines.flux.FluxPipelineOutput or tuple

~pipelines.flux.FluxPipelineOutput if return_dictis True, otherwise a tuple. When returning a tuple, the first element is a list with the generated images.

Function invoked when calling the pipeline for generation.

Examples:

import torch from diffusers import FluxPipeline

pipe = FluxPipeline.from_pretrained("black-forest-labs/FLUX.1-schnell", torch_dtype=torch.bfloat16) pipe.to("cuda") prompt = "A cat holding a sign that says hello world"

image = pipe(prompt, num_inference_steps=4, guidance_scale=0.0).images[0] image.save("flux.png")

Disable sliced VAE decoding. If enable_vae_slicing was previously enabled, this method will go back to computing decoding in one step.

Disable tiled VAE decoding. If enable_vae_tiling was previously enabled, this method will go back to computing decoding in one step.

Enable sliced VAE decoding. When this option is enabled, the VAE will split the input tensor in slices to compute decoding in several steps. This is useful to save some memory and allow larger batch sizes.

Enable tiled VAE decoding. When this option is enabled, the VAE will split the input tensor into tiles to compute decoding and encoding in several steps. This is useful for saving a large amount of memory and to allow processing larger images.

encode_prompt

< source >

( prompt: Union prompt_2: Union device: Optional = None num_images_per_prompt: int = 1 prompt_embeds: Optional = None pooled_prompt_embeds: Optional = None max_sequence_length: int = 512 lora_scale: Optional = None )

Parameters

FluxImg2ImgPipeline

class diffusers.FluxImg2ImgPipeline

< source >

( scheduler: FlowMatchEulerDiscreteScheduler vae: AutoencoderKL text_encoder: CLIPTextModel tokenizer: CLIPTokenizer text_encoder_2: T5EncoderModel tokenizer_2: T5TokenizerFast transformer: FluxTransformer2DModel )

Parameters

The Flux pipeline for image inpainting.

Reference: https://blackforestlabs.ai/announcing-black-forest-labs/

__call__

< source >

( prompt: Union = None prompt_2: Union = None image: Union = None height: Optional = None width: Optional = None strength: float = 0.6 num_inference_steps: int = 28 timesteps: List = None guidance_scale: float = 7.0 num_images_per_prompt: Optional = 1 generator: Union = None latents: Optional = None prompt_embeds: Optional = None pooled_prompt_embeds: Optional = None output_type: Optional = 'pil' return_dict: bool = True joint_attention_kwargs: Optional = None callback_on_step_end: Optional = None callback_on_step_end_tensor_inputs: List = ['latents'] max_sequence_length: int = 512 ) → ~pipelines.flux.FluxPipelineOutput or tuple

Parameters

Returns

~pipelines.flux.FluxPipelineOutput or tuple

~pipelines.flux.FluxPipelineOutput if return_dictis True, otherwise a tuple. When returning a tuple, the first element is a list with the generated images.

Function invoked when calling the pipeline for generation.

Examples:

import torch

from diffusers import FluxImg2ImgPipeline from diffusers.utils import load_image

device = "cuda" pipe = FluxImg2ImgPipeline.from_pretrained("black-forest-labs/FLUX.1-schnell", torch_dtype=torch.bfloat16) pipe = pipe.to(device)

url = "https://raw.githubusercontent.com/CompVis/stable-diffusion/main/assets/stable-samples/img2img/sketch-mountains-input.jpg" init_image = load_image(url).resize((1024, 1024))

prompt = "cat wizard, gandalf, lord of the rings, detailed, fantasy, cute, adorable, Pixar, Disney, 8k"

images = pipe( ... prompt=prompt, image=init_image, num_inference_steps=4, strength=0.95, guidance_scale=0.0 ... ).images[0]

encode_prompt

< source >

( prompt: Union prompt_2: Union device: Optional = None num_images_per_prompt: int = 1 prompt_embeds: Optional = None pooled_prompt_embeds: Optional = None max_sequence_length: int = 512 lora_scale: Optional = None )

Parameters

FluxInpaintPipeline

class diffusers.FluxInpaintPipeline

< source >

( scheduler: FlowMatchEulerDiscreteScheduler vae: AutoencoderKL text_encoder: CLIPTextModel tokenizer: CLIPTokenizer text_encoder_2: T5EncoderModel tokenizer_2: T5TokenizerFast transformer: FluxTransformer2DModel )

Parameters

The Flux pipeline for image inpainting.

Reference: https://blackforestlabs.ai/announcing-black-forest-labs/

__call__

< source >

( prompt: Union = None prompt_2: Union = None image: Union = None mask_image: Union = None masked_image_latents: Union = None height: Optional = None width: Optional = None padding_mask_crop: Optional = None strength: float = 0.6 num_inference_steps: int = 28 timesteps: List = None guidance_scale: float = 7.0 num_images_per_prompt: Optional = 1 generator: Union = None latents: Optional = None prompt_embeds: Optional = None pooled_prompt_embeds: Optional = None output_type: Optional = 'pil' return_dict: bool = True joint_attention_kwargs: Optional = None callback_on_step_end: Optional = None callback_on_step_end_tensor_inputs: List = ['latents'] max_sequence_length: int = 512 ) → ~pipelines.flux.FluxPipelineOutput or tuple

Parameters

Returns

~pipelines.flux.FluxPipelineOutput or tuple

~pipelines.flux.FluxPipelineOutput if return_dictis True, otherwise a tuple. When returning a tuple, the first element is a list with the generated images.

Function invoked when calling the pipeline for generation.

Examples:

import torch from diffusers import FluxInpaintPipeline from diffusers.utils import load_image

pipe = FluxInpaintPipeline.from_pretrained("black-forest-labs/FLUX.1-schnell", torch_dtype=torch.bfloat16) pipe.to("cuda") prompt = "Face of a yellow cat, high resolution, sitting on a park bench" img_url = "https://raw.githubusercontent.com/CompVis/latent-diffusion/main/data/inpainting_examples/overture-creations-5sI6fQgYIuo.png" mask_url = "https://raw.githubusercontent.com/CompVis/latent-diffusion/main/data/inpainting_examples/overture-creations-5sI6fQgYIuo_mask.png" source = load_image(img_url) mask = load_image(mask_url) image = pipe(prompt=prompt, image=source, mask_image=mask).images[0] image.save("flux_inpainting.png")

encode_prompt

< source >

( prompt: Union prompt_2: Union device: Optional = None num_images_per_prompt: int = 1 prompt_embeds: Optional = None pooled_prompt_embeds: Optional = None max_sequence_length: int = 512 lora_scale: Optional = None )

Parameters

FluxControlNetInpaintPipeline

class diffusers.FluxControlNetInpaintPipeline

< source >

( scheduler: FlowMatchEulerDiscreteScheduler vae: AutoencoderKL text_encoder: CLIPTextModel tokenizer: CLIPTokenizer text_encoder_2: T5EncoderModel tokenizer_2: T5TokenizerFast transformer: FluxTransformer2DModel controlnet: Union )

Parameters

The Flux controlnet pipeline for inpainting.

Reference: https://blackforestlabs.ai/announcing-black-forest-labs/

__call__

< source >

( prompt: Union = None prompt_2: Union = None image: Union = None mask_image: Union = None masked_image_latents: Union = None control_image: Union = None height: Optional = None width: Optional = None strength: float = 0.6 padding_mask_crop: Optional = None timesteps: List = None num_inference_steps: int = 28 guidance_scale: float = 7.0 control_guidance_start: Union = 0.0 control_guidance_end: Union = 1.0 control_mode: Union = None controlnet_conditioning_scale: Union = 1.0 num_images_per_prompt: Optional = 1 generator: Union = None latents: Optional = None prompt_embeds: Optional = None pooled_prompt_embeds: Optional = None output_type: Optional = 'pil' return_dict: bool = True joint_attention_kwargs: Optional = None callback_on_step_end: Optional = None callback_on_step_end_tensor_inputs: List = ['latents'] max_sequence_length: int = 512 ) → ~pipelines.flux.FluxPipelineOutput or tuple

Parameters

Returns

~pipelines.flux.FluxPipelineOutput or tuple

~pipelines.flux.FluxPipelineOutput if return_dictis True, otherwise a tuple. When returning a tuple, the first element is a list with the generated images.

Function invoked when calling the pipeline for generation.

Examples:

import torch from diffusers import FluxControlNetInpaintPipeline from diffusers.models import FluxControlNetModel from diffusers.utils import load_image

controlnet = FluxControlNetModel.from_pretrained( ... "InstantX/FLUX.1-dev-controlnet-canny", torch_dtype=torch.float16 ... ) pipe = FluxControlNetInpaintPipeline.from_pretrained( ... "black-forest-labs/FLUX.1-schnell", controlnet=controlnet, torch_dtype=torch.float16 ... ) pipe.to("cuda")

control_image = load_image( ... "https://huggingface.co/InstantX/FLUX.1-dev-Controlnet-Canny-alpha/resolve/main/canny.jpg" ... ) init_image = load_image( ... "https://raw.githubusercontent.com/CompVis/latent-diffusion/main/data/inpainting_examples/overture-creations-5sI6fQgYIuo.png" ... ) mask_image = load_image( ... "https://raw.githubusercontent.com/CompVis/latent-diffusion/main/data/inpainting_examples/overture-creations-5sI6fQgYIuo_mask.png" ... )

prompt = "A girl holding a sign that says InstantX" image = pipe( ... prompt, ... image=init_image, ... mask_image=mask_image, ... control_image=control_image, ... control_guidance_start=0.2, ... control_guidance_end=0.8, ... controlnet_conditioning_scale=0.7, ... strength=0.7, ... num_inference_steps=28, ... guidance_scale=3.5, ... ).images[0] image.save("flux_controlnet_inpaint.png")

encode_prompt

< source >

( prompt: Union prompt_2: Union device: Optional = None num_images_per_prompt: int = 1 prompt_embeds: Optional = None pooled_prompt_embeds: Optional = None max_sequence_length: int = 512 lora_scale: Optional = None )

Parameters

FluxControlNetImg2ImgPipeline

class diffusers.FluxControlNetImg2ImgPipeline

< source >

( scheduler: FlowMatchEulerDiscreteScheduler vae: AutoencoderKL text_encoder: CLIPTextModel tokenizer: CLIPTokenizer text_encoder_2: T5EncoderModel tokenizer_2: T5TokenizerFast transformer: FluxTransformer2DModel controlnet: Union )

Parameters

The Flux controlnet pipeline for image-to-image generation.

Reference: https://blackforestlabs.ai/announcing-black-forest-labs/

__call__

< source >

( prompt: Union = None prompt_2: Union = None image: Union = None control_image: Union = None height: Optional = None width: Optional = None strength: float = 0.6 num_inference_steps: int = 28 timesteps: List = None guidance_scale: float = 7.0 control_guidance_start: Union = 0.0 control_guidance_end: Union = 1.0 control_mode: Union = None controlnet_conditioning_scale: Union = 1.0 num_images_per_prompt: Optional = 1 generator: Union = None latents: Optional = None prompt_embeds: Optional = None pooled_prompt_embeds: Optional = None output_type: Optional = 'pil' return_dict: bool = True joint_attention_kwargs: Optional = None callback_on_step_end: Optional = None callback_on_step_end_tensor_inputs: List = ['latents'] max_sequence_length: int = 512 ) → ~pipelines.flux.FluxPipelineOutput or tuple

Parameters

Returns

~pipelines.flux.FluxPipelineOutput or tuple

~pipelines.flux.FluxPipelineOutput if return_dictis True, otherwise a tuple. When returning a tuple, the first element is a list with the generated images.

Function invoked when calling the pipeline for generation.

Examples:

import torch from diffusers import FluxControlNetImg2ImgPipeline, FluxControlNetModel from diffusers.utils import load_image

device = "cuda" if torch.cuda.is_available() else "cpu"

controlnet = FluxControlNetModel.from_pretrained( ... "InstantX/FLUX.1-dev-Controlnet-Canny-alpha", torch_dtype=torch.bfloat16 ... )

pipe = FluxControlNetImg2ImgPipeline.from_pretrained( ... "black-forest-labs/FLUX.1-schnell", controlnet=controlnet, torch_dtype=torch.float16 ... )

pipe.text_encoder.to(torch.float16) pipe.controlnet.to(torch.float16) pipe.to("cuda")

control_image = load_image("https://huggingface.co/InstantX/SD3-Controlnet-Canny/resolve/main/canny.jpg") init_image = load_image( ... "https://raw.githubusercontent.com/CompVis/stable-diffusion/main/assets/stable-samples/img2img/sketch-mountains-input.jpg" ... )

prompt = "A girl in city, 25 years old, cool, futuristic" image = pipe( ... prompt, ... image=init_image, ... control_image=control_image, ... control_guidance_start=0.2, ... control_guidance_end=0.8, ... controlnet_conditioning_scale=1.0, ... strength=0.7, ... num_inference_steps=2, ... guidance_scale=3.5, ... ).images[0] image.save("flux_controlnet_img2img.png")

encode_prompt

< source >

( prompt: Union prompt_2: Union device: Optional = None num_images_per_prompt: int = 1 prompt_embeds: Optional = None pooled_prompt_embeds: Optional = None max_sequence_length: int = 512 lora_scale: Optional = None )

Parameters

< > Update on GitHub