caikybaldo999/zyi-256-flash

ZYI-256 Flash

A small text-to-image model (343M parameters) that generates 256×256 images. Trained from scratch in ~5 hours on a single RTX PRO 6000 (96 GB).

Model

Architecture	DiT with rectified flow (flow matching)
Parameters	343M
Hidden size / depth / heads	1024 / 20 / 16
Resolution	256×256
Latent space	DC-AE f32c32 (mit-han-lab/dc-ae-f32c32-sana-1.0-diffusers), 32×8×8 = 64 tokens
Text encoder	google/flan-t5-base (frozen, 96 tokens)
Details	2D RoPE, QK-norm, adaLN-single, SwiGLU, cross-attention, EMA weights

Training

• Data: 1,226,758 images from LucasFang/FLUX-Reason-6M (quality score ≥ 7), random short or long caption per sample
• Steps: 9,544 at batch size 1024
• Optimizer: AdamW, lr 0.0002, cosine schedule, bf16
• Hardware: 1× NVIDIA RTX PRO 6000, ~5 h total (data download, encoding and training)

Fine-tuning

These weights continue from the base model (revision v1.0-base, 32,748 steps).

• Extra steps: 10,233 (total 42,981) at batch size 1024
• New data: 2,139,181 images from LucasFang/FLUX-Reason-6M not seen by the base model
• Optimizer: AdamW, lr 0.0001, cosine schedule, bf16
• Hardware: 1× NVIDIA RTX PRO 6000, 5.7 h so far

The sample grid above is regenerated from these weights.

Usage

pip install torch diffusers transformers safetensors huggingface_hub sentencepiece

from huggingface_hub import hf_hub_download
import importlib.util

path = hf_hub_download("caikybaldo999/zyi-256-flash", "zyi256.py")
spec = importlib.util.spec_from_file_location("zyi256", path)
zyi256 = importlib.util.module_from_spec(spec); spec.loader.exec_module(zyi256)

pipe = zyi256.ZyiPipeline("caikybaldo999/zyi-256-flash")
images = pipe("a misty mountain lake at sunrise with pine trees", n=4, steps=28, cfg=4.5, seed=0)
images[0].save("out.png")

English prompts work best (the training captions are in English).

Limitations

• Good at landscapes, interiors, food, objects and painting styles.
• Weak at animal anatomy, faces, hands and readable text.
• Does not follow spatial relations reliably ("a red ball on top of a blue box").