FWKV-Image
A ~40M‑parameter (trained) diffusion transformer for text‑to‑image generation that replaces standard self‑attention with a bidirectional per‑channel leaky integrator (FWKV) and is trained as a rectified‑flow velocity field over the latent space of a frozen VAE. It is a research experiment exploring how far linear‑RNN token mixing can go on visual generation tasks.
The model generates 256×256 images from text prompts encoded by a frozen CLIP ViT‑B/32 text encoder, using a DiT‑style patchify‑unpatchify pipeline with adaLN‑zero conditioning.
Model Description
- Architecture: 12 stacked FWKV‑DiT blocks at width 384, with 6 heads each. Each block replaces standard self‑attention with a bidirectional decayed accumulator:
- Forward scan:
stateₜ = W·stateₜ₋₁ + kₜ·vₜ - Backward scan: same recurrence run in reverse
- Output: sum of both directions (every patch sees every other patch with distance‑weighted decay)
- Computed exactly via a vectorised O(log T) parallel scan (no approximations, no O(T²) attention matrix).
- Forward scan:
- Cross‑attention: Standard multi‑head cross‑attention to the 77 CLIP text token embeddings (512‑dim) is retained inside each block.
- Conditioning: adaLN‑zero modulation derived from sinusoidal timestep embedding plus pooled CLIP text embedding.
- Patchify: Latent 4×32×32 is split into 256 patches of 2×2, mapped to 384‑d tokens. Positional embeddings are learned.
- Rectified flow: The model predicts a velocity field
v(xₜ, t, text)that transports Gaussian noisex₀to datax₁along a straight line. Training minimizes MSE against the ground‑truth velocityx₁ − x₀. - Factorised design: 4‑channel VAE latent → 384‑dim patch tokens → 4× expansion MLP.
- Context: 256 patch tokens (non‑causal, bidirectional). No image‑level positional embeddings beyond patch positions.
- Text encoder: Frozen CLIP ViT‑B/32 (
openai/clip-vit-base-patch32). - VAE: Frozen
stabilityai/sd-vae-ft-mse(4×32×32 latents, scaling factor 0.18215).
Uses
Direct Use
FWKV‑Image is intended for research on efficient diffusion transformers and for educational demonstrations of linear‑RNN architectures applied to visual generation. You can generate images from text prompts using the provided inference code.
Out‑of‑Scope Use
- This model is not suitable for any production or safety‑critical application.
- It has not been aligned with RLHF or other safety filters and may generate inappropriate or harmful content.
- The limited size and training data mean image fidelity, prompt adherence, and diversity are far below commercial text‑to‑image systems.
- Generated images should not be relied upon for factual or medical accuracy.
Bias, Risks, and Limitations
- Trained on a relatively small dataset (100 k pairs), the model can produce repetitive outputs, artifacts, or fail to follow complex prompts.
- Biases present in the training data (e.g. stereotypical depictions of people, occupations, or cultures) are likely reflected in generated images.
- The 256×256 output resolution and small DiT capacity limit fine detail and text rendering quality.
- As a research checkpoint, sampling hyperparameters (CFG scale, steps) have not been exhaustively tuned for all prompt categories.
How to Get Started
The model relies on a custom architecture. To load it, you must provide the modeling_fwkv_vision.py file (found in the repository) and trust the remote code:
from transformers import AutoModel
from diffusers import AutoencoderKL
from transformers import CLIPTokenizer
model = AutoModel.from_pretrained(
"FWKV/FWKV-Image",
trust_remote_code=True
).eval().cuda()
tokenizer = CLIPTokenizer.from_pretrained("openai/clip-vit-base-patch32")
image = model.generate(
tokenizer=tokenizer,
prompt="a red fox sitting in a snowy forest, digital art",
steps=50,
cfg_scale=4.0,
seed=42
)
image.save("output.png")
The unified checkpoint contains the trained DiT, the frozen CLIP text encoder, and the frozen VAE weights in a single model.safetensors, so no separate AutoencoderKL.from_pretrained(...) calls are required.
Training Details
Dataset
- Primary source:
jackyhate/text-to-image-2M(streamed, first 100 k valid pairs) - Fallback source:
HuggingFaceM4/COCO - Pre‑processing: Images resized to 256×256, encoded once to VAE latents (deterministic posterior mean) and cached to disk to avoid repeated VAE forward passes.
Training Procedure
| Hyperparameter | Value |
|---|---|
| Architecture | 12 FWKV‑DiT blocks, d_model=384, patch=2, 256 tokens |
| Heads | 6 |
| FFN multiplier | 4 |
| WKV decay floor | 0.05 |
| Objective | Rectified flow (velocity MSE) |
| Batch size | 32 |
| Learning rate | 1×10⁻⁴ (cosine schedule) |
| Weight decay | 0.0 |
| Gradient clipping | 1.0 |
| Optimizer | AdamW (β₁=0.9, β₂=0.95) |
| Precision | bfloat16 mixed (CUDA only) |
| Epochs | 5 |
| Effective examples | 100 000 |
| Hardware | 1× NVIDIA GPU |
Note: the training script itself is not public; only the final weights and inference code are released.
Evaluation
See Tiny T2I Leaderboard.
The model is intended as an architectural proof‑of‑concept rather than a competitive production image generator.
Environmental Impact
The training run consumed a single consumer/entry‑level NVIDIA GPU for a small number of epochs on 100 k examples. The total energy footprint is estimated to be well under 1 kWh and corresponding CO₂eq emissions are negligible (on the order of 0.1–0.3 kg CO₂eq assuming average grid carbon intensity).
Technical Specifications
- Model type: Diffusion transformer (DiT) with linear‑RNN token mixing
- Trained parameters: ~40 million (FWKV‑DiT backbone only)
- Total checkpoint size: Larger (includes frozen CLIP + VAE weights bundled in the same
model.safetensors) - Checkpoint format: PyTorch
safetensors - Required files in the repo:
config.jsonmodel.safetensorsmodeling_fwkv_vision.py
- Auto‑mapping: The
config.jsonincludes"auto_map": { "AutoModel": "modeling_fwkv_vision.FWKVVisionModel" }, so loading withtrust_remote_code=Truewill automatically locate the correct class.