PaperlayerView exact revision on arXiv

Source-linked AI summary

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

Kailai Feng, Yuxiang Wei, Bo Chen, Yang Pan, Hu Ye, Songwei Liu, Chenqian Yan, Yuan Gao

arXiv:2603.28713v1cs.CV

TL;DR

Existing diffusion models achieve strong generation and editing but are often too large and latency-intensive for consumer devices, while compact on-device models mainly support T2I and lack editing. DreamLite unifies both tasks in a compact on-device diffusion model, achieving competitive benchmark performance and sub-second 1024 × 1024 inference on Xiaomi 14.

  • Problem

    Existing diffusion models typically use billions of parameters and iterative denoising, while compact on-device models predominantly support T2I generation rather than image editing.

  • Method

    DreamLite uses a compact unified diffusion network with in-context spatial concatenation, task-progressive T2I-to-edit-to-joint training, and step distillation for mobile inference.

  • Results

    GenEval 0.72 and ImgEdit 4.11 outperform prior on-device baselines, while 4-step inference generates or edits 1024 × 1024 images in less than 1s on Xiaomi 14.

  • Takeaways & Limitations

    A single DreamLite model supports both text-to-image generation and text-guided image editing, eliminating the need to deploy two separate models for these tasks.

  • Takeaways & Limitations

    DreamLite has relatively lower GEdit scores and qualitative artifacts in challenging text, structural, and portrait identity scenarios, which the authors associate with information loss from its compact 1.2M VAE.

Abstract

from arXiv · show

Diffusion models have made significant progress in both text-to-image (T2I) generation and text-guided image editing. However, these models are typically built with billions of parameters, leading to high latency and increased deployment challenges. While on-device diffusion models improve efficiency, they largely focus on T2I generation and lack support for image editing. In this paper, we propose DreamLite, a compact unified on-device diffusion model (0.39B) that supports both T2I generation and text-guided image editing within a single network. DreamLite is built on a pruned mobile U-Net backbone and unifies conditioning through in-context spatial concatenation in the latent space. It concatenates images horizontally as input, using a (target | blank) configuration for generation tasks and (target | source) for editing tasks. To stabilize the training of this compact model, we introduce a task-progressive joint pretraining strategy that sequentially targets T2I, editing, and joint tasks. After high-quality SFT and reinforcement learning, DreamLite achieves GenEval (0.72) for image generation and ImgEdit (4.11) for image editing, outperforming existing on-device models and remaining competitive with several server-side models. By employing step distillation, we further reduce denoising processing to just 4 steps, enabling our DreamLite could generate or edit a 1024 x 1024 image in less than 1s on a Xiaomi 14 smartphone. To the best of our knowledge, DreamLite is the first unified on-device diffusion model that supports both image generation and image editing.

1 Introduction

DreamLite addresses the lack of unified on-device models for text-to-image generation and image editing with a compact single network. Its in-context conditioning and progressive training support competitive quality alongside subsecond mobile inference.

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

2 Related Work

Related work improves diffusion efficiency through compact architectures, optimized attention, compression, alignment, and few-step sampling. DreamLite extends this efficiency focus to unified on-device generation and editing under strict memory and latency constraints.

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

3 Method

The DreamLite training pipeline covers architecture design, task-progressive joint pretraining, post-training, and few-step distillation.

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

3.1 Model Architecture

DreamLite combines a compact UNet, lightweight VAE, text encoder, and in-context conditioning to support unified generation and editing efficiently. Its pruned backbone reduces the 2.5B baseline to 389M parameters while retaining generative performance.

  • 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.
  • 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.
  • 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.
  • 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.
  • 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.
  • 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].

3.2 Task-progressive Joint Pretraining

DreamLite uses task-progressive pretraining to stabilize unified generation and editing in a compact model. The process progresses from T2I pretraining to editing alignment and joint training, with region-aware weighting for local edits.

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

3.3 Post Training

DreamLite uses supervised fine-tuning and reinforcement learning after progressive pretraining to improve behavior, perceptual quality, and instruction following. The post-training data and rewards are tailored to high-quality generation and editing.

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

3.4 Step Distillation

DreamLite applies distribution matching distillation to reduce the multi-step sampling burden for mobile generation and editing. The distilled model produces images in four sampling steps without classifier-free guidance.

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

4 Experiments

DreamLite is evaluated on datasets, generation and editing benchmarks, ablations, and mobile inference efficiency. It combines compact unified modeling with competitive quality and substantially reduced sampling cost.

  • 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.
  • 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.
  • Quantitative Results: 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.
  • 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.
  • Ablation Study: 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.

5 Limitations

DreamLite remains constrained by components and tasks that limit fully optimized deployment and perceptual fidelity. The authors identify text-encoder latency, compact-VAE information loss, specialized-task artifacts, and incomplete alignment as future-work targets.

  • 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.
  • 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.
  • Reconstruction Fidelity and Specialized Tasks: 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.
  • 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.

6 Conclusion

DreamLite unifies text-to-image generation and text-guided image editing in one compact on-device diffusion network. Its progressive training and step distillation support competitive performance and sub-second mobile inference.

  • 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.
  • 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.
  • 4 denoising steps enable efficient mobile inference after step distillation.Sampling is compressed for deployment on mobile devices.
  • Less than 1 second is required to generate or edit a 1024 × 1024 image on a Xiaomi 14 smartphone.
Loading 2603.28713v1…