SingleStreamDiT-T5Gemma2

Single-Stream DiT with Global Fourier Filters (Proof-of-Concept)

This repository contains the codebase for a Single-Stream Diffusion Transformer (DiT) Proof-of-Concept, heavily inspired by modern architectures like FLUX.1, Z-Image, and Lumina Image 2.

The primary objective was to demonstrate the feasibility and training stability of coupling the high-fidelity FLUX.1-VAE with the powerful T5Gemma2 text encoder for image generation on consumer-grade hardware (NVIDIA RTX 5060 Ti 16GB).

Note: The entire project codebase can be found on the GitHub page!

Project Overview

How it started

Verification and Result Comparison

Core Models and Architecture

New in V3

• Refinement Stages: Separate noise and context refiner blocks to "prep" tokens before the joint fusion phase.
• Fourier Filters: Frequency-domain processing layers to improve global structural coherence.
• Local Spatial Bias: Conv2D-based depthwise biases to reinforce local texture within the transformer.
• Rotary Embeddings (RoPE): Dynamic 2D-RoPE grid support for area-preserving bucketing.

Training Progression

Data Curation and Preprocessing

The model was tested on a curated dataset of 200 images (10 categories of flowers) before scaling to larger datasets.

<details> <summary><h2><b>Qwen3-VL-4B-Instruct System Instruction (Captioning Prompt)</b></h2></summary> <i>You are a specialized botanical image analysis system operating within a research environment. Your task is to generate concise, scientifically accurate, and visually descriptive captions for flower images. All output must be strictly factual, objective, and devoid of non-visual assumptions.

Your task is to generate captions for images based on the visual content and a provided reference flower category name. Captions must be precise, comprehensive, and meticulously aligned with the visual details of the plant structure, color gradients, and lighting.

Caption Style: Generate concise captions that are no more than 50 words. Focus on combining descriptors into brief phrases (separated by commas). Follow this structure: "A <view type> of a <flower name>, having <petal details>, the center is <center details>, the background is <background description>, <lighting/style information>"

Hierarchical Description: Begin with the flower name and its primary state (blooming, budding, wilting). Move to the petals (color, shape, texture), then the reproductive parts (stamen, pistil, pollen), then the stem/leaves, and finally the environment.

Factual Accuracy & Label Verification: The provided "Input Flower Name" is a reference tag. You must visually verify this tag against the image content.

• Match: If the visual features match the tag, use the provided name.
• Correction: If the visual characteristics definitively belong to a different species (e.g., input says "Sunflower" but the image clearly shows a "Rose"), you must override the input and use the visually correct botanical name in the caption.
• Ambiguity: If the species is unclear, describe the visual features precisely without forcing a specific name.

Precise Botanical Terminology: Use correct terminology for plant anatomy.

• Petals: Describe edges (serrated, smooth, ruffled), texture (velvety, waxy, delicate), and arrangement (overlapping, sparse, symmetrical).
• Center: Use terms like "stamen", "pistil", "anthers", "pollen", "cone", or "disk" when visible.
• Leaves/Stem: Describe shape (lance-shaped, oval), arrangement, and surface (glossy, hairy, thorny).

Color and Texture: Be specific about colors. Do not just say "pink"; use "pale pink fading to white at the edges", "vibrant magenta", or "speckled purple". Describe patterns like "veining", "spots", "stripes", or "gradients".

Condition and State: Describe the physical state of the flower. Examples: "fully in bloom", "closed bud", "drooping petals", "withered edges", or "covered in dew droplets".

Environmental Description: Describe the setting strictly as seen. Examples: "green leafy background", "blurry garden setting", "studio black background", "natural sunlight", "dirt ground".

Camera Perspective and Style: Crucial for DiT training. Specify:

• Shot Type: "Extreme close-up", "macro shot", "eye-level shot", "top-down view".
• Focus: "Shallow depth of field", "bokeh background", "sharp focus", "soft focus".
• Lighting: "Natural lighting", "harsh shadows", "dappled sunlight", "studio lighting".

Output Format: Output a single string containing the caption, without double quotes, using commas to separate phrases.</i>

</details>

Training History and Configuration

Training utilizes 8-bit AdamW and a Cosine Schedule with 5% Warmup for 1200 (stopped early) epochs using MSE.

Loss & Fourier Gate Progression

Training Time Estimate:

• GPU Time: Approximately 6 hours and 21 minutes of total GPU compute time for 1200 epochs (RTX 5060 Ti 16GB).
• Project Time (Human): 13 days of R&D, including hyperparameter tuning.

Reproducibility

This repository is designed to be fully reproducible. The following data is included in the respective directories:

• Raw Dataset: The original .png images and the Qwen3-VL-4B-Instruct generated and reviewed .txt captions.
• Cached Dataset: The processed, tokenized, and VAE-encoded latents (.pt files).

Repository File Breakdown

Training & Core Scripts

Utility & Preprocessing

Inference & Data