Echo-Vision-1

A two-stage coarse-to-fine latent diffusion model (rectified flow), trained entirely from scratch on a single consumer GPU (RTX 4060, 8 GB VRAM).

Echo-Vision-1 generates 768×768 images by first "sketching" a low-resolution draft of the scene and then refining it into a high-resolution image with a dedicated super-resolution-style diffusion stage. This plan-then-paint design improves compositional coherence (fewer overlapping/phantom objects) and sharpness compared to a single small UNet trained end-to-end.

TL;DR

• 🧠 2 UNets: Draft (384 px, 36.2 M params) → Final (768 px, 37.6 M params) ≈ 74 M total
• Objective: rectified flow (velocity prediction), uniform timestep sampling, plain MSE
• ✍️ Text conditioning: frozen CLIP ViT-B/32 (77 tokens, 512-d)
• 🖼️ Latent space: frozen sd-vae-ft-ema
• 🏋️ Trained from scratch on 45,251 captioned images in ~6.5 h on one RTX 4060
• ⚡ Inference: ~35 Euler steps total, CFG ≈ 2–3, runs comfortably in 8 GB VRAM

---

1. Model Description

Most small text-to-image models try to generate a full-resolution image in a single pass. At low capacity this forces the network to solve layout, object identity, textures and fine detail simultaneously – which leads to the classic failure modes of small diffusion models: texture stitching, overlapping objects, phantom structures and blurry ("MSE-averaged") outputs.

Echo-Vision-1 instead splits the problem into two easier sub-tasks:

1. Draft stage – generates a tiny 48×48 latent (384 px) that fixes the global composition: where is the sky, the ground, the subject? At this scale the network cannot stitch textures; it must learn genuine scene-layout statistics.
2. Final stage – a conditional refiner that receives the upsampled draft latent plus the text embedding and "paints" the 96×96 latent (768 px) with real shapes, colors and sharp detail (a diffusion-based super-resolution step).

Both stages are trained with rectified flow (straight noise→data paths, velocity prediction) rather than DDPM ε-prediction, giving faster convergence and crisper samples with fewer inference steps.

Architecture

Component	Details
Stage 1 "draft"	UNet2DConditionModel, in=4 ch, out=4 ch, channels (64,128,192,256), 2 layers/block, 36.2 M params, latent 48×48 (384 px)
Stage 2 "final"	UNet2DConditionModel, in=8 ch (4 noise + 4 upsampled draft), out=4 ch, channels (96,128,192,256), gradient checkpointing, 37.6 M params, latent 96×96 (768 px)
Text encoder	frozen openai/clip-vit-base-patch32 (77 tokens, 512-d)
VAE	frozen stabilityai/sd-vae-ft-ema, scaling factor 0.18215
Attention	AttnProcessor2_0 (Flash/SDPA), channels-last memory format
Total params	≈ 74 M (both UNets)

Training objective

Rectified flow with linear interpolation x_t = (1−t)·x₀ + t·ε, uniform t ~ U(0,1), velocity target v = ε − x₀, loss MSE(v_pred, v). Classifier-free guidance dropout 15 % (text) and 10 % (draft-condition on the refiner). Horizontal-flip augmentation 50 % (applied consistently to target and condition).

---

2. Training Details

Hyperparameter	Value
Optimizer	AdamW (β 0.9/0.999, wd 1e-2, fused)
Learning rate	3e-4, cosine decay, 6 % warmup, min-LR ratio 0.05
Precision	bf16 autocast (training), fp16 (latents/VAE)
Batch size	48 (draft) / 4 (final)
EMA	decay 0.9995 (used for released weights)
Grad clip	1.0
Epochs	26–29 (draft) / 4 (final)
Wall time	≈ 2 h 15 m (draft) + ≈ 4 h 10 m (final) ≈ 6.5 h
Hardware	single NVIDIA RTX 4060 (8 GB), TF32 enabled
Throughput	~175 samples/s (draft stage)

The model was trained fully from scratch – no pretrained UNet weights, no distillation, no LoRA on a bigger model.

---

3. Dataset

