diff --git a/README.md b/README.md index a9d347d..25a84c8 100644 --- a/README.md +++ b/README.md @@ -30,9 +30,10 @@ We believe openness drives innovation, and we invite the research community to i 3. [Performance](#performance) 4. [Quick Start](#quick-start) 5. [Model Summary](#model-summary) -6. [Prompting Guide](#prompting-guide) -7. [Documentation](#documentation) -8. [Citation](#citation) +6. [Ideogram 4 vs Krea 2](#ideogram-4-vs-krea-2) +7. [Prompting Guide](#prompting-guide) +8. [Documentation](#documentation) +9. [Citation](#citation) ## News @@ -261,6 +262,34 @@ For full architecture details, see how the pipeline components fit together, see [docs/pipeline.md](docs/pipeline.md). +## Ideogram 4 vs Krea 2 + +Ideogram 4 and Krea 2 are both modern latent flow / rectified-flow image models +built around Diffusion Transformer backbones and Qwen3-VL-class conditioning. +The practical difference is where each system spends its complexity budget: +Ideogram 4 exposes a concrete design-control inference stack in this repo, while +Krea 2's technical report emphasizes broad aesthetic exploration, style +references, prompt-expander training, and a large multi-stage post-training +pipeline. + +```mermaid +flowchart LR + Shared["Shared modern recipe: DiT + flow + VLM conditioning"] --> I["Ideogram 4"] + Shared --> K["Krea 2"] + I --> IC["Structured JSON captions"] + I --> IL["Bounding-box layout"] + I --> IT["Typography and palette control"] + I --> IA["Asymmetric CFG inference"] + K --> KD["Data curation pipeline"] + K --> KP["Prompt expander"] + K --> KS["Style references"] + K --> KR["SFT / PO / RL stack"] +``` + +See [docs/krea2_comparison.md](docs/krea2_comparison.md) for the full +side-by-side comparison and [docs/diagrams.md](docs/diagrams.md) for the diagram +gallery. + ## Prompting Guide Ideogram 4 is trained exclusively on **structured JSON captions**. While @@ -312,6 +341,8 @@ See [docs/prompting.md](docs/prompting.md) for a full guide. | [docs/inference.md](docs/inference.md) | Sampler presets, parameter reference, resolutions, optimization tips | | [docs/model_architecture.md](docs/model_architecture.md) | Architecture diagram, DiT spec, component details | | [docs/pipeline.md](docs/pipeline.md) | Conceptual pipeline walkthrough — how all components fit together | +| [docs/krea2_comparison.md](docs/krea2_comparison.md) | Side-by-side comparison of Ideogram 4 and Krea 2 | +| [docs/diagrams.md](docs/diagrams.md) | Mermaid diagrams for the repo, inference flow, CFG, and model comparison | | [docs/development.md](docs/development.md) | Dev setup, pre-commit hooks, contributing | | [docs/safety.md](docs/safety.md) | Pre-training, post-training, and inference-time safety mitigations; how to report violations | diff --git a/docs/ARQUITECTURA.md b/docs/ARQUITECTURA.md new file mode 100644 index 0000000..d99e626 --- /dev/null +++ b/docs/ARQUITECTURA.md @@ -0,0 +1,487 @@ +# Ideogram 4 — arquitectura explicada + +Notas propias sobre cómo funciona el modelo por dentro, con los números reales del +`Ideogram4Config` y del autoencoder de este repo. Pensado para alguien que viene de +LLMs/transformers pero no ha tocado difusión. + +--- + +## 1. Visión general + +Ideogram 4 es un modelo **texto → imagen** de tipo **flow matching** (no difusión DDPM +clásica). El pipeline de inferencia es: + +``` +frase corta + │ (magic prompt: una LLM) + ▼ +caption JSON estructurado + │ (Qwen3-VL-8B: text encoder) + ▼ +tokens de condición ─────────────┐ + ▼ +ruido (32,64,64) ──► DiT (×N pasos, flow matching) ──► latente final + │ (VAE decoder) + ▼ + imagen RGB +``` + +Tres piezas entrenadas: + +| Pieza | Qué es | Params | +|---|---|---| +| **VAE** (autoencoder) | comprime imagen ↔ latente | pequeño | +| **DiT** | el modelo de difusión (la estrella) | **~9.3B** | +| **Text encoder** | Qwen3-VL-8B-Instruct (un VLM entero) | ~8B | + +El "magic prompt" es una llamada a otra LLM (hosted o vía OpenRouter), no es parte de +los pesos. + +--- + +## 2. Flow matching (en vez de difusión clásica) + +La difusión clásica (DDPM) entrena la red para **predecir el ruido** `ε` y revierte una +cadena estocástica de cientos de pasos. Ideogram 4 usa **flow matching / rectified flow**: + +- Defines un camino recto entre un latente real `x₀` y ruido `x₁`: + `x_t = (1−t)·x₀ + t·x₁`. +- La velocidad objetivo a lo largo de ese camino es **constante**: `v = x₁ − x₀`. +- La red aprende a predecir esa velocidad `v_θ(x_t, t, condición)`. +- En inferencia resuelves una **ODE**: integras de `t=1` (ruido) a `t=0` (imagen). + Cada "paso de denoising" es un paso del solver (tipo Euler): `x ← x − Δt · v`. + +Como los caminos son casi rectos, bastan **pocos pasos** (12 en turbo, 48 en calidad), +frente a los 50–1000 de DDPM. + +Parámetros de muestreo (presets en `sampler_configs.py`): +- `num_steps`: pasos del solver (12 / 20 / 48). +- `guidance_schedule`: **CFG** variando por paso. CFG corre el modelo con condición y sin + ella y extrapola: `v = v_uncond + w·(v_cond − v_uncond)`. `w` alto = más obediente. +- `mu` / `std`: el **time-shift** del schedule (a más resolución, más shift). + +--- + +## 3. El VAE: por qué pasamos de 3 a 32 canales + +### Qué hace +El VAE (autoencoder variacional) es un **compresor aprendido** entre píxeles y un espacio +latente más pequeño: + +- **Encoder**: `imagen (3, H, W)` → `latente (32, H/8, W/8)` +- **Decoder**: `latente (32, H/8, W/8)` → `imagen (3, H, W)` + +Config real (`autoencoder.py`): `ch_mult = [1, 2, 4, 4]` ⇒ 3 etapas de *downsample* ⇒ +**8× espacial**. `z_channels = 32`. + +### Por qué 32 canales y no 3 +Es un **trade espacial ↔ canales**. Al bajar la resolución 8× pierdes muchísima +información espacial; para no perder calidad, esa información se "reempaqueta" en más +canales. Ejemplo a 512×512: + +| | shape | nº de valores | +|---|---|---| +| Imagen RGB | (3, 512, 512) | 786.432 | +| Latente VAE | (32, 64, 64) | 131.072 | + +Resultado: **~6× menos valores** en total, pero —lo importante para el DiT— **64× menos +posiciones espaciales** (512² → 64²). Como el coste de la atención es cuadrático en el +número de posiciones, eso lo cambia todo. + +Los 3 canales RGB son una representación "tonta" (rojo/verde/azul por píxel). Los 32 +canales del latente son **features aprendidas**: cada canal codifica patrones útiles +(bordes, texturas, color a distintas frecuencias). Más canales = latente de mayor +fidelidad. Referencia: SD1.5 usaba 4 canales, FLUX usa 16, Ideogram 4 usa 32 (latentes +más ricos → mejor detalle y texto). + +### Por qué "variacional" (VAE y no AE a secas) +El encoder produce una distribución (media + varianza) con una regularización KL ligera +que empuja el latente hacia algo suave y bien-comportado (cercano a gaussiano). Eso hace +que el espacio latente sea un **objetivo agradable de modelar** para el DiT. Un +autoencoder plano podría tener un latente arbitrariamente "picudo", más difícil de +aprender. En la práctica el peso KL es pequeño; los "VAE" modernos para difusión son casi +AEs con regularización suave. + +--- + +## 4. Patchify: cómo se trocea el latente + +El DiT es un transformer y necesita una **secuencia de tokens**, no una rejilla 2D. El +patchify convierte la rejilla latente en tokens, igual que un ViT pero sobre el latente. + +Config: `patch_size = 2`, y por eso `in_channels = z_channels(32) × patch²(4) = 128`. + +Paso a paso, a 512px → latente `(C=32, H=64, W=64)`: + +1. **Trocear en bloques 2×2 no solapados.** La rejilla 64×64 se divide en una rejilla de + `32×32 = 1024` bloques, cada uno de 2×2 posiciones. +2. **Aplanar cada bloque con todos sus canales.** Un bloque tiene `32 canales × 2 × 2 = + 128` valores → se aplana en **un vector de 128**. +3. Quedan **1024 tokens** de dimensión 128. +4. **Proyección lineal** `128 → 4608`: cada token pasa a la dimensión del modelo. + +En notación einops: +```python +# (c, h, w) -> (num_tokens, c*p1*p2) +rearrange(x, "c (h p1) (w p2) -> (h w) (c p1 p2)", p1=2, p2=2) +``` + +Al final de la red se hace lo inverso (**unpatchify**): de `(1024, 128)` se reconstruye la +rejilla `(32, 64, 64)`. + +Número de tokens = `(píxeles / 16)²` (8× del VAE × 2 del patch): + +| Resolución | rejilla latente | tokens | +|---|---|---| +| 512² | 32×32 | 1.024 | +| 768² | 48×48 | 2.304 | +| 1024² | 64×64 | 4.096 | +| 2048² | 128×128 | 16.384 | + +--- + +## 5. ¿Por qué un VAE? ¿No se puede hacer la inversa del flow? + +Pregunta clave, porque mezcla dos "inversas" distintas: + +- **El flow (DiT)** mapea **ruido ↔ latente**, ambos en el *mismo* espacio + `(32,64,64)`. Es (aproximadamente) invertible porque es una ODE determinista. Pero + **nunca toca píxeles**. +- **El VAE** mapea **píxeles ↔ latente**. Su "inversa" es el **decoder** (que es una red + aparte, entrenada para reconstruir; el encode→decode es *lossy*, no identidad exacta). + +Es decir, hay **dos pares inversa-ish, en dominios diferentes**: + +``` +píxeles ⇄ latente (VAE encoder / VAE decoder) +ruido ⇄ latente (flow forward / flow reverse) +``` + +"Invertir el flow" te lleva de ruido a **latente**, no a píxeles. Para llegar a píxeles +necesitas sí o sí el **decoder del VAE**. Son operaciones distintas. + +### ¿Y por qué no correr el flow directamente sobre píxeles (y olvidarse del VAE)? +Se puede (existe la *pixel-space diffusion*), pero es carísimo a alta resolución: + +- A 512px en píxeles, con patch 2, tendrías `512×512/4 = 65.536` tokens, frente a `1.024` + en latente. +- La atención es O(n²): `65.536²` vs `1.024²` ≈ **~4000× más** cómputo de atención. + +Además el VAE descarta detalle de alta frecuencia imperceptible, así el DiT gasta su +capacidad en semántica y composición en vez de en reproducir cada píxel. Por eso el +estándar es **difusión latente** (VAE + DiT). + +--- + +## 6. El DiT por dentro + +Config real (`Ideogram4Config`): + +| Campo | Valor | +|---|---| +| `emb_dim` (dimensión del modelo) | 4608 | +| `num_layers` | 34 | +| `num_heads` | 18 (head_dim = 256) | +| `intermediate_size` (MLP SwiGLU) | 12288 | +| `adanln_dim` (embedding del timestep) | 512 | +| `in_channels` | 128 | +| `rope_theta` | 5.000.000 | +| `mrope_section` | (24, 20, 20) — RoPE multimodal 3D | +| Total params | ~9.3B | + +Params por bloque (×34): atención ≈ 85M (`qkv` 4608→13824, `o` 4608→4608), MLP SwiGLU +≈ 170M (`w1,w3` 4608→12288, `w2` 12288→4608), adaLN ≈ 9.4M → **~264M/bloque**. + +### El bloque (lo que lo hace un *Diffusion* Transformer: adaLN) +A diferencia de un transformer normal, el timestep `t` **no** entra como token. Una Linear +`512 → 4×4608` genera 4 señales de modulación que escalan las normalizaciones y abren/ +cierran los residuales: + +``` +x ──► RMSNorm ──► ×(1+scale_msa) ──► Self-Attention ──► ×gate_msa ──► (+) ──► x' +x' ─► RMSNorm ──► ×(1+scale_mlp) ──► MLP SwiGLU ──► ×gate_mlp ──► (+) ──► x'' + ▲ ▲ + └──────────── modulación desde t (adaLN) ───────────┘ +``` + +Así el mismo peso "sabe" si está en el paso 1 (mucho ruido) o en el 11 (casi imagen). + +### Single-stream (vs MMDiT de SD3/FLUX) +Texto e imagen se concatenan en **una sola secuencia** procesada por **los mismos pesos** +en las 34 capas, sin ramas separadas. Más interacción cross-modal y escala como un LLM. + +### Recorrido de shapes (a 512px) + +| Etapa | Shape | +|---|---| +| Ruido latente | (32, 64, 64) | +| Patchify 2×2 | (1024, 128) | +| Proyección | (1024, 4608) | +| + tokens de condición | (L, 4608) | +| qkv (por bloque) | (L, 13824) → 18×(L,256) | +| atención (por cabeza) | (L, L) | +| MLP SwiGLU | (L,4608)→(L,12288)→(L,4608) | +| Cabeza de salida (velocidad) | (L, 128) | +| Unpatchify | (32, 64, 64) | +| (×12 pasos) → VAE decoder | (3, 512, 512) | + +--- + +## 7. El text encoder: un VLM entero + +En vez de CLIP/T5, usa **Qwen3-VL-8B-Instruct**. Se extraen hidden states de **13 capas** +(0, 3, 6, …, 33, 35 → `llm_features_dim = 4096 × 13 = 53248`), que se proyectan y entran +al DiT como tokens de condición. Esto da comprensión profunda del prompt, multilingüe, y +sobre todo **renderizado de texto** correcto dentro de la imagen. + +--- + +## 8. Magic prompt + +El modelo se entrenó con **captions JSON estructurados**, no con frases planas. El "magic +prompt" es una LLM que expande tu frase a ese JSON. Configs disponibles: `ideogram-4-v1` +(API hosted gratis de Ideogram, lee `IDEOGRAM_API_KEY`), `claude-opus-v1`, +`claude-sonnet-v1` (vía OpenRouter). + +Esquema del JSON: +```jsonc +{ + "high_level_description": "...resumen de la escena...", + "compositional_deconstruction": { + "background": "...descripción del fondo...", + "elements": [ + { "type": "obj", "desc": "...objeto..." }, + { "type": "text", "text": "TEXTO EXACTO", "desc": "...cómo/dónde se renderiza..." } + ] + } +} +``` + +Clave: los elementos `type: "text"` llevan **la cadena literal** a renderizar → por eso +Ideogram escribe texto tan bien. Contrapartida: la LLM **inventa composición** (objetos, +fondo) que tú no pediste. Con `--no-magic-prompt` el modelo recibe tu frase plana (fuera +de su distribución → peor); o puedes escribir tú el JSON a mano para control total. + +--- + +## 9. Notas de ejecución en Mac (Apple Silicon) + +Footprint real de los pesos (fp8), medido en caché: + +| Componente | fp8 | +|---|---| +| transformer (condicional) | 8.7 GB | +| unconditional_transformer | 8.7 GB | +| text_encoder (Qwen3-VL) | 8.2 GB | +| vae | 0.16 GB | +| **total** | **~26 GB** | + +Son **dos DiTs completos** (CFG asimétrico, ver §10.6) + el encoder → **26 GB > 24 GB** +de RAM. En una M-Pro de 24 GB esto **no cabe** ni en fp8 (hace *thrashing* de swap). + +Limitaciones de plataforma: +- **fp8 NO corre en MPS**: PyTorch no soporta el dtype `float8_e4m3fn` en Metal + (`Fp8Linear.forward` revienta en `weight.to(bf16)`). +- La ruta `nf4` (bitsandbytes 4-bit) es **CUDA-only**. + +Solución implementada en [`run_metal.py`](../run_metal.py): ver §11. + +Licencia **Non-Commercial**: solo uso personal/experimental. + +--- + +## 10. Las matemáticas, en detalle + +Notación: `d = 4608` (emb_dim), `d_ff = 12288`, `h = 18` cabezas, `d_h = 256` +(head_dim), `L` = longitud de secuencia, `c` = condición (tokens del encoder). + +### 10.1 Flow matching (rectified flow) + +Dos extremos: un latente real `x₀ ~ p_data` y ruido `x₁ ~ N(0, I)`. Se define el +**camino recto**: + +$$x_t = (1-t)\,x_0 + t\,x_1, \qquad t \in [0,1]$$ + +Su derivada temporal (la **velocidad**) es constante a lo largo del camino: + +$$\frac{dx_t}{dt} = x_1 - x_0 =: u_t$$ + +La red `v_θ` se entrena para regresar esa velocidad (condicionada en `t` y en `c`): + +$$\mathcal{L} = \mathbb{E}_{t,\,x_0,\,x_1}\Big[\;\big\lVert v_\theta(x_t, t, c) - (x_1 - x_0) \big\rVert_2^2 \;\Big]$$ + +En **inferencia** se resuelve la ODE de `t=1` (ruido) a `t=0` (imagen): + +$$\frac{dx}{dt} = v_\theta(x, t, c)$$ + +Con un solver de Euler y paso `Δt` (lo que hace el bucle, `z = z + v·Δt` con `Δt<0`): + +$$x_{t-\Delta t} = x_t - \Delta t \cdot v_\theta(x_t, t, c)$$ + +Como el camino es recto, la trayectoria es casi una línea → bastan pocos pasos (12–48). + +### 10.2 Schedule de `t` (logit-normal + shift por resolución) + +Durante el entrenamiento `t` no se muestrea uniforme sino **logit-normal** (params +`μ = mu`, `σ = std`): + +$$t = \sigma\!\left(\mu + \sigma\,\varepsilon\right), \quad \varepsilon \sim N(0,1), \quad \sigma(a)=\tfrac{1}{1+e^{-a}}$$ + +En inferencia se aplica un **time-shift** `s` que depende de la resolución (más tokens → +más shift), que reescala `t`: + +$$t' = \frac{s\,t}{1 + (s-1)\,t}$$ + +Esto concentra pasos donde más hacen falta a alta resolución. + +### 10.3 VAE (ELBO) + +Encoder probabilístico `q_φ(z|x) = N(μ_φ(x), σ_φ²(x))`, con reparametrización +`z = μ_φ(x) + σ_φ(x) ⊙ ε`, `ε ~ N(0,I)`. Decoder `p_θ(x|z)`. Se maximiza el ELBO: + +$$\log p(x) \;\ge\; \underbrace{\mathbb{E}_{q_\phi}[\log p_\theta(x|z)]}_{\text{reconstrucción}} \;-\; \underbrace{\mathrm{KL}\!\big(q_\phi(z|x)\,\Vert\,N(0,I)\big)}_{\text{regularización}}$$ + +$$\mathrm{KL} = \tfrac{1}{2}\sum_i \big(\mu_i^2 + \sigma_i^2 - \log\sigma_i^2 - 1\big)$$ + +En la práctica la reconstrucción es `L1/L2 + perceptual + GAN`, y el peso de la KL es +pequeño (β≪1) → "VAE" casi-determinista. Antes de difundir, el latente se normaliza +(buffers `latent_shift`, `latent_scale`): `ẑ = (z − shift)/scale`; al decodificar se +invierte: `z = ẑ·scale + shift`. + +### 10.4 Patchify (conteo) + +Latente `(C, H_l, W_l)` con `C=32`, `H_l=W_l=H/8`. Con `patch=p=2`: + +$$\text{tokens} = \frac{H_l}{p}\cdot\frac{W_l}{p}, \qquad \dim_\text{token} = C\cdot p^2 = 32\cdot4 = 128$$ + +$$\text{rearrange: } (C, H_l, W_l) \;\to\; \Big(\tfrac{H_l}{p}\tfrac{W_l}{p},\; C p^2\Big) \;\xrightarrow{\text{Linear } 128\to d}\; (L_\text{img}, d)$$ + +### 10.5 Atención + RoPE 3D + +Por cabeza, atención escalada: + +$$\mathrm{Attn}(Q,K,V) = \mathrm{softmax}\!\Big(\frac{QK^\top}{\sqrt{d_h}}\Big)\,V, \qquad d_h = 256$$ + +`qkv` proyecta `d → 3d` y se reparte en `h=18` cabezas de `d_h=256` (`h·d_h = d`). + +**RoPE** rota pares de dimensiones por un ángulo dependiente de la posición: + +$$\theta_k = \text{pos}\cdot \text{base}^{-2k/d_h}, \qquad \text{base} = 5\times10^{6}$$ + +**mRoPE 3D**: el `head_dim/2 = 128` se parte en secciones `(24, 20, 20)` para las +coordenadas `(temporal, alto, ancho)`. Tokens de imagen usan `pos = (0, h, w)`; tokens +de texto usan `pos = (p, p, p)`. Así la posición 2D del patch entra en la atención. + +Coste: tiempo `O(L²·d)`, memoria de la matriz de scores `O(h·L²)` → cuadrático en `L` +(ver §10.8). + +### 10.6 adaLN y CFG + +**adaLN.** El embedding del timestep `t` (dim `adanln_dim=512`) pasa por una +`Linear(512 → 4d)` que produce 4 señales: `scale_msa, gate_msa, scale_mlp, gate_mlp`. +Cada sub-capa se modula así (RMSNorm tipo *sandwich*, con `gate = tanh(·)` y +`scale ← 1+scale`, según el código): + +$$x' = x + \mathrm{gate_{msa}} \odot \mathrm{Norm_2}\big(\mathrm{Attn}(\mathrm{Norm_1}(x)\odot(1+\mathrm{scale_{msa}}))\big)$$ + +$$x'' = x' + \mathrm{gate_{mlp}} \odot \mathrm{Norm_2}\big(\mathrm{FF}(\mathrm{Norm_1}(x')\odot(1+\mathrm{scale_{mlp}}))\big)$$ + +Con `RMSNorm(x) = \dfrac{x}{\sqrt{\overline{x^2}+\epsilon}}\odot\gamma`. + +**CFG asimétrico.** Se combinan la velocidad condicional y la incondicional. El código +usa la forma convexa (con `gw` = peso por paso del `guidance_schedule`): + +$$v = \mathrm{gw}\cdot v_\text{cond} + (1-\mathrm{gw})\cdot v_\text{uncond} = v_\text{uncond} + \mathrm{gw}\,(v_\text{cond} - v_\text{uncond})$$ + +"Asimétrico" = la rama incondicional es **solo imagen** (sin tokens de texto) y usa un +**segundo DiT** (`unconditional_transformer`) con la condición a cero. Por eso hay dos +modelos en memoria. + +### 10.7 SwiGLU + +$$\mathrm{FF}(x) = W_2\big(\mathrm{SiLU}(W_1 x) \odot (W_3 x)\big), \qquad \mathrm{SiLU}(a)=a\,\sigma(a)$$ + +Con `W₁, W₃ : d→d_ff` y `W₂ : d_ff→d` (`d=4608`, `d_ff=12288`). + +### 10.8 Conteo de parámetros y FLOPs + +**Params por bloque** (`d=4608`, `d_ff=12288`, `d_a=512`): + +$$\underbrace{3d^2 + d^2}_{\text{attn }=4d^2} + \underbrace{3\,d\,d_{ff}}_{\text{MLP}} + \underbrace{d_a\cdot 4d}_{\text{adaLN}} = 84.9\text{M} + 169.9\text{M} + 9.4\text{M} \approx 264\text{M}$$ + +$$\times\,34 \text{ bloques} \approx 8.99\text{B} \;\;(+\text{ embeds/cabeza}) \approx 9.3\text{B}$$ + +**FLOPs por forward** (dominado por los Linear): `≈ 2·N_params·L`. A 512px (`L≈1024`): + +$$\approx 2 \cdot 9.3\times10^{9} \cdot 1024 \approx 1.9\times10^{13}\ \text{FLOP/forward}$$ + +La parte de atención `≈ 4 L^2 d` por capa: a `L=1024` son `~6.6×10¹¹` en total (las 34 +capas) — un orden de magnitud por debajo de los Linear *a esta resolución*. A 2048px +(`L≈16384`) la atención crece `×256` y pasa a dominar. + +**Por imagen**: `num_steps × (1 forward si no-CFG, 2 si CFG)`. Turbo (12) sin CFG ≈ +`12 × 1.9×10¹³ ≈ 2.3×10¹⁴` FLOP (230 TFLOP); con CFG, el doble. + +### 10.9 Cuantización fp8 (e4m3) y la LUT + +Formato `e4m3`: 1 bit signo, 4 exponente (bias 7), 3 mantisa → máx normal `448`. +**Escala por fila** (por canal de salida `i`): + +$$s_i = \frac{\max_j |W_{ij}|}{448}, \qquad Q_{ij} = \mathrm{round_{fp8}}\!\Big(\frac{W_{ij}}{s_i}\Big), \qquad W_{ij} \approx Q_{ij}\cdot s_i$$ + +Como cada `Q_{ij}` es **1 byte**, la descuantización es una **LUT** de 256 entradas +`T[b] = \text{fp8\_to\_float}(b)`: + +$$W \approx T\big[\text{bytes}(Q)\big] \odot s$$ + +Esto es justo lo que hace [`run_metal.py`](../run_metal.py) para correr en MPS: guarda +`Q` como `uint8` (MPS soporta `uint8`, no `float8`) y aplica `T` con un *gather*. + +--- + +## 11. Optimización para Metal (24 GB) — `run_metal.py` + +El problema es doble: (a) MPS no soporta `float8_e4m3fn`; (b) los pesos suman ~26 GB y +no caben en 24 GB (§9). `run_metal.py` ataca ambos: + +**1. Shim fp8 → MPS por LUT.** Se monkeypatchea `Fp8Linear` para guardar el peso como +`uint8` (los bytes crudos del e4m3, que MPS sí almacena) y descuantizar en el forward con +una LUT de 256 entradas (§10.9): +```python +w = lut[self.weight.long()] * self.weight_scale.unsqueeze(1) # uint8 -> bf16 +``` +También se parchea `load_fp8_state_dict` para hacer `tensor.view(torch.uint8).to("mps")` +en lugar de mover el `float8` (que falla). + +**2. Carga por fases liberando memoria.** Nunca se tienen los 26 GB a la vez: +``` +Fase A: cargar encoder (8.2GB) → codificar prompt → LIBERAR encoder +Fase B: cargar DiT (8.7GB) → bucle de denoising +Fase C: cargar VAE (0.16GB) → decode +``` +Pico ≈ tamaño del componente mayor (~8.7 GB) en vez de 26 GB. + +**3. Sin CFG por defecto.** Se salta el segundo DiT (`--cfg` lo reactiva, pero va al +límite de RAM). Sin CFG, `v = v_cond` directamente; con CFG se carga el +`unconditional_transformer` y se aplica §10.6. + +Medidas (M5 Pro, 24 GB, 512px, turbo 12, sin CFG): +- Fase A (encoder, carga + encode en MPS): ~124–131 s. +- Denoising: paso 1 ~16 s (warmup/compilación de kernels MPS), pasos siguientes + **~3 s/paso** en régimen. +- 12 pasos ≈ ~50 s de cómputo. Total extremo-a-extremo ~3 min. +- Comparativa: la ruta CPU (`run_inference.py --device cpu`) no terminaba en 27 min + por *thrashing* de swap (26 GB de pesos en 24 GB de RAM). + +**Calidad: CFG es obligatorio.** Sin CFG (un solo DiT) la imagen colapsa a gris/ruido: +el modelo condicional solo no da una velocidad válida (se entrenó para usarse en la +combinación de §10.6). Con `--cfg --magic` (los 2 DiTs + caption JSON) la imagen sale +correcta, pero los ~17 GB residentes hacen *swap*: +- Con CFG: **~60–95 s/paso** (vs 3 s sin CFG), total **~21 min** para 512px/turbo. +- Conclusión: en 24 GB hay que elegir **velocidad XOR calidad**. La versión buena + (CFG + magic) es correcta pero impracticable por RAM → para uso real, **GPU CUDA** + (ruta `nf4`, entra en 24 GB de VRAM y va en segundos). + +Alternativa no implementada para CFG completo en 24 GB: **carga en streaming** de los +shards (mover tensor a tensor a MPS y liberar) para evitar el pico transitorio +`state_dict (CPU) + modelo (MPS)`, manteniendo los 2 DiTs (~17.4 GB) residentes. diff --git a/docs/diagrams.md b/docs/diagrams.md new file mode 100644 index 0000000..ecc0a76 --- /dev/null +++ b/docs/diagrams.md @@ -0,0 +1,127 @@ +# Diagrams + +This page collects the model and inference diagrams that are useful when +orienting around the Ideogram 4 codebase and comparing it with Krea 2. + +## Repo Map + +```mermaid +flowchart TD + R["ideogram4 repo"] --> Run["run_inference.py / run_metal.py"] + R --> Src["src/ideogram4"] + R --> Docs["docs"] + R --> Assets["assets"] + + Src --> Pipe["pipeline_ideogram4.py"] + Src --> Model["modeling_ideogram4.py"] + Src --> Sched["scheduler.py"] + Src --> AE["autoencoder.py"] + Src --> Safety["safety.py"] + Src --> Magic["magic_prompt.py"] + + Pipe --> Model + Pipe --> Sched + Pipe --> AE + Pipe --> Magic + Pipe --> Safety +``` + +## Ideogram 4 Inference + +```mermaid +sequenceDiagram + participant User + participant Magic as Magic Prompt + participant Qwen as Qwen3-VL + participant DiT as Ideogram4Transformer + participant Sampler as Euler Sampler + participant VAE + + User->>Magic: plain prompt + Magic->>Qwen: structured JSON caption + Qwen->>DiT: concatenated hidden-state taps + Sampler->>DiT: noisy image tokens + timestep + DiT->>Sampler: conditional velocity + Sampler->>DiT: image-only negative branch + DiT->>Sampler: unconditional velocity + Sampler->>Sampler: asymmetric CFG update + Sampler->>VAE: final latents + VAE->>User: image +``` + +## Transformer Block + +```mermaid +flowchart TD + X["Token sequence"] --> N1["RMSNorm"] + T["Timestep embedding"] --> M["AdaLN projection"] + M --> SA["Scale attention branch"] + N1 --> SA + SA --> A["Attention: QK-RMSNorm + MRoPE"] + A --> G1["Tanh gate + residual"] + X --> G1 + G1 --> N2["RMSNorm"] + M --> SM["Scale MLP branch"] + N2 --> SM + SM --> FF["SwiGLU MLP"] + FF --> G2["Tanh gate + residual"] + G1 --> G2 +``` + +## Asymmetric CFG + +```mermaid +flowchart LR + Z["Current latents"] --> Pos["Conditional transformer: text + image tokens"] + Z --> Neg["Unconditional transformer: image tokens only"] + Text["Qwen3-VL features"] --> Pos + Zero["Zero text features"] --> Neg + Pos --> PV["positive velocity"] + Neg --> NV["negative velocity"] + PV --> Mix["v = guidance * pos + (1 - guidance) * neg"] + NV --> Mix + Mix --> Step["Euler flow step"] +``` + +## Ideogram 4 vs Krea 2 + +```mermaid +flowchart TD + Root["Modern image model stack"] + + Root --> I["Ideogram 4"] + I --> ICode["Open inference code"] + I --> IParams["9.3B released variants"] + I --> IJson["JSON captions"] + I --> IBox["Bounding boxes"] + I --> IColor["Color palettes"] + I --> ICFG["Asymmetric CFG"] + I --> I2K["Native 2K focus"] + + Root --> K["Krea 2"] + K --> KReport["Technical report focus"] + K --> KData["Data curation"] + K --> KMid["Midtraining and SFT"] + K --> KPref["Preference optimization"] + K --> KRL["Multi-reward RL"] + K --> KPrompt["Prompt expander"] + K --> KStyle["Style references"] +``` + +## Krea 2 Training Stack + +```mermaid +flowchart TD + Raw["Large real-image corpus"] --> Filter["Dedup, filtering, OCR, captioning"] + Filter --> Pre["Progressive pretraining: 256, 512, 1024"] + Pre --> Mid["Midtraining for domain coverage"] + Mid --> SFT["High-aesthetic SFT"] + SFT --> Merge["Model merging"] + Merge --> PO["Preference optimization"] + PO --> RL["Multi-reward RL"] + RL --> Distill["Optional timestep distillation"] + Distill --> Serve["Generation stack"] + + Expand["Prompt expander"] --> Serve + Style["Style-reference system"] --> Serve +``` diff --git a/docs/krea2_comparison.md b/docs/krea2_comparison.md new file mode 100644 index 0000000..5e4bbed --- /dev/null +++ b/docs/krea2_comparison.md @@ -0,0 +1,188 @@ +# Ideogram 4 vs Krea 2 + +This document compares the open Ideogram 4 implementation in this repository +with the architecture and training stack described in the public +[Krea 2 Technical Report](https://www.krea.ai/blog/krea-2-technical-report). + +The short version: both models sit in the same modern family of latent +flow-matching / rectified-flow Diffusion Transformers, and both lean on +Qwen3-VL-style visual-language representations. Ideogram 4 is more explicitly +productized around structured layout control, text rendering, asymmetric CFG, +and native 2K inference in the released code. Krea 2 is presented as a broader +foundation-model series optimized for aesthetic diversity, creative exploration, +prompt expansion, style-reference control, and a heavy multi-stage post-training +stack. + +## At A Glance + +| Axis | Ideogram 4 in this repo | Krea 2 technical report | +| --- | --- | --- | +| Release shape | Open inference code and gated open weights | Weights and inference are described as released under a permissive license | +| Main goal | Design-first generation with strong typography, layout, palettes, and JSON control | Creative exploration across broad aesthetics with text and image-based control | +| Backbone | Fully single-stream DiT | Final architecture uses a single-stream transformer block | +| Objective / sampler | Flow-matching velocity prediction with Euler sampling and a logit-normal schedule | Standard rectified-flow loss under `v`-parameterization, shifted logit-normal schedules | +| Parameters | 9.3B for released Ideogram 4 transformer variants | Not specified in the technical report | +| Text encoder | Qwen3-VL-8B-Instruct hidden states from 13 layers are concatenated | Qwen 3 VL final encoder with feature aggregation across layers | +| Attention | Multi-head attention with QK-RMSNorm and MRoPE | GQA with gated sigmoid attention | +| MLP | SwiGLU | SwiGLU | +| Norm | RMSNorm inside blocks, final LayerNorm | Zero-centered RMSNorm and QKNorm | +| Positional encoding | 3D MRoPE with text and image positions in one coordinate system | 3D axial RoPE | +| Timestep conditioning | Per-block AdaLN-style modulation from timestep embedding | Lightweight timestep modulation with bias | +| Autoencoder | Repo loads the released VAE from `vae/diffusion_pytorch_model.safetensors` | Qwen Image VAE early, FLUX 2 VAE for larger models | +| Prompt interface | Structured JSON captions; magic prompt converts plain text to JSON | Prompt expander maps short prompts into richer model-friendly captions | +| Extra controls | Bounding boxes, color palettes, JSON composition, asymmetric CFG | Style-reference system with strength and weighted style mixing | +| Post-training emphasis | Repo focuses on inference; README discusses safety and product controls | Pretraining, midtraining, SFT, preference optimization, RL, optional timestep distillation | +| License posture | Ideogram 4 Non-Commercial model license | Report says permissive license for weights and inference | + +## System Shape + +```mermaid +flowchart LR + subgraph I["Ideogram 4"] + I0["Plain text prompt"] --> I1["Magic prompt"] + I1 --> I2["Structured JSON caption"] + I2 --> I3["Qwen3-VL hidden-state taps"] + I3 --> I4["Single-stream DiT"] + I4 --> I5["Asymmetric CFG Euler sampler"] + I5 --> I6["VAE decode"] + I6 --> I7["Image"] + end + + subgraph K["Krea 2"] + K0["User prompt or style references"] --> K1["Prompt expander / style-reference system"] + K1 --> K2["Qwen 3 VL features"] + K2 --> K3["Single-stream DiT"] + K3 --> K4["Rectified-flow generation"] + K4 --> K5["Image"] + end +``` + +## Ideogram 4 Architecture From This Repo + +```mermaid +flowchart TD + P["Prompt"] --> T["Tokenizer + Qwen3-VL-8B-Instruct"] + T --> H["13 hidden-state taps: 0, 3, ..., 33, 35"] + H --> C["Concatenate and project language features"] + + N["Noise latents"] --> Patch["Patchified latent tokens"] + Time["Flow timestep"] --> Ada["AdaLN conditioning"] + + C --> S["Unified text + image token sequence"] + Patch --> S + S --> B["34 x Ideogram4TransformerBlock"] + Ada --> B + B --> V["Velocity prediction"] + V --> E["Euler flow update"] + E --> D["VAE decode"] +``` + +Key implementation anchors: + +- `Ideogram4Config` fixes the released transformer shape: `emb_dim=4608`, + `num_layers=34`, `num_heads=18`, `intermediate_size=12288`, and + `in_channels=128`. +- `QWEN3_VL_ACTIVATION_LAYERS` selects 13 intermediate Qwen3-VL layers and + concatenates them before projection into the DiT hidden size. +- `Ideogram4Attention` uses one QKV projection, Q/K RMSNorm, scaled dot-product + attention, and MRoPE. +- `Ideogram4TransformerBlock` uses timestep-conditioned scale and gate terms + for attention and MLP branches. +- `Ideogram4Pipeline` loads separate conditional and unconditional transformer + weights and applies asymmetric classifier-free guidance. + +## Krea 2 Architecture From The Report + +```mermaid +flowchart TD + D["Curated real-image data"] --> PT["Pretraining: 256 -> 512 -> 1024"] + PT --> MT["Midtraining"] + MT --> SFT["Domain SFT + model merging"] + SFT --> PO["Preference optimization"] + PO --> RL["Multi-reward RL"] + RL --> TD["Optional timestep distillation"] + TD --> K2["Krea 2 model family"] + + PE["Prompt expander SFT + RL"] --> K2 + SR["Style-reference system"] --> K2 +``` + +Architecture choices described by Krea: + +- Single-stream transformer block for the final model, after testing + single-stream, dual-stream, and hybrid-stream variants. +- GQA with gated sigmoid attention, chosen for efficiency and stability. +- SwiGLU MLPs. +- Lightweight timestep modulation with bias instead of heavier per-block MLPs. +- Qwen 3 VL as the final text encoder, with layerwise feature aggregation and + lightweight bidirectional layers over token features. +- Qwen Image VAE and FLUX 2 VAE are called out as the autoencoders that scaled + best for their runs. + +## The Real Difference + +```mermaid +flowchart LR + A["Shared modern recipe"] --> B["Latent flow / rectified flow"] + A --> C["DiT backbone"] + A --> D["Qwen3-VL-class conditioning"] + A --> E["SwiGLU + RMSNorm family choices"] + + B --> I["Ideogram 4 difference: structured design control"] + C --> I + D --> I + E --> I + + B --> K["Krea 2 difference: exploration and post-training stack"] + C --> K + D --> K + E --> K +``` + +The biggest difference is not that one is "a transformer" and the other is not: +both are transformer-based image generators. The difference is where each system +spends its complexity budget. + +Ideogram 4 spends it on a highly explicit inference contract: JSON captions, +bounding boxes, color palettes, multilingual text rendering, dual conditional / +unconditional branches, and resolution-aware sampling. Its released code makes +the inference graph concrete. + +Krea 2 spends it on distribution shaping: data curation, progressive-resolution +training, midtraining, SFT, preference optimization, RL, prompt-expander RL, and +style-reference control. Its report gives more detail about training systems and +post-training than about the exact inference API. + +## Practical Interpretation + +Use Ideogram 4 as the clearer reference when you care about: + +- explicit layout control; +- prompt schemas and JSON conditioning; +- typography-heavy design work; +- open local inference mechanics; +- reading actual implementation details. + +Use Krea 2 as the clearer reference when you care about: + +- broad aesthetic exploration; +- style references and style mixing; +- large-scale data infrastructure; +- post-training recipes for image models; +- prompt-expander training and RL. + +## Caveats + +This is an asymmetric comparison. Ideogram 4 is compared from executable repo +code plus local documentation. Krea 2 is compared from its public technical +report, not from a line-by-line inspection of its model source in this repo. +Where Krea gives training-level or product-level information but not a concrete +class or tensor shape, this document keeps the wording at that same level. + +## Sources + +- Local Ideogram 4 files: `src/ideogram4/modeling_ideogram4.py`, + `src/ideogram4/pipeline_ideogram4.py`, `src/ideogram4/scheduler.py`, + `src/ideogram4/constants.py`, and `README.md`. +- Krea 2 Technical Report, published June 23, 2026: + https://www.krea.ai/blog/krea-2-technical-report diff --git a/out_metal.png b/out_metal.png new file mode 100644 index 0000000..aac0baa Binary files /dev/null and b/out_metal.png differ diff --git a/out_metal_cfg.png b/out_metal_cfg.png new file mode 100644 index 0000000..84be84e Binary files /dev/null and b/out_metal_cfg.png differ diff --git a/run_metal.py b/run_metal.py new file mode 100644 index 0000000..6c89661 --- /dev/null +++ b/run_metal.py @@ -0,0 +1,280 @@ +#!/usr/bin/env python +"""Ideogram 4 en Apple Silicon (Metal/MPS), optimizado para 24 GB. + +Por qué este script y no run_inference.py: + * PyTorch MPS no soporta el dtype float8_e4m3fn -> aquí los pesos fp8 se + guardan como uint8 y se descuantizan con una LUT (e4m3->bf16) en el forward. + * El modelo completo son ~26 GB de pesos (2 DiTs + encoder) y no cabe en 24 GB. + Aquí se carga por fases liberando memoria: encoder -> (libera) -> DiT. + * Por defecto NO usa CFG (un solo DiT) para caber con holgura. CFG=1 carga + también el unconditional (va MUY justo de RAM en 24 GB). + +Uso: + ./.venv/bin/python run_metal.py --prompt "..." [--size 512] [--preset V4_TURBO_12] +""" +from __future__ import annotations + +import argparse +import gc +import time + +import torch +import torch.nn as nn +import torch.nn.functional as F + +# -------------------------------------------------------------------------- +# Shim: hacer que la ruta fp8 (weight-only e4m3) funcione en MPS via LUT uint8 +# -------------------------------------------------------------------------- +from ideogram4 import quantized_loading as ql + +FP8 = ql.FP8_WEIGHT_DTYPE # torch.float8_e4m3fn +_LUT: dict = {} + + +def _lut(device: torch.device, dtype: torch.dtype) -> torch.Tensor: + """Tabla de 256 entradas: byte e4m3 -> valor en `dtype`. Construida en CPU.""" + key = (str(device), dtype) + if key not in _LUT: + table = torch.arange(256, dtype=torch.uint8).view(FP8).to(torch.float32) + _LUT[key] = table.to(device=device, dtype=dtype) + return _LUT[key] + + +def _fp8_init(self, in_features, out_features, bias, compute_dtype): + nn.Module.__init__(self) + self.in_features = in_features + self.out_features = out_features + self.compute_dtype = compute_dtype + # peso guardado como uint8 (los bytes crudos del e4m3): MPS sí lo soporta + self.register_buffer("weight", torch.empty(out_features, in_features, dtype=torch.uint8)) + self.register_buffer("weight_scale", torch.empty(out_features, dtype=torch.float32)) + if bias: + self.register_buffer("bias", torch.empty(out_features, dtype=compute_dtype)) + else: + self.bias = None + + +def _fp8_forward(self, x): + lut = _lut(self.weight.device, x.dtype) + # gather: uint8 -> bf16 (descuantiza una capa cada vez, transitorio). + # MPS acepta índices int32; evita el temporal int64 de weight.long(). + w = lut[self.weight.to(torch.int32)] * self.weight_scale.to(x.dtype).unsqueeze(1) + bias = self.bias.to(x.dtype) if self.bias is not None else None + return F.linear(x, w, bias) + + +ql.Fp8Linear.__init__ = _fp8_init +ql.Fp8Linear.forward = _fp8_forward + + +def _patched_load(model, state_dict, device, dtype, *, assign=False, strict=True): + """Como load_fp8_state_dict original, pero convierte fp8 -> uint8 (para MPS).""" + import warnings + + prepared = {} + for k, v in state_dict.items(): + if v.dtype == FP8: + prepared[k] = v.view(torch.uint8).to(device=device) # bytes crudos a MPS + elif k.endswith(ql.FP8_SCALE_SUFFIX): + prepared[k] = v.to(device=device, dtype=torch.float32) + elif v.is_floating_point(): + prepared[k] = v.to(device=device, dtype=dtype) + else: + prepared[k] = v.to(device=device) + missing, unexpected = model.load_state_dict(prepared, strict=False, assign=assign) + if unexpected: + raise RuntimeError(f"unexpected keys after fp8 load: {unexpected[:10]}") + if missing: + if strict: + raise RuntimeError(f"missing keys after fp8 load: {missing[:10]}") + warnings.warn(f"missing keys after fp8 load: {missing[:10]}", stacklevel=2) + model.to(device) + + +ql.load_fp8_state_dict = _patched_load +# pipeline importó el nombre por valor -> parchear también ahí +import ideogram4.pipeline_ideogram4 as P # noqa: E402 + +P.load_fp8_state_dict = _patched_load + +# -------------------------------------------------------------------------- +from ideogram4 import PRESETS # noqa: E402 +from ideogram4.caption_verifier import CaptionVerifier # noqa: E402 +from ideogram4.latent_norm import get_latent_norm # noqa: E402 +from ideogram4.modeling_ideogram4 import Ideogram4Config # noqa: E402 +from ideogram4.pipeline_ideogram4 import ( # noqa: E402 + Ideogram4Pipeline, + Ideogram4PipelineConfig, + _build_transformer, + _load_autoencoder, + _load_indexed_or_single_state_dict, + _load_qwen3_vl, +) +from ideogram4.scheduler import ( # noqa: E402 + get_schedule_for_resolution, + make_step_intervals, +) +from huggingface_hub import hf_hub_download # noqa: E402 + + +def _free(): + gc.collect() + if torch.backends.mps.is_available(): + torch.mps.empty_cache() + + +def main(): + ap = argparse.ArgumentParser() + ap.add_argument("--prompt", required=True) + ap.add_argument("--output", default="out.png") + ap.add_argument("--size", type=int, default=512) + ap.add_argument("--preset", choices=sorted(PRESETS), default="V4_TURBO_12") + ap.add_argument("--seed", type=int, default=0) + ap.add_argument("--cfg", action="store_true", help="usa los 2 DiTs (MUY justo en 24GB)") + ap.add_argument("--magic", action="store_true", help="expande el prompt a JSON (IDEOGRAM_API_KEY)") + args = ap.parse_args() + + prompt = args.prompt + if args.magic: + import os + + from ideogram4 import DEFAULT_MAGIC_PROMPT, MAGIC_PROMPTS, aspect_ratio_from_size + + key = os.environ.get("IDEOGRAM_API_KEY") or os.environ.get("MAGIC_PROMPT_API_KEY") + if not key: + raise SystemExit("--magic necesita IDEOGRAM_API_KEY (en .env)") + print("[magic] expandiendo prompt a JSON...", flush=True) + m = MAGIC_PROMPTS[DEFAULT_MAGIC_PROMPT](api_key=key) + prompt = m.expand(args.prompt, aspect_ratio=aspect_ratio_from_size(args.size, args.size)) + + device = torch.device("mps") + dtype = torch.bfloat16 + h = w = args.size + cfg = Ideogram4PipelineConfig(weights_repo="ideogram-ai/ideogram-4-fp8") + tcfg = Ideogram4Config() + preset = PRESETS[args.preset] + t0 = time.time() + + # ---- Fase A: text encoder (una sola pasada), luego se libera ---- + print("[fase A] cargando text encoder...", flush=True) + tok, enc = _load_qwen3_vl( + cfg.weights_repo, + device, + dtype, + tokenizer_subfolder=cfg.tokenizer_subfolder, + text_encoder_subfolder=cfg.text_encoder_subfolder, + ) + pipe = Ideogram4Pipeline.__new__(Ideogram4Pipeline) + pipe.config = cfg + pipe.device = device + pipe.dtype = dtype + pipe.text_encoder = enc + pipe.text_tokenizer = tok + pipe.caption_verifier = CaptionVerifier() + shift, scale = get_latent_norm() + pipe.latent_shift = shift.to(device) + pipe.latent_scale = scale.to(device) + + inputs = pipe._build_inputs([prompt], height=h, width=w) + print("[fase A] codificando prompt...", flush=True) + llm_features = pipe._encode_text( + inputs["token_ids"], inputs["text_position_ids"], inputs["indicator"] + ).detach() + + pipe.text_encoder = None + del enc + _free() + print(f"[fase A] hecho ({time.time() - t0:.0f}s). Encoder liberado.", flush=True) + + # ---- Fase B: DiT(s) y bucle de denoising ---- + num_steps = preset.num_steps + schedule = get_schedule_for_resolution((h, w), known_mean=preset.mu, std=preset.std) + # OJO: el schedule usa float64/ndtri (no soportados en MPS) -> dejar en CPU. + # t_val/s_val son escalares Python, así que no hace falta que estén en device. + step_intervals = make_step_intervals(num_steps) + gw = torch.as_tensor(preset.guidance_schedule, dtype=torch.float32, device=device) + + num_image_tokens = inputs["num_image_tokens"] + grid_h, grid_w = inputs["grid_h"], inputs["grid_w"] + max_text = inputs["max_text_tokens"] + latent_dim = tcfg.in_channels + + print("[fase B] cargando DiT condicional...", flush=True) + cond_sd = _load_indexed_or_single_state_dict(cfg.weights_repo, cfg.conditional_index_filename) + cond = _build_transformer(tcfg, cond_sd, device, dtype) + del cond_sd + _free() + + uncond = None + neg_pos = neg_seg = neg_ind = neg_feat = None + if args.cfg: + print("[fase B] cargando DiT incondicional (CFG, memoria al límite)...", flush=True) + uncond_sd = _load_indexed_or_single_state_dict( + cfg.weights_repo, cfg.unconditional_index_filename + ) + uncond = _build_transformer(tcfg, uncond_sd, device, dtype) + del uncond_sd + _free() + neg_pos = inputs["position_ids"][:, max_text:] + neg_seg = inputs["segment_ids"][:, max_text:] + neg_ind = inputs["indicator"][:, max_text:] + neg_feat = torch.zeros( + 1, num_image_tokens, llm_features.shape[-1], dtype=llm_features.dtype, device=device + ) + + g = torch.Generator() # en CPU para evitar rarezas del generador MPS + g.manual_seed(args.seed) + z = torch.randn(1, num_image_tokens, latent_dim, dtype=torch.float32, generator=g).to(device) + pos_z = torch.empty(1, max_text + num_image_tokens, latent_dim, dtype=torch.float32, device=device) + pos_z[:, :max_text].zero_() + + print(f"[fase B] denoising: {num_steps} pasos (CFG={'on' if args.cfg else 'off'})", flush=True) + with torch.inference_mode(): + for i in range(num_steps - 1, -1, -1): + ts = time.time() + t_val = float(schedule(step_intervals[i + 1].unsqueeze(0)).item()) + s_val = float(schedule(step_intervals[i].unsqueeze(0)).item()) + t = torch.full((1,), t_val, dtype=torch.float32, device=device) + + pos_z[:, max_text:].copy_(z) + pos_out = cond( + llm_features=llm_features, + x=pos_z, + t=t, + position_ids=inputs["position_ids"], + segment_ids=inputs["segment_ids"], + indicator=inputs["indicator"], + ) + pos_v = pos_out[:, max_text:] + + if args.cfg: + neg_v = uncond( + llm_features=neg_feat, + x=z, + t=t, + position_ids=neg_pos, + segment_ids=neg_seg, + indicator=neg_ind, + ) + gw_i = gw[i] + v = gw_i * pos_v + (1.0 - gw_i) * neg_v + else: + v = pos_v # sin CFG: velocidad condicional cruda + + z = z + v * (s_val - t_val) + print(f" paso {num_steps - i}/{num_steps} ({time.time() - ts:.1f}s)", flush=True) + + del cond, uncond + _free() + + # ---- Fase C: decode con el VAE ---- + print("[fase C] decodificando con VAE...", flush=True) + ae_path = hf_hub_download(repo_id=cfg.weights_repo, filename=cfg.autoencoder_filename) + pipe.autoencoder = _load_autoencoder(ae_path, device, dtype) + imgs = pipe._decode(z, grid_h=grid_h, grid_w=grid_w) + imgs[0].save(args.output) + print(f"OK -> {args.output} (total {time.time() - t0:.0f}s)", flush=True) + + +if __name__ == "__main__": + main() diff --git a/try.sh b/try.sh new file mode 100755 index 0000000..51f7cdd --- /dev/null +++ b/try.sh @@ -0,0 +1,37 @@ +#!/usr/bin/env bash +# Prueba rápida de Ideogram 4 en Mac. Uso: +# ./try.sh "tu prompt aqui" +# Variables (opcionales): +# SIZE=512 resolución (múltiplo de 16) +# PRESET=... V4_TURBO_12 | V4_DEFAULT_20 | V4_QUALITY_48 +# DEVICE=cpu cpu funciona con fp8; mps NO soporta float8_e4m3fn +# MAGIC=1 activa magic prompt (necesita IDEOGRAM_API_KEY en .env) +set -euo pipefail +cd "$(dirname "$0")" + +# Carga credenciales desde .env si existe (HF_TOKEN, IDEOGRAM_API_KEY) +if [ -f .env ]; then set -a; . ./.env; set +a; fi + +PROMPT="${1:-a ginger cat wearing a tiny wizard hat reading a spellbook}" +SIZE="${SIZE:-512}" +PRESET="${PRESET:-V4_TURBO_12}" +DEVICE="${DEVICE:-cpu}" +MAGIC="${MAGIC:-0}" + +# Con MAGIC=1 expandimos el prompt a JSON (API hosted de Ideogram); si no, prompt plano. +if [ "$MAGIC" = "1" ]; then + MAGIC_ARGS=(--magic-prompt) +else + MAGIC_ARGS=(--no-magic-prompt --warn-on-caption-issues) +fi + +./.venv/bin/python run_inference.py \ + --prompt "$PROMPT" \ + --output out.png \ + --quantization fp8 \ + --device "$DEVICE" \ + --height "$SIZE" --width "$SIZE" \ + --sampler-preset "$PRESET" \ + "${MAGIC_ARGS[@]}" + +echo "Hecho -> $(pwd)/out.png"