DreamLite: A Lightweight On-Device Unified Model for Image Generation and Editing

Motivation

• Existing lightweight on-device diffusion models predominantly support T2I generation and lack image editing.Deploying separate models also increases system complexity and resource consumption.Source: §1 · p. 3

Approach

• DreamLite unifies generation and editing in one compact diffusion network using spatially concatenated target and condition images.Generation uses a blank condition, whereas editing uses the source image.Source: §1 · p. 3

Training

• Task-progressive pretraining sequences T2I, editing, and unified joint training to stabilize optimization in the compact model.The intermediate editing stage aligns visual conditioning with the generative latent space before joint optimization.Source: §1 · p. 3

Results

• GenEval (0.72), DPG (85.8), ImgEdit (4.11), and GEdit (6.88) demonstrate competitive generation and editing performance.DreamLite outperforms specialized lightweight baselines while remaining competitive with larger unified models.Source: §1 · p. 3

Deployment

• 4 denoising steps enable DreamLite to generate or edit 1024 × 1024 images in less than 1 second on Xiaomi 14.This deployment result follows sampling-process compression with DMD2.Source: §1 · p. 3

Contribution

• DreamLite is presented as the first unified on-device diffusion model supporting both generation and editing within one network.The stated contribution eliminates the need to deploy two separate models.Source: §1 · p. 3

Unified Models

• Large unified image models support generation and instruction-based editing, but DreamLite targets the same capabilities for on-device deployment.The paper positions its compact model against larger or cloud-oriented unified systems.Source: §2.1 · p. 4

Efficient Diffusion

• Efficiency research reduces diffusion cost through architectural optimization, attention mechanisms, and training or inference recipes.Examples include key-value compression and linear attention for high-resolution generation.Source: §2.2 · p. 4

On-Device Deployment

• On-device systems use quantization, pruning, and knowledge distillation to reduce model size and latency.Prior compact U-Net systems demonstrated mobile generation, including 1024-pixel output in SnapGen.Source: §2.3 · p. 4

Post-Training Alignment

• Post-training alignment uses reward models and reinforcement-learning variants to improve aesthetics, prompt faithfulness, and instruction following.Editing-specific rewards and GRPO variants are among the approaches reviewed.Source: §2.4 · pp. 4–5

Few-Step Sampling

• Few-step sampling methods distill multi-step diffusion behavior into compact samplers for interactive and on-device applications.The paper employs DMD2 to compress its sampling process.Source: §2.5 · p. 5

Training Pipeline

• The pipeline first describes the model architecture, then presents task-progressive joint pretraining, post-training, and few-step distillation.These stages organize DreamLite’s training and deployment procedure.Source: §3 · p. 5

Architecture

• DreamLite uses a UNet backbone, VAE, and text encoder with in-context conditioning at the input level.The architecture is designed to unify image generation and editing.Source: §3.1 · p. 5

VAE

• The lightweight VAE maps images to four-channel latents with an 8 × 8 downsampling factor for efficient training and inference.The implementation uses TinyVAE for image tokenization.Source: §3.1 · pp. 5–6

Compact UNet

• The UNet is made shallower and thinner by reducing transformer blocks from [0, 2, 10] to [0, 2, 4] and channels from [320, 640, 1280] to [256, 512, 896].Additional optimizations include removing high-resolution self-attention, separable convolutions, reduced feed-forward expansion, and single-KV-head MQA.Source: §3.1 · p. 6

Efficiency

• 389M parameters compress the 2.5B baseline while significantly reducing FLOPs and preserving generative performance.The architectural evolution is evaluated using parameter counts and GFLOPs of the U-Net.Source: §3.1 · p. 6, Figure 3 · p. 6

Text Encoder

• Qwen3-VL-2B provides text conditioning and multimodal instruction comprehension for semantic alignment.The text encoder processes complex user instructions as part of the conditioning stack.Source: §3.1 · p. 6

In-Context Conditioning

• For generation, the target latent is paired with a blank latent; for editing, it is paired with the source latent.The concatenated latent is fed directly into the UNet, with task tokens distinguishing [Generate] and [Edit].Source: §3.1 · pp. 6–7

Progressive training

• T2I → Edit → Unified Joint Training progressively aligns generation and editing before unified optimization.The intermediate editing stage aligns visual conditioning with the generative latent space.Source: §3.2 · p. 7, §3.2.2 · p. 7, §3.2.3 · p. 8

Text-to-image pretraining

• The T2I stage trains a flow-matching diffusion model across 256 × 256, 512 × 512, and 1024 × 1024 resolutions.Logit-normal noise sampling and dynamic time shifting are also used to improve convergence and establish a generative prior.Source: §3.2.1 · p. 7

Edit pretraining

• Editing pretraining activates in-context conditioning and aligns visual conditions with the pretrained generative latent space.It uses paired text-guided editing data and a flow-matching loss.Source: §3.2.2 · p. 7

Region-aware weighting

