je-suis-tm/tron_ares_style_lora_flux2_nf4

export MODEL_NAME="diffusers/FLUX.2-dev-bnb-4bit"
export INSTANCE_DIR="/pvol/tron_ares_style"
export OUTPUT_DIR="/pvol/tron_ares_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="Tron Ares style" \
  --caption_column="text" \
  --aspect_ratio_buckets="800,1920" \
  --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" 

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/tron_ares_style_lora_flux2_nf4",
                       weight_name='pytorch_lora_weights.safetensors')
pipe.enable_model_cpu_offload()

prompt = "Tron Ares Style, two blurred motorcycles streak through a rain slicked downtown street, two searing red laser beams erupting from their tail lights, cutting through the darkness. Tall buildings loom, while parked cars and a few pedestrians populate the scene. Christmas lights twinkle amongst the tree branches overhead, a festive contrast to the urban grit. The motorcycles, embodies the city's relentless energy, a vibrant, unsettling pulse in the night."

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("tron_ares_style.png")