dual-stream-dit

Asuka-FM: 200M Dual-Stream SPRINT DiT (Flow Matching)

An ultra-efficient, 200M-parameter Dual-Stream Diffusion Transformer (DiT) trained from scratch on consumer hardware (1x NVIDIA RTX 4090 in ~50 hours) using Rectified Flow / Flow Matching, SPRINT token-dropping, and Continuous Aspect-Normalized RoPE.

Despite its compact size, the model generates high-fidelity anime art with accurate character reproduction, prompt alignment, and zero-shot camera control (panning and zoom) at inference time.

• Developed by: aipracticecafe
• Repository: deeplearningcafe/asuka-fm
• Dataset: curated-danbooru-2026 (~337K samples)
• Architecture: Dual-Stream DiT (SPRINT-aligned, SwiGLU, Sandwich RMSNorm)
• Parameters: ~200M (Backbone) + 230M (Text Encoder)
• Objective: Conditional Flow Matching (Logit-Normal timestep sampling)
• License: OpenRAIL++

<Gallery />

---

Key Architecture & Features

1. SPRINT Dual-Stream Transformer Backbone

The base architecture is based on the results from (i1: A Simple and Fully Open Recipe for Strong Text-to-Image Models)[https://arxiv.org/abs/2606.11289]

• Minimalist Conditioning: Uses prepended time tokens rather than heavy per-block AdaLN modulation, reducing parameter overhead.
• Sandwich Normalization & SwiGLU: RMSNorm placed before and after attention and FFN operations ensures training stability in bfloat16.
• SPRINT Acceleration: Employs structured group-wise 2x2 patch dropping (75% drop ratio during early pretraining), cutting compute demands by half while retaining global semantic context.

2. HDM Continuous Aspect-Normalized RoPE (Camera Control)

Based on the Patch Diffusion: Faster and More Data-Efficient implementation proposed in the HDM paper. Early experiments with discrete integer position IDs caused severe multi-head replication when deviating from trained aspect ratios. We resolve this by adopting continuous aspect-normalized coordinate mapping which was the original idea proposed in HDM:

$$r_H = \sqrt{\frac{H}{W}}, \quad r_W = \sqrt{\frac{W}{H}}$$

$$y_{\text{pos}} = \left(\frac{y_{\text{center}}}{H}\right) \cdot (2 r_H) - r_H, \quad x_{\text{pos}} = \left(\frac{x_{\text{center}}}{W}\right) \cdot (2 r_W) - r_W$$

Manipulating $(x_{\text{pos}}, y_{\text{pos}})$ directly controls the visual framing at inference time:

• X-Shift / Y-Shift: Translates the virtual camera viewport horizontally or vertically.
• Zoom: Scales spatial coordinates around the focal center.

3. Components

• Text Encoder: LiquidAI/LFM2.5-230M (unidirectional causal LM hidden states passed via a 2-layer Transformer Text Adapter).
• Autoencoder: kaiyuyue/FLUX.2-dev-vae (C=32, f=8, patch size p=2, effective image token size 16x16).

---

Training Recipe & Multi-Stage Progression (1x RTX 4090, ~100h)

The model uses progressive multi-resolution pretraining on 337,000 curated Danbooru samples with shifted square cropping and in-RAM latent caching.

• Optimizer: 8-bit AdamW (bitsandbytes), $\beta=(0.9, 0.95)$, WD=0.01
• Learning Rate: 5e-4 with warmup and cosine annealing only for last stage
• Effective Batch Size: 256 (Batch size 64x4 gradient accumulation)
• Precision: Mixed Precision (torch.bfloat16 with torch.compile)

Stage 1: 256px Base Pretraining (Epochs 0 – 170)

• Goal: Establish character concept learning, global structure, and RoPE spatial alignment.

Stage 2: 512px Resolution Adaptation (Epochs 170 – 209)

• Goal: Adapt patch receptive fields to higher frequency textures and sharper lineart.

Stage 3: 1MP Native Training (Active Run)

• Scaling from 512px to 1024px aspect ratio buckets.

---

Checkpoint Guide

---

512px Generation Outputs

Figure: 512px generations demonstrating improved micro-texture, eyes, and line definition achieved in Stage 2.

---

Inference and WebUI

The training repository has a simple Gradio UI to perform sampling with camera controls. ui_app.py

Recommended Settings

• Sampling Algorithm: Forward Euler ODE Solver (Rectified Flow)
• Inference Steps: 25 – 35
• CFG Scale: 4.5 – 6.0
• Time Shift: 1.0 – 3.0 (higher values enhance contrast/saturation)

Tag Formatting Template

1girl/1boy, character name, series name, artists, rating, quality tags, actions/clothing

Negative prompt:

sketch, monochrome, english text, realistic, traditional media, very displeasing, displeasing, bad score, worse score, bad quality,worst quality, worst detail, watermark, signature, blurry, bad hands, bad anatomy, bad fingers, extra fingers, extra limbs, deformed limbs, comic, speech bubble, lowres

Python Sampling Script

import torch
from omegaconf import OmegaConf
from src.models.factory import load_trainable_model
from src.diffusion.schedules import LinearSchedule
from src.diffusion.sampling import generate_samples
cfg = OmegaConf.load("config.yaml")
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
dtype = torch.bfloat16
unet, text_encoder, vae, tokenizer, _ = load_trainable_model(
    models_path=cfg.paths.models,
    device=device,
    dtype=dtype,
    resume_from_checkpoint="path_to_checkpoint",
    model_type="sprint_dual",
    model_cfg=cfg.models,
    autocast_dtype=dtype,
)
schedule = LinearSchedule(device=device)
sample_configs = [
    {
        "prompt": "1girl, souryuu asuka langley, neon genesis evangelion, "
                  "masterpiece, absurdres",
        "negative_prompt": "very displeasing, displeasing, bad score",
        "height": 256,
        "width": 384,
        "sample_steps": 30,
        "cfg_scale": 6.0,
        "shift": 1.0,
        "zoom": 1.0,
        "x_shift": 0.0,
        "y_shift": 0.0,
        "seed": 42,
    }
]
images = generate_samples(
    unet=unet,
    text_encoder=text_encoder,
    tokenizer=tokenizer,
    vae=vae,
    schedule=schedule,
    sample_configs=sample_configs,
    global_batch_size=1,
    diffusion_type="flow_matching",
    device=device,
    dtype=dtype,
    autocast_dtype=dtype,
    use_unet_mult=False,
    vae_mean=cfg.models.vae_mean,
    vae_std=cfg.models.vae_std,
    in_channels=cfg.models.in_channels,
    coord_system="aspect_norm",
)
images[0].save("output.png")

---

Content Advisory & Limitations

• Resolution Scope: Checkpoints up to epoch_209_* are natively optimized for 512px generation and reliably scale up to 768px/1024px with aspect-normalized RoPE. Native 1024px fidelity will be unlocked upon completion of Stage 3.
• Complex Hands: The compact 200M parameter count prioritizes prompt faithfulness and character identity; complex hand interactions may still require multiple seed attempts.
• NSFW & Sensitive Content: The model was trained on a dataset that contains uncensored anime-style art, including sensitive, questionable, and explicit material (g, s, q, e ratings). Use appropriate tag filtering during training if a SFW-only model is required.
• Legal Disclaimer: Images are sourced from Danbooru for research, educational, and generative modeling purposes under fair use. Copyright of all underlying illustrations belongs to their respective creators.

---

References

• i1: i1: A Simple and Fully Open Recipe for Strong Text-to-Image Models
• Patch Diffusion: Patch Diffusion: Faster and More Data-Efficient
• Home-made Diffusion Model from Scratch to Hatch: HDM
• Danbooru Dataset Toolkit: Dataset processed and formatted using the Danbooru Dataset Toolkit.

License

This model is licensed under the CreativeML Open RAIL++-M License. You are free to use, modify, and distribute the model weights provided you include the license text and respect the behavioral use restrictions.