• Local-edit training weights changed regions more heavily using x = Atotal/Aedit and w(x) = log2(x) + 1.Global editing and style transfer retain uniform weighting to preserve global distribution alignment.Source: §3.2.2 · p. 8

Unified joint training

• Explicit [Generate] and [Edit] task tokens reduce ambiguity during unified joint training on mixed T2I and editing data.The joint stage consolidates the generative prior and editing instruction-following capabilities into one parameter set.Source: §3.2.3 · p. 8

Post-training strategy

• SFT and RL form a two-stage post-training strategy for stabilizing and improving DreamLite after joint pretraining.The strategy addresses behavior instability caused by high-variance pretraining data.Source: §3.3 · p. 8

Supervised fine-tuning

• Approximately 0.5M curated samples improve visual quality and caption diversity during supervised fine-tuning.The data are selected for high visual quality and caption diversity.Source: §3.3.1 · p. 8

Reinforcement learning

• Task-specific reward models guide RLHF: HPSv3 for T2I generation and EditReward for image editing.The ReFL framework uses scalar reward signals to guide the denoising trajectory, with b = 11 for generation and b = 2.5 for editing.Source: §3.3.2 · p. 8

Outcome

• Post-training consistently enhances perceptual quality and instruction following, enabling DreamLite to outperform prior mobile diffusion baselines.The stated outcome covers both capabilities optimized by the task-specific rewards.Source: §3.3.2 · p. 8

Sampling reduction

• 4 denoising steps replace the tens of steps required for high-quality mobile generation and editing.Distribution Matching Distillation reduces sampling overhead without compromising visual quality.Source: §3.4 · p. 9

Distillation objective

• DMD distills a multi-step diffusion teacher into a few-step generator by minimizing approximate KL divergence between real and generated distributions.The loss matches the generator output distribution to the real-image distribution at diffusion timesteps.Source: §3.4 · p. 9

Distilled model

• The distilled model supports high-quality generation and editing in four sampling steps without CFG.GAN loss is additionally used to enhance diversity and realism.Source: §3.4 · p. 9

Dataset Information

• DreamLite is trained on approximately 20M text-to-image pairs and 1.7M image-editing samples, with roughly balanced T2I and editing data during joint training.The editing corpus covers understanding, local, global, view, and style edits.Source: §4.1 · p. 9, §4.2 · p. 10

Quantitative Results

• DreamLite is evaluated across GenEval, DPG, ImgEdit, and GEdit-EN, using 1024 × 1024 outputs for image-generation comparisons.The experiments compare unified, lightweight, and on-device model groups.Source: §4.3 · p. 10, Table 2 · p. 10
• DreamLite achieves SOTA among lightweight models on ImgEdit and supports on-device editing at a model scale significantly under 0.5B parameters.The paper describes this as among the first successful on-device image-editing systems at this scale.Source: §4.3 · p. 11, Table 4 · p. 11

Ablation Study

• The in-context mechanism improves editing from 3.67 to 3.88 and unified-training generation from 0.61 to 0.71 compared with channel concatenation.It represents generation as in-context referencing with a blank condition image.Source: §4.5 · p. 15
• Task-progressive pretraining avoids the losses of direct joint training, reaching GenEval 0.71 and ImgEdit 3.91 versus 0.65 and 3.14.The progressive sequence is T2I → Edit → Unified.Source: §4.5 · p. 15

Text Encoder Scale

• The 0.39B U-Net still relies on a standard 2B-parameter text encoder, which introduces non-negligible latency during on-device deployment.The authors are developing a lightweight text encoder below 1B parameters.Source: §5 · p. 17

Reconstruction Fidelity and Specialized Tasks

• The compact 1.2M VAE is associated with information loss or reconstruction blurriness in complex structural details.Lower GEdit scores and artifacts occur in text generation, text editing, and portrait identity preservation.Source: §5 · p. 17
• DreamLite requires specialized fine-tuning for text and facial generation or editing to improve perceptual quality and instruction following.This is proposed alongside training a slightly larger, high-fidelity VAE.Source: §5 · p. 17

Multi-modal Alignment

• Future work targets more advanced step distillation combined with a reward model to improve alignment with human aesthetic preferences and complex spatial instructions.This limitation concerns multi-modal alignment during post-training.Source: §5 · p. 17

• DreamLite supports text-to-image generation and text-guided image editing within a single network.Its unified conditioning uses in-context spatial concatenation and explicit task tokens.Source: §6 · p. 17
• The task-progressive curriculum trains DreamLite through T2I, editing, and joint stages before high-quality SFT and reinforcement learning.This strategy is designed to stabilize training of the small-capacity unified model.Source: §6 · p. 17
• 4 denoising steps enable efficient mobile inference after step distillation.Sampling is compressed for deployment on mobile devices.Source: §6 · p. 17
• Less than 1 second is required to generate or edit a 1024 × 1024 image on a Xiaomi 14 smartphone.Source: §6 · p. 17