The training is based on https://github.com/huggingface/diffusers/blob/main/examples/dreambooth/test_dreambooth_lora_flux2.py. Everything in this training script needs to be set at the same torch dtype as the script was designed for unquantized version which will export in float32. The training took 11.5 hours on A100 80GB with max VRAM consumption at 32GB. The inference consumes 35GB of VRAM. To avoid running low on VRAM, both transformers and text_encoder were quantized.

Train

export MODEL_NAME="diffusers/FLUX.2-dev-bnb-4bit"
export INSTANCE_DIR="/pvol/marvel_zombies_style"
export OUTPUT_DIR="/pvol/marvel_zombies_style_lora_flux2_nf4"
export Q_DIR="/pvol/quantization_config.json"
accelerate config default
accelerate launch train_dreambooth_lora_flux2.py \
  --pretrained_model_name_or_path=$MODEL_NAME  \
  --mixed_precision="bf16" \
  --dataset_name=$INSTANCE_DIR \
  --output_dir=$OUTPUT_DIR \
  --bnb_quantization_config_path=$Q_DIR \
  --gradient_checkpointing \
  --cache_latents \
  --instance_prompt="Marvel Zombies style" \
  --caption_column="text" \
  --aspect_ratio_buckets="720,1728" \
  --center_crop \
  --train_batch_size=1 \
  --guidance_scale=1 \
  --use_8bit_adam \
  --offload \
  --checkpointing_steps=100 \
  --gradient_accumulation_steps=4 \
  --optimizer="adamW" \
  --learning_rate=1e-4 \
  --lr_scheduler="constant" \
  --lr_warmup_steps=100 \
  --max_train_steps=1500 \
  --rank=4 \
  --seed="0"

Usage

python

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.float32 #only supports float32 when using train_dreambooth_lora_flux2.py 

transformer = Flux2Transformer2DModel.from_pretrained(
  repo_id, subfolder="transformer", torch_dtype=torch_dtype, device_map="cuda:0"
)
text_encoder = Mistral3ForConditionalGeneration.from_pretrained(
  repo_id, subfolder="text_encoder", dtype=torch_dtype, device_map="cuda:0"
)

pipe = Flux2Pipeline.from_pretrained(
  repo_id, transformer=transformer, text_encoder=text_encoder, torch_dtype=torch_dtype
)
pipe.load_lora_weights("je-suis-tm/marvel_zombies_style_lora_flux2_nf4",
                       weight_name='pytorch_lora_weights.safetensors')
pipe.enable_model_cpu_offload()

prompt = "Marvel Zombies Style, an Aztec god with a skull face, depicted in a full-body view. He stands firmly, gazing intensely and pointing directly at the viewer, conveying a powerful sense of anger. The god wears an intricately designed headdress made of vibrant feathers and jewels, featuring ancient symbols. His elaborate armor, decorated with traditional Aztec motifs, covers his entire body, showcasing intricate details. His eyes glow a menacing red, enhancing his intimidating presence. The background is a deep, shadowy darkness, emphasizing the eerie atmosphere, with subtle glimpses of ancient ruins barely visible in the dim light. The textures of his skin and the shine of the jewels contribute an extra layer of realism, while the full-body composition emphasizes his power and stature"

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,
  height=800, 
  width=1920, #the movie was in a weird resolution 800*1920, the resolution needs to be divided by 16
).images[0]

image.save("marvel_zombies_style.png")

Trigger words

You should use Marvel Zombies Style to trigger the image generation.

Download model

Download them in the Files & versions tab.