UniMAGE: Unified Imaginative Audio-Video Generation

• UniMAGE is a unified model that integrates script drafting and key-shot design by routing text and image tokens through specialized expert transformers.
• It employs a 'first interleaving, then disentangling' training paradigm to separately optimize narrative reasoning and visual generation for better multimodal outcomes.
• The model achieves state-of-the-art performance with measurable improvements, including enhanced character identity similarity and prompt adherence compared to previous methods.

UniMAGE is a unified director model for imaginative audio-video generation, integrating the traditionally disjoint processes of script drafting and key-shot design into a single framework. By leveraging the Mixture-of-Transformers (MoT) architecture, UniMAGE routes text and image tokens to specialized “expert” transformer branches, bridging user prompts to long-context, multi-shot film scripts and visually consistent keyframe images. The model introduces a “first interleaving, then disentangling” training paradigm and achieves state-of-the-art performance on open-source benchmarks for narrative coherence and visual quality (Zhang et al., 29 Dec 2025).

1. Architecture and Token Routing

UniMAGE is constructed atop the Mixture-of-Transformers (MoT) backbone derived from Bagel, employing $E = 2$ expert sub-transformers per layer:

• Understanding Expert: Text-oriented, responsible for script reasoning.
• Generation Expert: Image-oriented, responsible for keyframe synthesis.

A lightweight router computes the gating distribution $r(x)$ for each token (or block), modulating the expert contributions via

$$\mathrm{MoT}(x) = \sum_{i=1}^{E} r_i(x)\,\mathrm{Expert}_i(x), \quad \sum_{i=1}^E r_i(x)=1,$$

where $r_i(x)$ is the scalar gating probability and $\mathrm{Expert}_i(x)$ is the output from expert $i$.

All token types—text tokens $y_t$ (shared BPE vocabulary), ViT tokens $v_j$ (reference frames, frozen SigLIP2 ViT), and VAE tokens $z_k$ (latents from FLUX-based VAE)—are projected to a shared hidden size $d$, mixed via multimodal self-attention, and routed accordingly.

2. Script and Data Representation

Scripts are linearized into sequences encapsulating user prompts, character/environment/frame identities, frame and video text annotations, and interleaved keyframe latents. A typical sequence is

$r(x)$0

where $r(x)$1 is the prompt, $r(x)$2 are text annotations, and $r(x)$3 are VAE latents. Embeddings for all tokens are unified for multimodal processing.

In-Context ID Prompting is implemented by interspersing identity tokens among image tokens, facilitating character/scene consistency throughout the narrative.

3. Training Paradigm: Interleaving then Disentangling

The UniMAGE training regime is organized into two sequential stages:

3.1 Interleaved Concept Learning (ICL)

• Joint training over multi-shot scripts $r(x)$4 (where $r(x)$5 is text, $r(x)$6 is keyframe latents).
• Unified autoregressive objective: $r(x)$7 where $r(x)$8 denotes the position preceding the $r(x)$9th image block.
• Both $$\mathrm{MoT}(x) = \sum_{i=1}^{E} r_i(x)\,\mathrm{Expert}_i(x), \quad \sum_{i=1}^E r_i(x)=1,$$0 (understanding) and $$\mathrm{MoT}(x) = \sum_{i=1}^{E} r_i(x)\,\mathrm{Expert}_i(x), \quad \sum_{i=1}^E r_i(x)=1,$$1 (generation) receive gradient updates.

3.2 Disentangled Expert Learning (DEL)