• 45,251 square-cropped web images (384 px and 768 px encodings), filtered for resolution and sharpness.
• Captions are machine-generated:
  • ~7 k high-quality detailed single-sentence captions from Qwen2-VL-2B-Instruct
  • the remainder from Florence-2-large (<MORE_DETAILED_CAPTION>) / BLIP fallbacks
• Latents and CLIP embeddings were pre-computed once and cached (fp16) for training.

The dataset reflects general web imagery; the model is intended as a research demo of small-scale from-scratch cascade training.

---

4. Usage

The repo ships two subfolders, draft/ and final/, each a standard diffusers UNet. A minimal inference pipeline:

import torch, torch.nn.functional as F
from diffusers import UNet2DConditionModel, AutoencoderKL
from transformers import CLIPTokenizer, CLIPModel

DEVICE = "cuda"
REPO = "<your-username>/echo-vision-1"

tokenizer = CLIPTokenizer.from_pretrained("openai/clip-vit-base-patch32")
clip   = CLIPModel.from_pretrained("openai/clip-vit-base-patch32").to(DEVICE)
vae    = AutoencoderKL.from_pretrained("stabilityai/sd-vae-ft-ema").to(DEVICE)
draft  = UNet2DConditionModel.from_pretrained(REPO, subfolder="draft").to(DEVICE)
final  = UNet2DConditionModel.from_pretrained(REPO, subfolder="final").to(DEVICE)

@torch.inference_mode()
def embed(text):
    ids = tokenizer(text, padding="max_length", max_length=77,
                    truncation=True, return_tensors="pt").input_ids.to(DEVICE)
    return clip.text_model(ids).last_hidden_state

@torch.inference_mode()
def sample(unet, cond, uncond, lat, steps, guidance, cond_latent=None):
    x = torch.randn(1, 4, lat, lat, device=DEVICE)
    for i in range(steps):
        t = 1.0 - i / steps
        tt = torch.full((1,), t * 1000, device=DEVICE)
        inp = x if cond_latent is None else torch.cat([cond_latent, x], dim=1)
        v_c = unet(inp, tt, encoder_hidden_states=cond).sample
        v_u = unet(inp, tt, encoder_hidden_states=uncond).sample
        x = x - (v_u + guidance * (v_c - v_u)) / steps
    return x

prompt = "a red fox in a neon-lit city at night"
emb, u = embed(prompt), embed("")

z = sample(draft, emb, u, 48, steps=20, guidance=3.0)          # Stage 1
z = sample(final, emb, u, 96, steps=15, guidance=2.0,          # Stage 2
           cond_latent=F.interpolate(z, size=96, mode="bilinear", align_corners=False))

img = (vae.decode(z / 0.18215).sample / 2 + 0.5).clamp(0, 1)   # 768x768 RGB

Recommended inference settings: draft 20 steps @ guidance 3.0, final 12–20 steps @ guidance 2.0.

---

5. Limitations & Biases

Echo-Vision-1 is a small research model and should be judged accordingly:

• Quality is below Stable Diffusion 1.5 on open-domain prompts; the gap to frontier models (Flux, DALL·E 3, Gemini image) is large and expected at this scale.
• No text rendering; struggles with complex scenes, counting, anatomy (hands) and rare concepts.
• Inherits biases of web-scraped imagery and machine-generated captions.
• Outputs may still show artifacts on prompts far outside the training distribution.

Its value lies in demonstrating that a coherent coarse-to-fine cascade can be trained from scratch on consumer hardware in a single night.

---

6. Ethical Considerations

Intended for research and personal experimentation. Do not use for generating misleading or harmful content. Training data was collected for personal research only.

---

7. Citation

@misc{echovision1,
  title  = {Echo-Vision-1: A Two-Stage Coarse-to-Fine Rectified-Flow
            Text-to-Image Model Trained from Scratch on a Consumer GPU},
  author = {<Your Name>},
  year   = {2026},
  howpublished = {\url{https://huggingface.co/<your-username>/echo-vision-1}}
}

---

Built with ❤️ on one RTX 4060 – proof that you don't need a datacenter to train your own diffusion model.