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

Image editing (mega update)

The same model now does text-to-image and instruction-based image editing at 256×256, still with flan-T5-base. The source image is encoded by the DC-AE and its 64 tokens join the same sequence as the 64 target tokens (in-context conditioning, shared 2D RoPE positions, one learned "source" embedding — the only new weights). Editing uses text + image classifier-free guidance (InstructPix2Pix-style).

Each row: source · edited · reference.

from PIL import Image
pipe = zyi256.ZyiPipeline("caikybaldo999/zyi-256-flash")
pipe("a red fox sitting in a snowy forest", n=4, seed=0)             # text -> image (unchanged)
pipe.edit(Image.open("photo.jpg"), "make it winter, add snow")[0]  # image editing (English instructions)

• Training: +18,613 steps from the previous weights (revision v1.2-edit), 40% text-to-image / 60% editing
• Data: 1,000,000 images from LucasFang/FLUX-Reason-6M and 190,988 edit pairs from TIGER-Lab/OmniEdit-Filtered-1.2M (MIT): add, remove, swap, attribute, environment and style edits
• Hardware: 1× NVIDIA RTX PRO 6000, 6.5 h
• Images are center-cropped to a square; low resolution limits fine edits (text, small objects, faces).

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