---
language: en
license: apache-2.0
library_name: diffusers
pipeline_tag: text-to-image
tags:
- fwkv
- fwkv-vision
- text-to-image
- rectified-flow
- diffusion-transformer
- dit
- custom-architecture
- linear-transformer
---
This image was not generated by this model
# 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).
- **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 noise `x₀` to data `x₁` along a straight line. Training minimizes MSE against the ground‑truth velocity `x₁ − 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:
```python
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](https://huggingface.co/spaces/FlameF0X/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.json`
- `model.safetensors`
- `modeling_fwkv_vision.py`
- **Auto‑mapping:** The `config.json` includes `"auto_map": { "AutoModel": "modeling_fwkv_vision.FWKVVisionModel" }`, so loading with `trust_remote_code=True` will automatically locate the correct class.