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