shortcut to loading flux using diffusers/accelerate on 8gb vram pc, with minimal ram usage by offloading main transformer block to disk. (with nvme ssd generation speed slowdowns appear to be similar to cpu -> gpu transfer speed slowdowns, although this will tank speeds on slower storage drives)
import torch
from diffusers import FluxPipeline
from accelerate import disk_offload
import gc
import os
def load_flux_with_disk_offload():
model_id = "black-forest-labs/FLUX.1-dev"
offload_dir = "./flux_offload"
# Create offload directory
os.makedirs(offload_dir, exist_ok=True)
pipe = FluxPipeline.from_pretrained(
model_id,
torch_dtype=torch.bfloat16,
device_map=None,
low_cpu_mem_usage=True,
)
pipe.text_encoder.to("cuda:0")
pipe.text_encoder_2.to("cuda:0")
pipe.vae.to("cuda:0")
print("Offloading transformer to disk...")
disk_offload(
model=pipe.transformer,
offload_dir=os.path.join(offload_dir, "transformer"),
execution_device="cuda:0", # Will be moved to GPU when needed
)
return pipe
def main():
pipe = load_flux_with_disk_offload()
# Enable memory optimizations
pipe.enable_attention_slicing()
pipe.enable_vae_slicing()
torch.cuda.empty_cache()
gc.collect()
prompt = "a fantasy landscape with mountains and rivers, trending on artstation"
# Generate with conservative settings
image = pipe(
prompt,
num_inference_steps=20,
guidance_scale=3.0,
height=512,
width=512,
max_sequence_length=128,
).images[0]
image.save("out.png")
if name == "__main__":
main()