---
title: 'OmniPSD: Unified PSD Generation & Decomposition'
url: https://www.emergentmind.com/topics/omnipsd
type: topic
---

# OmniPSD: Unified PSD Generation & Decomposition

OmniPSD is a unified generative and decomposition framework for layered PSD (Photoshop Document) files, which supports both text-to-PSD generation and image-to-PSD decomposition while preserving semantically coherent and hierarchically structured RGBA layers. Built upon the Flux ecosystem and leveraging a DiT-style diffusion transformer in an RGBA-VAE latent space, OmniPSD establishes a new paradigm in editable design synthesis, enabling high-fidelity, transparency-aware layered designs from textual prompts or single input images [2512.09247].

## 1. Problem Formulation: Dual Generative and Decomposition Tasks

OmniPSD addresses two core tasks, both within a shared latent space induced by an RGBA-VAE encoder \(E_\alpha\) and decoder \(D_\alpha\):

- **Text→PSD Generation**: The input is a hierarchical prompt \(\mathbf{c} = \{\text{caption}_\text{full}, \text{caption}_\text{fg}, \text{caption}_\text{mid}, \text{caption}_\text{bg}\}\). The output is \(L=4\) RGBA layers \(\{I^{(l)}\in\mathbb{R}^{H\times W\times 4}\}_{l=1}^4\), corresponding to full poster, foreground, midground, and background. These are visually arranged as  
  $$
  G = 
  \begin{bmatrix}
  I^{(1)} & I^{(2)} \\
  I^{(3)} & I^{(4)}
  \end{bmatrix} \in \mathbb{R}^{2H\times 2W\times 4}
  $$
  The grid \(G\) is encoded via \(E_\alpha\) into latent tokens \(\mathbf{z}_0\), and a diffusion/flow transformer is trained to map noise to this latent grid, i.e., learning \(p_\theta(\mathbf{z}_0|\mathbf{c})\).

- **Image→PSD Decomposition**: The input is a single flattened RGBA poster \(\mathbf{I}_0\). The output is a stack of layers \(\{I_\text{text}, I_{\text{fg}^{(1)}}, ..., I_{\text{fg}^{(K)}}, I_\text{bg}\}\). Each flattened input and target layer is encoded: \(\mathbf{z}_0 = E_\alpha(\mathbf{I}_0)\), \(\mathbf{z}_1 = E_\alpha(I_y)\). A flow-matching vector field
  $$
  \frac{d\mathbf{z}_t}{dt} = v_\theta(\mathbf{z}_t, t|\mathbf{z}_0), \quad \mathbf{z}_t = (1-t)\mathbf{z}_0 + t\mathbf{z}_1, \; t\in[0,1]
  $$
  is learned for layerwise extraction and erasure via iterative LoRA-adapters.

This dual problem formulation enables both end-to-end design synthesis from descriptions and full layer decomposition for direct PSD editability.

## 2. Diffusion-Transformer Architecture and Attention Mechanisms

OmniPSD utilizes a two-stage DiT-style transformer operating over RGBA-VAE latent tokens:

- **Token Embedding**: The RGBA-VAE encodes each image into a spatial grid \(\mathbf{Z}\in \mathbb{R}^{N\times D}\) of tokens.

- **Multi-head Self-attention**: Given token matrix \(\mathbf{X}\):
  $$
  Q = XW_Q,\quad K = XW_K,\quad V = XW_V,\quad \mathrm{Attn}(Q,K,V) = \mathrm{softmax}\!\left( \frac{QK^\top}{\sqrt{d}} \right) V
  $$
  with outputs concatenated and projected.

- **Cross-attention**: For conditional generation, queries are drawn from the diffusion stream, and keys/values from condition tokens (text or flattened images):
  $$
  \mathrm{CrossAttn}(Q_{\rm diff}, K_{\rm cond}, V_{\rm cond}) = \mathrm{softmax}\!\left( \frac{Q_{\rm diff}K_{\rm cond}^\top}{\sqrt{d}} \right) V_{\rm cond}
  $$

- **Spatial In-context Reasoning**: By arranging RGBA layers into a \(2\times2\) grid, self-attention connects across all layers (e.g., foreground↔background, foreground↔full poster), increasing compositional capacity without specialized modules.

This configuration allows the same transformer backbone to serve both generation and decomposition in a unified architecture.

## 3. Diffusion and Flow-matching Objectives

OmniPSD employs distinct objectives depending on task:

- **DDPM-style Diffusion (Text→PSD)**:
  - *Forward noising*:
    $$
    q(\mathbf{z}_t | \mathbf{z}_0) = \mathcal{N}(\mathbf{z}_t; \alpha_t \mathbf{z}_0, (1-\alpha_t^2)\mathbf{I})
    $$
    with $\alpha_t = \sqrt{\bar{\alpha}_t}$.
  - *Reverse denoising* by the conditional DiT:
    $$
    p_\theta(\mathbf{z}_{t-1} | \mathbf{z}_t, \mathbf{c}) = \mathcal{N}(\mathbf{z}_{t-1}; \mu_\theta(\mathbf{z}_t, t, \mathbf{c}), \Sigma_\theta(\mathbf{z}_t, t, \mathbf{c}))
    $$
  - *Loss*: Simplified \(\varepsilon\) parameterization:
    $$
    \mathcal{L}_\mathrm{simple} = \mathbb{E}_{t,\,\mathbf{z}_0,\,\varepsilon}\; \|\varepsilon - \varepsilon_\theta(\alpha_t\mathbf{z}_0+\sqrt{1-\alpha_t^2}\,\varepsilon,\,t,\,\mathbf{c})\|^2
    $$