• Script-only (pure text) samples update $$\mathrm{MoT}(x) = \sum_{i=1}^{E} r_i(x)\,\mathrm{Expert}_i(x), \quad \sum_{i=1}^E r_i(x)=1,$$2; text+image pairs update $$\mathrm{MoT}(x) = \sum_{i=1}^{E} r_i(x)\,\mathrm{Expert}_i(x), \quad \sum_{i=1}^E r_i(x)=1,$$3 with stop-gradient on $$\mathrm{MoT}(x) = \sum_{i=1}^{E} r_i(x)\,\mathrm{Expert}_i(x), \quad \sum_{i=1}^E r_i(x)=1,$$4.
• Objectives:
  • $$\mathrm{MoT}(x) = \sum_{i=1}^{E} r_i(x)\,\mathrm{Expert}_i(x), \quad \sum_{i=1}^E r_i(x)=1,$$5
  • $$\mathrm{MoT}(x) = \sum_{i=1}^{E} r_i(x)\,\mathrm{Expert}_i(x), \quad \sum_{i=1}^E r_i(x)=1,$$6 (MSE over drift fields)
  • Joint optimization:

  $$\mathrm{MoT}(x) = \sum_{i=1}^{E} r_i(x)\,\mathrm{Expert}_i(x), \quad \sum_{i=1}^E r_i(x)=1,$$7

  with $$\mathrm{MoT}(x) = \sum_{i=1}^{E} r_i(x)\,\mathrm{Expert}_i(x), \quad \sum_{i=1}^E r_i(x)=1,$$8.

• Pre-Context Script Splitting augments text-only phase: scripts are randomly partitioned, and the understanding expert learns narrative continuation from prefixed context.

4. Optimization and Losses

Optimization employs AdamW with learning rate $$\mathrm{MoT}(x) = \sum_{i=1}^{E} r_i(x)\,\mathrm{Expert}_i(x), \quad \sum_{i=1}^E r_i(x)=1,$$9 and a total of $r_i(x)$0 training steps. Primary losses include:

• Next Token Prediction:

$r_i(x)$1

• Rectified Flow for Images:

$r_i(x)$2

• Expert-Load Balancing (Optional):

$r_i(x)$3

The total stage-wise loss adds balancing terms as needed.

5. Inference Workflow

The inference pipeline follows these steps:

1. Given user prompt $r_i(x)$4, the Understanding Expert ($r_i(x)$5) generates the full script autoregressively: $r_i(x)$6.
2. Optionally, an extension or continuation is requested via designated tokens.
3. Script $r_i(x)$7 is segmented into $r_i(x)$8 shots.
4. For each shot $r_i(x)$9, the Generation Expert ($\mathrm{Expert}_i(x)$0) samples VAE keyframe latents $\mathrm{Expert}_i(x)$1, conditioned on corresponding text.
5. Latents $\mathrm{Expert}_i(x)$2 are decoded into images.
6. Resulting script and images $\mathrm{Expert}_i(x)$3 are forwarded to downstream audio-video generators (e.g., Veo3), along with extracted dialogue and sound tokens.

High-level pseudocode formally describes both training and inference routines, specifying dataset partitions and parameter initialization.

6. Datasets and Evaluation Metrics

Three principal datasets underpin the training stages:

• ICL: 450K multi-shot text–image scripts.
• DEL (text expert): 250K pure text scripts.
• DEL (image expert): 250K single-shot text–image pairs.

UniMAGE is evaluated on ViStoryBench using metrics: | Metric | Description | Reported UniMAGE SOTA | |-----------------------------|--------------------------------------------------|-----------------------| | Style Similarity (CSD) | Consistency of visual style | - | | Character ID Similarity (CIDS) | Character identity preservation | 59.2 (vs. 57.0) | | Prompt Adherence (Alignment)| Fidelity to user instructions | 80.8 (vs. 62.5) | | OCCM | On-stage character count matching | 88.07 (vs. 87.0) | | Image Quality (Inception) | Standard image quality measure | - | | Aesthetics | Human-rated visual appeal | ≈4.55 (vs. 5.76) |

A human study with 50 participants ranks UniMAGE highest for narrative logic (GSB=0.72), character consistency, and overall quality (Zhang et al., 29 Dec 2025).

7. Significance and Implementation Context

UniMAGE demonstrates the technical feasibility of unifying imaginative reasoning and visual generation in a single scalable framework, substantially improving character consistency and narrative logic in script-to-video pipelines. The architecture, grounded in expertly routed multimodal attention with explicit expert specializations, enables direct re-implementation using Bagel/MoT codebases, provided access to the well-described datasets and hyperparameters. This suggests broad applicability for automated film production and creative multimedia authoring, especially in contexts requiring long-context story coherence and complex visual composition (Zhang et al., 29 Dec 2025).