README.md

---
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
---

<div align="center">
    <img src="https://cdn-uploads.huggingface.co/production/uploads/6615494716917dfdc645c44e/O_kbP5ljh6m1v_Zrpe1bp.png" alt="FWKV-Image Banner" style="width: 100%; max-width: 100%; height: auto; display: inline-block; margin-bottom: 0.5em; margin-top: 0.5em; reading-order: 20px; border-radius: 20px;"/>
    <i> This image was <b>not</b> generated by this model</i>
</div>
<br>

# 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.