- **ODE Flow-matching (Image→PSD)**:
  - *Continuous trajectory*: $\mathbf{z}_t = (1-t)\mathbf{z}_0 + t\mathbf{z}_1$
  - *Vector field* $v_\theta$ is learned to match:
    $$
    \mathcal{L}_\mathrm{flow} = \mathbb{E}_{t\sim U(0,1),\,\mathbf{z}_0,\,\mathbf{z}_1}\; \|v_\theta(\mathbf{z}_t, t | \mathbf{z}_0) - (\mathbf{z}_1 - \mathbf{z}_0)\|^2
    $$
  - This flow-matching approach is deterministic, enabling faster inference in decomposition.

The framework incorporates LoRA-adapters for task specialization in the iterative Image→PSD extraction/erasure pipeline.

## 4. RGBA-VAE Representation: Transparency Preservation

The RGBA-VAE, an “AlphaVAE” extension, is tailored to editable design with four-channel (RGBA) support:

- **Encoder** $q_\phi(z_{RGB}, z_A|I)$:
  $$
  z_{RGB} \sim \mathcal{N}(\mu_{RGB}(I), \Sigma_{RGB}(I)), \quad z_{A} \sim \mathcal{N}(\mu_A(I), \Sigma_A(I))
  $$
- **Decoder** $p_\phi(\hat I|z_{RGB}, z_A)$: decodes latent representations of color and alpha channels to reconstruct the original image.
- **Loss Function**:
  $$
  \mathcal{L}_{\text{VAE}} = 
    \lambda_{\rm pix}\, \mathbb{E}\|I - \hat I\|_1 + 
    \lambda_{\rm patch}\, \mathbb{E}\|\phi(I) - \phi(\hat I)\|_1 +
    \lambda_{\rm perc}\, \mathbb{E}\|\psi(I) - \psi(\hat I)\|^2 +
    \lambda_{\rm KL} \left( \mathrm{KL}(q(z_{RGB})\|p) + \mathrm{KL}(q(z_A)\|p) \right)
  $$
  with $\phi$, $\psi$ representing patch and perceptual feature extractors; $p$ an isotropic Gaussian prior.

Retraining on layered-poster data yields high PSNR/SSIM and low LPIPS, outperforming alternative Alpha-VAEs (MSE=$9.82\times10^{-4}$, PSNR=32.5 dB, SSIM=0.945) [2512.09247].

## 5. In-context Grid and Iterative Editing Algorithms

OmniPSD’s in-context learning leverages a \(2\times2\) grid structure for compositional reasoning during both tasks:

- **Text→PSD, Single-pass:**
  1. Form a \(2\times2\) grid \(G\) from the four target RGBA layers.
  2. Encode \(G\) with RGBA-VAE; apply diffusion or flow-matching process in latent space.
  3. Decode the recovered latent tokens to obtain all four RGBA output layers.

- **Image→PSD, Multi-step Iterative Extraction (pseudocode excerpt):**
  ```
  Input: flattened RGBA poster I0
  z0 = Eα(I0)
  For k=1 to K:
    1) Foreground extraction:
       z_fg = FM_extract( z0 | mask m_k )
       If not text-layer:
         I_fg^(k) = Dα(z_fg)
       Else:
         I_text = OCR→font-recover→vector
    2) Foreground erasure / background restoration:
       z_bg = FM_erase( z0 | mask m_k )
       I_bg^(k) = Dα(z_bg)
       z0 ← z_bg
  Return layers {I_text, {I_fg^(k)}, I_bg^(K)} as PSD stack
  ```
  Here, each $FM_{\rm extract,erase}$ is a LoRA adapter trained via $\mathcal{L}_\mathrm{flow}$ on Flux-Kontext, ensuring reversible foreground/background separation.

This grid-based approach enables the transformer to reason over layer compositionality and semantic relationships directly within the attention mechanism.

## 6. Dataset and Empirical Evaluation

- **Layered Poster Dataset**: Approximately 200k professionally designed PSDs, each split into RGBA text, foreground, and background layers. For evaluation, 500 hierarchical text prompts and 500 flattened test images are utilized.

- **Metrics**:
  
  | Task         | Metric           | OmniPSD  | GPT-Image-1 | LayerDiffuse |
  |--------------|------------------|----------|-------------|--------------|
  | Text→PSD     | FID              | 30.43    | 53.21       | 89.35        |
  | Text→PSD     | CLIP-Score (%)   | 37.64    | 35.59       | 24.78        |
  | Text→PSD     | GPT-4 Score      | 0.90     | 0.84        | 0.66         |
  | Image→PSD    | MSE              | 1.14e-3  | 2.48e-2     | N/A          |
  | Image→PSD    | PSNR (dB)        | 24.0     | N/A         | N/A          |
  | Image→PSD    | SSIM             | 0.952    | N/A         | N/A          |
  | Image→PSD    | GPT-4 Score      | 0.92     | 0.86        | 0.84         |

- **Ablation Studies**: 
  - RGBA-VAE: MSE=$9.82\times10^{-4}$, PSNR=32.5 dB, SSIM=0.945 vs. Alpha-VAE: $4.15\times10^{-3}$/26.9dB/0.739.
  - Hierarchical prompts: Removal increases Text→PSD FID to 38.56 and reduces the GPT-4 score to 0.78.

User studies (18 participants) confirm the semantic coherence and practical editability of OmniPSD-generated and -decomposed layers. This suggests that transparency-aware latent representations and in-context transformer architectures are effective for editable design generation and decomposition at scale [2512.09247].

Source: https://www.emergentmind.com/topics/omnipsd