akrao9/Boomer-T2I

Boomer FLA

Boomer FLA is a 657M parameter text-to-image diffusion model that generates 1024×1024px images from text prompts.

Instead of standard quadratic self-attention, it uses GatedDeltaNet , a bidirectional Flash Linear Attention mixer , as the backbone of its transformer blocks. This keeps memory flat regardless of sequence length. Every 6th block adds a full SDPA layer for global spatial coherence.

Text conditioning uses Gemma 4 2B (1536-dim embeddings, up to 300 tokens). Decoding uses the DC-AE f32c32 VAE with 32× spatial compression, producing 32×32 latents from 1024px images. Inference uses STORK-2 flow sampling (Tan et al., 2025).

Trained on 3.8 Million JourneyDB 512 px latents and finetuned on maxu/Fine-T2i 600k high-resolution latents at 1024px.

---

Sample Outputs

1. Misty Pine Forest: A cinematic, wide-angle shot of a misty pine forest at sunrise, deep green valleys, soft morning light piercing through fog, photorealistic, 8k resolution.
2. Black Sand Beach: A dramatic black sand beach in Iceland, towering basalt columns, massive white waves crashing on the shore, moody overcast sky, high detail landscape architecture.
3. Alpine Lake: A serene alpine lake reflecting jagged snow-capped mountain peaks, crystal clear turquoise water, vibrant wildflower meadows in the foreground, golden hour lighting.
4. Tuscan Hills: A sweeping view of rolling terracotta hills in Tuscany, isolated cypress trees lining a dirt road, warm late afternoon sun casting long shadows, classic landscape photography.
5. Desert Oasis: A futuristic desert oasis at dusk, neon-infused crystalline structures contrasting with deep orange sand dunes, hyper-detailed, sprawling cinematic atmosphere.
6. Hidden Lagoon: A majestic waterfall cascading down a sheer mossy cliff into a hidden tropical lagoon, lush emerald foliage, sunbeams cutting through the canopy, long exposure water effect.

Architecture

---

Property	Value
Parameters	657M
Backbone	Bidirectional GatedDeltaNet (Flash Linear Attention)
Depth	24 layers
Hidden dim	896
Heads	14
Image attention	Every 6th layer (full SDPA + 2D RoPE)
Patch size	1 — one token per latent pixel (256 tokens at 512px, 1024 tokens at 1024px)
Text encoder	Gemma 4 2B (google/gemma-4-E2B-it)
VAE	DC-AE f32c32 (mit-han-lab/dc-ae-f32c32-sana-1.1-diffusers)
Sampler	STORK-2, 32 steps
Dtype	bfloat16

---

Training details

Setting	Value
Pre-train dataset	JourneyDB (~3.8M images, 512px, patch size 1)
Fine-tune dataset	FineT2I (~600k images, 1024px, patch size 1)
Optimizer	Fused AdamW
Hardware	RTX PRO 6000 BLACKWELL SERVER EDITION
Precision	bfloat16

---

VRAM and RAM requirements

Measured at 1024×1024px, bfloat16, STORK-2, CFG batch=2.

Component	VRAM
DiT weights (EMA, bf16)	1.25 GB
Gemma 4 2B text encoder	8.62 GB
Denoising peak (CFG on)	1.36 GB
VAE decode peak	3.51 GB

Mode	Peak VRAM	Minimum GPU
Condition-cache — pre-encoded embeddings, no text encoder in VRAM	4.76 GB	RTX 3060 8GB, T4
Fresh-prompt — text encoder + DiT + VAE together	13.38 GB	RTX 3090, A100

System RAM: loading the text encoder (Gemma 4 2B) requires ~9 GB of system RAM even when using GPU. For condition cache mode, encoding can be done on CPU with ~9 GB RAM the generation step then needs only 5 GB VRAM.

---

Usage

Install

pip install -U torch diffusers transformers accelerate safetensors torchvision scipy
pip install git+https://github.com/fla-org/flash-linear-attention.git

Generate

import torch
from diffusers import DiffusionPipeline

pipe = DiffusionPipeline.from_pretrained(
    "akrao9/Boomer-T2I",
    custom_pipeline="pipeline_boomer",
    trust_remote_code=True,
    torch_dtype=torch.bfloat16,
).to("cuda")

image = pipe("A hyper-detailed, cinematic landscape photography shot of a pristine, mirror-like alpine lake nestled deeply between towering, jagged snow-capped mountain peaks. The scene is captured during the perfect golden hour, with the low-angled warm sun casting deep amber and violet hues across the rugged granite rock faces. In the foreground, vibrant clusters of purple lupines and orange poppies dot a lush emerald meadow that meets the crystal-clear turquoise edge of the water. Wisps of soft, low-hanging morning mist drift lazily across the lake's surface, breaking the perfect reflection of the monumental peaks above. Shot on 35mm lens, ultra-sharp focus, dramatic depth of field, 8k resolution, path-traced lighting textures.")[0]
image.save("output.png")

Note

The transformer weights (1.25 GB) are downloaded from this repo. The VAE and Gemma 4 2B text encoder are fetched automatically from their upstream HuggingFace repos on first use (10 GB total). Accept the Gemma Terms of Use and run hf auth login before first use.

Parameters

image = pipe(
    "a rocky coastline at sunset with crashing waves",
    steps=32,        # denoising steps — 32 is recommended with STORK-2
    cfg_scale=4.5,   # classifier-free guidance scale (4.0–5.0)
    cfg_rescale=0.5, # reduces over-saturation at high CFG
    seed=42,
)[0]

Low VRAM — condition cache mode

Encode prompts once on any machine (including CPU), save the embedding, then generate with only the 1.25 GB DiT loaded. Peak VRAM drops from 13.38 GB → 4.76 GB.

# Step 1 — encode on any machine (even CPU with 9GB RAM)
import torch
from transformers import AutoModelForCausalLM, AutoProcessor

TE_REPO      = "google/gemma-4-E2B-it"
tokenizer    = AutoProcessor.from_pretrained(TE_REPO)
text_encoder = AutoModelForCausalLM.from_pretrained(
    TE_REPO, torch_dtype=torch.bfloat16
).get_decoder()

tokens = tokenizer(
    "a mountain lake surrounded by alpine peaks",
    max_length=300, padding="max_length",
    truncation=True, return_tensors="pt",
)
with torch.inference_mode():
    hidden = text_encoder(
        tokens["input_ids"], attention_mask=tokens["attention_mask"]
    )[0]
    idx  = [0] + list(range(-299, 0))
    emb  = hidden[:, idx]
    mask = tokens["attention_mask"][:, idx]

torch.save({"emb": emb.cpu(), "mask": mask.cpu()}, "condition.pt")

# Step 2 — generate on low-VRAM GPU
# Load only the DiT weights; skip text encoder by providing embeddings directly
import torch
from huggingface_hub import snapshot_download
import sys

local = snapshot_download("akrao9/Boomer-T2I")
sys.path.insert(0, local)
from pipeline_boomer import BoomerPipeline
from modeling_boomer_fla import BoomerFLADiT

transformer = BoomerFLADiT.from_pretrained(local, subfolder="transformer").to("cuda", dtype=torch.bfloat16)
pipe = BoomerPipeline(transformer=transformer).to("cuda")
# pipe.text_encoder and pipe.vae remain None until __call__ — override with pre-encoded emb manually

---

Capabilities and limitations

Trained on JourneyDB (pretrain, 512px) and FineT2I (finetune, 1024px).

• Strongest at — dramatic landscapes, natural environments, and scenic/architectural scenes
• Also works for — many everyday objects when embedded in a scene (quality varies)
• Not reliable — humans, portraits, and people (faces and anatomy are inconsistent)
• Less reliable — animals, legible text in images, very fine small details
• Landscapes can show a painterly/HDR bias from heavily post-processed stock images in the training set
• Not safety filtered — outputs may reflect biases in the training data
• Maximum tested resolution: 1024×1024px

---

Acknowledgements

• STORK-2 — inference sampling (Tan et al., 2025)
• Plateau logit-normal — training timestep distribution from FLUX.2 representation comparison (Black Forest Labs, 2025). Boomer uses μ=0, σ=1 with flow shift 1.5.
• Sana — flow shift, logit-normal sampling, forward noising, and related helpers adapted from NVlabs/Sana (boomer/sana_flow.py, boomer/sana_latent_cache.py); several DiT blocks follow Sana conventions (Xie et al., 2024)

---

License

The Boomer FLA model weights are released for research and personal use.

Upstream component licenses:

• DC-AE VAE: mit-han-lab/dc-ae-f32c32-sana-1.1-diffusers
• Gemma 4 text encoder: Gemma Terms of Use