Bài Lab tổng hợp — Qwen-Image-2.1 Architecture

Giới thiệu bài lab

Lab Qwen-Image-2.1: chạy thử vòng khử nhiễu 3 bước thu nhỏ từ 1A đến 6B bằng NumPy

Bài lab là một script Python dùng NumPy (thư viện tính toán trên mảng số) chạy ngay trong trình duyệt. Nó dựng lại cả 10 khối 1A đến 6B ở kích thước thu nhỏ: vector 8 chiều thay vì 4096, 2 lớp thay vì 32, 3 bước Euler thay vì 8. Chuỗi token gồm 11 phần tử: 7 token Prefix (chữ và ảnh tham chiếu, phần đứng trước) cộng 4 token của ảnh đích đang được khử nhiễu.

Ý tưởng cốt lõi của kiến trúc là tính phần Prefix một lần rồi dùng lại. Lab cho bạn tự thấy điều đó bằng số: kết quả dùng cache và kết quả tính lại toàn bộ chuỗi chỉ khác nhau cỡ 2.22e-16 (tức 2.22 × 10 mũ −16, gần như bằng 0).

Bài lab gồm

Sơ đồ 4 cột có khung tương tác ở đầu trang, để biết mỗi bảng số thuộc khối nào.
Bảng ánh xạ 1:1 cho từng mã 1A đến 6B: lab dựng gì và nhìn số nào để kiểm chứng.
Trình chạy trực tiếp với mã qwen-image-21.py sửa được, và các nút Chạy cell, Dừng, Đặt lại.
Output chia thành các khối [1A] đến [6B], rồi mục [KIỂM TRA & THỬ NGHIỆM] với 3 phép kiểm chứng tự động.
Trang bài đọc Lab đầy đủ (liên kết ở thanh công cụ) có hướng dẫn chi tiết và 4 bài tập kèm đáp án.

Làm xong, bạn có thể

Đọc bản đồ image_ids = [-1, -1, 0, 0, 0, 0, -1, 1, 1, 1, 1] và target_token_mask để nói token nào là chữ, ảnh tham chiếu hay ảnh đích.
Đếm khóa hợp lệ trên từng hàng của mặt nạ Block-Causal 11×11 và giải thích vì sao token chữ ở vị trí 1 thấy 2 khóa còn token ảnh đầu tiên ở vị trí 2 thấy 6 khóa.
Giải thích vì sao tắt điều kiện t = 0 khiến vector K (Key) của Prefix bị trôi (lệch cỡ 9.61e-01 đến 1.68e+00) làm cache mất tác dụng.
Giải thích vì sao Steps 2 và 3 chỉ cần tính 4 token ảnh đích mà vẫn ra kết quả khớp tính lại toàn bộ 11 token.
Nói được lab chứng minh gì (cơ chế cache và mặt nạ tính đúng) và không chứng minh gì (chất lượng ảnh, tốc độ GPU).

Phạm vi: lab chỉ chạy suy luận (forward) của một mô hình thu nhỏ và không có bước huấn luyện; trọng số là số ngẫu nhiên cố định bằng seed 21, nên ảnh RGBA 2×2 cuối cùng chỉ cho thấy luồng tính đúng, không phải ảnh đẹp. Muốn học huấn luyện, hãy xem khóa LLM Training (Text), vốn chỉ dạy mô hình văn bản; huấn luyện mô hình tạo ảnh từ nhiễu chưa có bài nào trên site.

Cách dùng bài lab

Xem sơ đồ 4 cột ở đầu trang: rê chuột hoặc bấm vào khung để mở pop-over, các nút lọc theo cột làm mờ khung của những cột còn lại, và nút Ẩn khung overlay cho bạn đọc ảnh gốc.
Đọc bảng ánh xạ bên dưới: cột thứ ba cho biết lab dựng gì ở mỗi mã, cột thứ tư cho biết nhìn số nào để kiểm chứng.
Đọc cấu hình thu nhỏ ở đầu mã: ảnh điều kiện 2×2 với 4 kênh RGBA, ngữ cảnh Qwen3-VL gồm 4 slot [chữ, chữ, ảnh, chữ], TURBO_SIGMAS = (1.0, 0.85, 0.45, 0.0) tức 3 bước Euler, và SEED = 21.
Dự đoán trước rồi bấm Chạy cell. Ở [1A] đến [2A], kiểm tra image_ids và target_token_mask: 7 chữ F (Prefix) đứng trước 4 chữ T (ảnh đích).
Ở [3A] đến [3C], đối chiếu tọa độ của hai khối ảnh, xem hàng t = 0.0 của bảng điều biến giữ nguyên khi t đổi từ 1000 sang 850, và đếm khóa hợp lệ trên từng hàng mặt nạ.
Ở [4A] đến [6B], so cached_k kích thước (7, 2, 4) của Step 1 với 4 token ảnh đích tính ở Steps 2 và 3, rồi đọc [KIỂM TRA & THỬ NGHIỆM]. Có thể sửa mã và chạy lại; một số hằng số như VLM_IMG_MASK có assert giữ chuỗi ở 7 + 4 token nên đổi sai sẽ báo lỗi, và nút Đặt lại khôi phục mã gốc.

Đọc kết quả đúng cách: 3 phép kiểm chứng chỉ nói rằng cơ chế cache và mặt nạ tính đúng với trọng số ngẫu nhiên; chúng không nói ảnh đẹp hay mô hình nhanh hơn.

Sơ đồ luồng dữ liệu 4 cột · 10 thành phần (1A → 6B)

Kiến trúc Qwen-Image-2.1 & Ánh xạ 1:1 với Bài Lab tổng hợp

Bắt đầu từ 1A ở góc dưới Cột 1 đi lên 1B → 2A, sang Cột 2 (3A → 3B → 3C), qua Cột 3 (4A ↔ 6A) và kết thúc ở Cột 4 (5A → 6B). Di chuột hoặc bấm vào từng khối để xem pop-over và đối chiếu với Lab.

Sơ đồ luồng dữ liệu 4 cột của kiến trúc Qwen-Image-2.1 gồm 10 thành phần từ 1A đến 6B

1A1 · RGBA VAE, VLM & Joint Packing

RGBA [B, 4, H, W] → AutoencoderKLQwenImage21 → ẑ = (z − mean)/std ∈ [B, N_img, 64] (patch_size=1) → img_in(64 → 4096)

1A · RGBA 64-Ch VAE & Unpatched Latents

Biến ảnh thành token mà Transformer đọc được: VAE nén ảnh RGBA thành một lưới latent nhỏ, mỗi ô có 64 kênh, rồi mỗi ô được chiếu thành một token 4096 chiều.

Mục đích: Transformer không thể xử lý trực tiếp hàng triệu điểm ảnh. Bước này nén ảnh xuống một lưới gọn nhưng vẫn giữ đủ thông tin, kể cả kênh trong suốt Alpha, để cuối pipeline (6B) giải nén lại thành ảnh.
Đầu vào: Ảnh tham chiếu 4 kênh R, G, B, Alpha có dạng [B, 4, H, W]. Ảnh đích không đi qua encoder: nó bắt đầu là nhiễu ngẫu nhiên tạo thẳng trên lưới latent cùng kích thước.
Đầu ra: Mỗi ảnh thành N_img token 4096 chiều (1 ô latent = 1 token vì patch_size = 1), sau khi chuẩn hóa từng kênh (z − mean) / std và chiếu img_in 64 → 4096. Các token này được lắp vào chuỗi chung ở 2A.

Bắt đầu đọc sơ đồ từ góc dưới bên trái (1A): AutoencoderKLQwenImage21 nhận ảnh 4 kênh R, G, B, Alpha và nén thành các vector latent 64 kênh. Trước khi đi vào mạng DiT, mỗi kênh được chuẩn hóa theo thống kê riêng ẑ = (z − latents_mean) / latents_std rồi chiếu thẳng qua img_in (64 → 4096) với patch_size = 1 (không gộp ô 2×2 tại đây) để bảo toàn chi tiết biên trong suốt.

Bài 1A · RGBA 64-Channel VAE & Unpatched Latents Xem bước [1A] trong Bài Lab tổng hợp

BẢNG ÁNH XẠ 1:1 · MỖI KHỐI KIẾN TRÚC Ở TAB 01 ĐƯỢC DỰNG VÀ KIỂM CHỨNG TRONG LAB NHƯ THẾ NÀO?

Mã	Thành phần ở Tab 01 (bấm để mở bài)	Lab NumPy dựng gì (`qwen-image-21.py`)	Cách đọc số liệu & phép kiểm chứng
1A	RGBA 64-Channel VAE & Unpatched Latents ↗	Chuẩn hóa 4 pixel RGBA (2×2×4) qua z = (z − mean)/std và chiếu thẳng qua img_in với PATCH_SIZE = 1 (4 → 8).	Đối chiếu vector sau chuẩn hóa và sau img_in: denoise luôn diễn ra trong không gian chuẩn hóa.
1B	Qwen3-VL Context & Zero-Center RMSNorm (w + 1) ↗	Chuẩn hóa chuỗi ngữ cảnh VLM bằng rsqrt(mean(x²) + ε) · (w + 1) và chiếu qua MLP 2 tầng GELU.	Kiểm tra các slot văn bản đi qua txt_in GELU MLP trong khi vector slot ảnh được giữ nguyên để mở rộng ở 2A.
2A	Mở rộng 4× Slot ảnh (2×2) & Ghép chuỗi đồng nhất ↗	Dùng repeats = [1, 1, 4, 1] bung slot ảnh VLM thành 4 vị trí và nối thêm 4 token ảnh đích → S_joint = 11 (prefix 7 + target 4).	Kiểm tra image_ids = [-1, -1, 0, 0, 0, 0, -1, 1, 1, 1, 1] và target_token_mask đánh dấu đúng 4 vị trí cuối.
3A	3D Axial RoPE (Frame, Height, Width) & Lưới tâm không ↗	Token chữ tăng đều cả 3 trục (0,0,0) → (1,1,1) → (2,2,2); hai khối ảnh 2×2 ghim frame = 0 và căn tâm [(-1,-1), (-1,0), (0,-1), (0,0)].	Xác nhận tọa độ không gian (h, w) của ảnh tham chiếu và ảnh đích trùng khớp, không phụ thuộc vị trí đứng trong câu.
3B	Điều kiện nhân quả (t=0 vs t) & Shared Modulation ↗	Tạo bảng modulation 2 hàng [t_step, 0.0]: 7 token Prefix luôn đọc hàng t = 0.0, còn 4 token Target đọc hàng t_step.	So sánh Step 1 (t=1000) và Step 2 (t=850): hàng 1 (t=0.0) lệch đúng 0.00e+00 trong khi hàng 0 đổi theo bước nhiễu.
3C	Mặt nạ Block-Causal 11×11 (Nhân quả + Hai chiều) ↗	Dựng ma trận 11×11 theo quy tắc (q >= kv) OR (q_img_id == kv_img_id >= 0).	Kiểm chứng (3): đổi token target (pos 7) gây lệch 0.00e+00 lên 7 token prefix; đổi pixel ảnh điều kiện (pos 5) lan truyền song hướng khắp khối 2..5.
4A	Single-Stream DiT Attention & QK-RMSNorm ↗	Điều biến LN(x)·(1 + scale1) + shift1, chiếu Q/K/V chung một luồng, chuẩn hóa QK-RMSNorm từng head, xoay 3D RoPE và cộng residual qua tanh(gate1).	Quan sát trọng số attention của token 7 trên cả 11 khóa hợp lệ và biên độ cổng tanh(gate1) ≤ 1.
5A	Mạng SwiGLU có điều biến (3× Ratio) & Cổng Tanh ↗	Mở rộng chiều ẩn 3× (8 → 24 → 8), nhân cổng SiLU(W_gate x) ⊙ (W_up x) và cộng residual qua tanh(gate2).	Đối chiếu 6 cột đầu của nhánh gate, nhánh up và tích SwiGLU trước khi thu về chiều ẩn gốc.
6A	Prefix KV Cache (Step 1 Extract → Steps 2..3 Cached) ↗	Step 1 (extract) lưu cached_k, cached_v shape (7, 2, 4); Steps 2–3 (cached) chỉ tính Q/K/V cho 4 token target và ghép cat([cached_k, target_k]).	Kiểm chứng (1) & (2): cached vs recompute toàn trình khớp tới 2.22e-16 (< 1e-9); nếu tắt t=0 trên prefix thì cache bị trôi tới 1.68e+00!
6B	Scale-Only Final AdaLN, Velocity & Vòng lặp Euler → RGBA ↗	Chuẩn hóa đầu ra LN(h)·(1 + scale_final) không có shift, chiếu ra vận tốc v_θ, tích phân 3 bước Euler (1.0 → 0.85 → 0.45 → 0.0) và giải mã ra ảnh RGBA 2×2×4.	Kiểm tra 4 kênh R, G, B, Alpha đầu ra khớp tuyệt đối (sai lệch 1.11e-16) giữa đường dùng cache và đường tính lại toàn chuỗi.

Trình chạy Lab trực tiếp · `qwen-image-21.py`

qwen-image-21.pyPYTHON

"""Lab Qwen-Image-2.1: chạy toàn trình 1A→6B bằng NumPy thuần.

Bài lab giải phẫu MỘT lượt diffusion (3 bước Euler Turbo) của kiến trúc
Qwen-Image-2.1 ở kích thước thu nhỏ, đồng bộ mã khối với sơ đồ kiến trúc
Tab 1 (1A..6B):

[1A] VAE RGBA 64 kênh: encode ảnh điều kiện, chuẩn hóa latent (z-mean)/std, img_in
  [1B] Ngữ cảnh Qwen3-VL: RMSNorm zero-center (w+1) và txt_in MLP GELU
  [2A] Mở rộng slot ảnh VLM 4× (repeats=[1,1,4,1]) & đóng gói chuỗi joint 11 token
  [3A] RoPE trục 3D (frame, height, width) với lưới ảnh zero-centered
  [3B] Điều kiện nhân quả theo t: bảng modulation [t_step, 0.0] — prefix đọc t=0
  [3C] Mask attention block-causal 11×11: text nhân quả, khối ảnh song hướng
  [4A] Attention DiT single-stream & QK-RMSNORM per-head & RoPE & cổng tanh
  [5A] FFN SwiGLU điều biến (ratio 3×) & cổng tanh
  [6A] Prefix KV cache: Step 1 extract → Steps 2-3 cached (chỉ 4 token target)
  [6B] Final AdaLN scale-only, velocity, vòng Euler 3 bước, decode RGBA

Tiến trình thực thi: TURBO_SIGMAS = (1.0, 0.85, 0.45, 0.0) — 3 bước Euler.
Step 1 (sigma=1.0) chạy kv_cache_mode="extract": tính cả 11 token và LƯU
K/V của 7 token prefix vào QwenImage21KVCache. Steps 2-3 (sigma 0.85, 0.45)
chạy kv_cache_mode="cached": chỉ chiếu Q/K/V cho 4 token target
(joint_hidden_states[7:]), K/V prefix lấy từ cache — rồi đối chiếu với lượt
tính lại toàn bộ 11 token bằng mask block-causal ở [KIỂM TRA & THỬ NGHIỆM].

Weights chưa học, sinh một lần bằng np.random.default_rng(21); chạy lại cho
ra cùng số. Vùng số ở mức micro (DIM 8 thay 4096, 2 lớp thay 32, 4 kênh
latent thay 64) nên chỉ kiểm chứng CƠ CHẾ; KHÔNG tái tạo checkpoint thật,
KHÔNG đo tốc độ GPU, KHÔNG chứng minh chất lượng ảnh sinh ra.
"""
import numpy as np

# ── Cấu hình thu nhỏ (sửa các số này để thí nghiệm tiếp) ────────────────────
RGBA_CHANNELS = 4        # đầu ra VAE có 4 kênh R, G, B, Alpha
LATENT_CHANNELS = 4      # đứng thay 64 kênh latent của AutoencoderKLQwenImage
PATCH_SIZE = 1           # thu nhỏ của patch 2×2 trong prod (patch 1×1 = không chia ô)
DIM = 8                  # đứng thay inner_dim = 4096
N_HEADS = 2              # số head attention
HEAD_DIM = 4             # đứng thay head_dim = 128
AXES_DIMS_ROPE = (2, 1, 1)  # RoPE chia (frame, height, width); prod là (16, 56, 56)
MLP_RATIO = 3
MLP_DIM = DIM * MLP_RATIO   # 24; prod FFN ratio ~3× của 4096
N_LAYERS = 2             # đứng thay 32 transformer blocks
IMG_TOKENS_PER_SLOT = 4  # mỗi slot ảnh = lưới latent 2×2 = 4 token
VLM_SLOTS = 4            # ngữ cảnh VLM: [text, text, cond_img, text]
VLM_IMG_MASK = [False, False, True, False]  # slot nào là ảnh trong ngữ cảnh VLM
IMG_SHAPES = [(1, 2, 2), (1, 2, 2)]         # ảnh điều kiện 2×2 và ảnh target 2×2
TURBO_SIGMAS = (1.0, 0.85, 0.45, 0.0)       # 3 bước Euler Turbo
TIME_DIM = 8             # chiều embedding timestep
ROPE_BASE = 10000.0
EPS_NORM = 1e-6
MASK_NEG = -1e9          # điểm attention cho vị trí bị che
SEED = 21

# ── Suy ra chiều chuỗi joint từ các hằng số trên ────────────────────────────
PREFIX_LEN = sum(IMG_TOKENS_PER_SLOT if s else 1 for s in VLM_IMG_MASK)  # 2+4+1 = 7
TARGET_LEN = IMG_SHAPES[1][1] * IMG_SHAPES[1][2]                        # 2×2 = 4
S_JOINT = PREFIX_LEN + TARGET_LEN                                       # 11
assert PREFIX_LEN == 7 and TARGET_LEN == 4 and S_JOINT == 11
assert sum(VLM_IMG_MASK) == 1
assert N_HEADS * HEAD_DIM == DIM

# ── Weights (chưa học): sinh MỘT lần, mọi pass dùng chung ───────────────────
rng = np.random.default_rng(SEED)
def _n(shape, s):
    return rng.normal(0.0, s, shape)

# VAE RGBA→latent và latent→RGBA
LATENT_MEAN = _n((LATENT_CHANNELS,), 0.20)        # thống kê latent của VAE
LATENT_STD = 0.5 + np.abs(_n((LATENT_CHANNELS,), 0.25))
W_ENC = _n((RGBA_CHANNELS, LATENT_CHANNELS), 0.45)
B_ENC = _n((LATENT_CHANNELS,), 0.20)
W_DEC = _n((LATENT_CHANNELS, RGBA_CHANNELS), 0.45)
B_DEC = _n((RGBA_CHANNELS,), 0.20)

# Ảnh điều kiện demo cố định: (1, 2, 2, 4) = 4 pixel × 4 kênh R, G, B, A
IMG_DEMO = np.array([[[[0.86, 0.10, 0.14, 0.95], [0.18, 0.55, 0.87, 0.90]],
                      [[0.95, 0.83, 0.20, 0.85], [0.10, 0.78, 0.36, 0.80]]]])

W_IMG_IN = _n((LATENT_CHANNELS, DIM), 0.35)       # img_in: latent (4,) → inner (8,)

# Ngữ cảnh VLM (đóng vai hidden states của Qwen3-VL): 4 token × 8 chiều
CONTEXT = _n((VLM_SLOTS, DIM), 0.40)
RMS_CTXW = _n((DIM,), 0.30)                       # RMSNorm zero-center của ngữ cảnh
W_TXT1 = _n((DIM, DIM), 0.30)                     # txt_in MLP GELU: lớp 1
W_TXT2 = _n((DIM, DIM), 0.30)                     # txt_in MLP GELU: lớp 2

# Bảng modulation timestep: 2 hàng [t_step, 0.0] × (6 nhóm × DIM)
W_T_MOD = _n((TIME_DIM, 6 * DIM), 0.20)
SCALE_SHIFT_TABLE = _n((2, 6 * DIM), 0.15)        # production: AdaLN scale-shift
W_OUT_MOD = _n((TIME_DIM, DIM), 0.20)             # final AdaLN SCALE-ONLY (micro)

# Per-layer DiT
W_Q = [_n((DIM, DIM), 0.30) for _ in range(N_LAYERS)]
W_K = [_n((DIM, DIM), 0.30) for _ in range(N_LAYERS)]
W_V = [_n((DIM, DIM), 0.30) for _ in range(N_LAYERS)]
W_O = [_n((DIM, DIM), 0.30) for _ in range(N_LAYERS)]
GATE1 = [_n((DIM,), 0.30) for _ in range(N_LAYERS)]       # cổng tanh attention
W_SWIGLU_GATE = [_n((DIM, MLP_DIM), 0.30) for _ in range(N_LAYERS)]
W_SWIGLU_UP = [_n((DIM, MLP_DIM), 0.30) for _ in range(N_LAYERS)]
W_SWIGLU_DOWN = [_n((MLP_DIM, DIM), 0.25) for _ in range(N_LAYERS)]
GATE2 = [_n((DIM,), 0.30) for _ in range(N_LAYERS)]       # cổng tanh FFN
RMS_QW = _n((HEAD_DIM,), 0.30)                    # QK-RMSNORM zero-center (Q)
RMS_KW = _n((HEAD_DIM,), 0.30)                    # QK-RMSNORM zero-center (K)

W_UNPATCH = _n((DIM, LATENT_CHANNELS), 0.35)      # inner (8,) → latent (4,)

# ── Hàm số học dùng chung ───────────────────────────────────────────────────
def rmsnorm(x, w):
    """RMSNorm zero-center (phong cách Qwen3): x · rsqrt(mean(x²)+eps) · (w+1)."""
    ms = np.mean(x * x, axis=-1, keepdims=True)
    return x / np.sqrt(ms + EPS_NORM) * (1.0 + w)

def layernorm(x):
    """LayerNorm không affine nghĩa (norm trong block DiT của prod)."""
    mu = x.mean(axis=-1, keepdims=True)
    var = ((x - mu) ** 2).mean(axis=-1, keepdims=True)
    return (x - mu) / np.sqrt(var + EPS_NORM)

def gelu(x):
    """GELU gần đúng tanh, dùng trong txt_in (mã [1B])."""
    return 0.5 * x * (1.0 + np.tanh(np.sqrt(2.0 / np.pi) * (x + 0.044715 * x ** 3)))

def silu(x):
    return x / (1.0 + np.exp(-x))

def softmax(v, axis=-1):
    v = v - v.max(axis=axis, keepdims=True)
    e = np.exp(v)
    return e / e.sum(axis=axis, keepdims=True)

def standardize(z):
    """Chuẩn hóa latent theo thống kê VAE: (z - mean) / std (mã [1A])."""
    return (z - LATENT_MEAN) / LATENT_STD

def unstandardize(z_std):
    """Đưa latent về không gian VAE gốc: z · std + mean (mã [6B])."""
    return z_std * LATENT_STD + LATENT_MEAN

# ── [3A] RoPE trục 3D: tọa độ (frame, height, width) zero-centered ──────────
# Prod dùng 64 cặp xoay (frame 8 + height 28 + width 28) với head_dim 128.
# Micro chỉ có 2 cặp (head_dim 4): cặp 0 xoay theo TRỤC FRAME; cặp 1 chia sẻ
# cho hai trục height và width (góc cộng dồn) để cả 3 trục đều tham gia.
ROPE_TF = 1.0                                    # trục frame (chiếm 2 chiều RoPE)
ROPE_TH = ROPE_BASE ** (-2.0 * 2.0 / HEAD_DIM)   # trục height (chiếm 1 chiều)
ROPE_TW = ROPE_BASE ** (-2.0 * 3.0 / HEAD_DIM)   # trục width (chiếm 1 chiều)

def rope_coords():
    """Tọa độ (frame, height, width) của 11 token.

Text TIẾN DẦN cả 3 trục (t, t, t); hai khối ảnh 2×2 ghim frame=0 và căn
    giữa (height, width) về lưới [(-1,-1), (-1,0), (0,-1), (0,0)].
    """
    grid = [(h, w) for h in (-1, 0) for w in (-1, 0)]
    coords = []
    text_i = 0
    for is_img in VLM_IMG_MASK:
        if is_img:
            coords += [(0, h, w) for (h, w) in grid]
        else:
            coords.append((text_i, text_i, text_i))
            text_i += 1
    coords += [(0, h, w) for (h, w) in grid]     # khối target — frame cũng 0
    return np.array(coords, dtype=float)         # (11, 3) = (fhw)

ROPE_COORDS = rope_coords()
ANGLE_PAIR = np.stack([                           # (2, 11) — 2 cặp góc
    ROPE_COORDS[:, 0] * ROPE_TF,                  # cặp 0: trục frame
    ROPE_COORDS[:, 1] * ROPE_TH + ROPE_COORDS[:, 2] * ROPE_TW,  # cặp 1: h + w
])

def rope_rotate_heads(x, angles_per_pos):
    """Xoay (heads, S, head_dim) theo 2 cặp góc (2, S).

Cặp 0 xoay chiều 0↔1, cặp 1 xoay chiều 2↔3 (biểu diễn phức của RoPE).
    """
    c0, s0 = np.cos(angles_per_pos[0]), np.sin(angles_per_pos[0])
    c1, s1 = np.cos(angles_per_pos[1]), np.sin(angles_per_pos[1])
    out = np.empty_like(x)
    out[..., 0] = x[..., 0] * c0 - x[..., 1] * s0
    out[..., 1] = x[..., 0] * s0 + x[..., 1] * c0
    out[..., 2] = x[..., 2] * c1 - x[..., 3] * s1
    out[..., 3] = x[..., 2] * s1 + x[..., 3] * c1
    return out

# ── [3C] Mask block-causal 11×11 ────────────────────────────────────────────
def image_block_ids():
    ids = np.full(S_JOINT, -1, dtype=int)
    ids[2:6] = 0    # khối ảnh điều kiện (4 token liền sau 2 token text)
    ids[7:11] = 1   # khối ảnh target (4 token ở cuối chuỗi)
    return ids

IMAGE_BLOCK_IDS = image_block_ids()

def block_causal_mask():
    """Cho phép (q >= kv) hoặc cùng khối ảnh; text thì nhân quả thuần."""
    allowed = np.zeros((S_JOINT, S_JOINT), dtype=bool)
    for q in range(S_JOINT):
        for kv in range(S_JOINT):
            same = IMAGE_BLOCK_IDS[q] >= 0 and IMAGE_BLOCK_IDS[q] == IMAGE_BLOCK_IDS[kv]
            allowed[q, kv] = (q >= kv) or same
    return allowed

BLOCK_MASK = block_causal_mask()

# ── [3B] Embedding timestep & bảng modulation 2 hàng ───────────────────────
def timestep_emb(t):
    """Embedding sin-cos của timestep t (t = 1000·sigma theo FlowMatch)."""
    freqs = ROPE_BASE ** (-2.0 * np.arange(TIME_DIM // 2) / TIME_DIM)
    ang = float(t) * freqs
    return np.concatenate([np.cos(ang), np.sin(ang)])   # (TIME_DIM,)

def modulation_table(t_step):
    """Bảng modulation 2 HÀNG [t_step, 0.0]; mỗi hàng (6 × DIM) = 48 cột.

Hàng 0 dành cho token target (đang denoise ở bước này); hàng 1 là
    t = 0.0 — TẤT CẢ token prefix đọc hàng này, để prefix constant.
    """
    emb = np.stack([timestep_emb(t_step), timestep_emb(0.0)])   # (2, TIME_DIM)
    return emb @ W_T_MOD + SCALE_SHIFT_TABLE                    # (2, 6*DIM)

def chunk_mod(rows):
    """(S, 6*DIM) → (S, 6, DIM): shift1, scale1, gate1, shift2, scale2, gate2."""
    return rows.reshape(rows.shape[0], 6, DIM)

# ── Thành phần production được micro hóa ───────────────────────────────────
class AutoencoderKLQwenImage21:
    """VAE của Qwen-Image-2.1 (micro): encode RGBA→latent, decode latent→RGBA.

Production dùng 64 kênh latent và chỉ xuất RA RGBA 4 kênh khi decode.
    Micro dùng 4 kênh đơn giản hóa, giữ đúng vai trò encode → chuẩn hóa → decode.

"""

def encode(self, rgba):
        flat = np.asarray(rgba).reshape(-1, RGBA_CHANNELS)
        return flat @ W_ENC + B_ENC            # (4, LATENT_CHANNELS)

def decode(self, z):
        flat = np.asarray(z).reshape(-1, LATENT_CHANNELS)
        rgba = np.tanh(flat @ W_DEC + B_DEC)   # tanh đưa về [-1, 1]
        return rgba.reshape(IMG_SHAPES[1][1], IMG_SHAPES[1][2], RGBA_CHANNELS)

class QwenImage21KVCache:
    """Cache K/V của PREFIX (7 token) theo từng lớp.

Mode "extract" (Step 1): ghi prefix K/V — mỗi lớp một mảng (7, 2, 4)
    = (prefix, heads, head_dim). Mode "cached" (Steps 2-3): CHỈ ĐỌC.
    """

def __init__(self, n_layers):
        self.k = [None] * n_layers
        self.v = [None] * n_layers
        self.mode = "extract"

def store(self, layer, k_heads, v_heads):
        # k_heads (2, 11, 4) → lưu phần prefix dưới dạng (7, 2, 4) theo S-major.
        self.k[layer] = k_heads[:, :PREFIX_LEN, :].transpose(1, 0, 2).copy()
        self.v[layer] = v_heads[:, :PREFIX_LEN, :].transpose(1, 0, 2).copy()

def prefix_k_heads(self, layer):
        return self.k[layer].transpose(1, 0, 2)   # đọc lại: (2, 7, 4)

def prefix_v_heads(self, layer):
        return self.v[layer].transpose(1, 0, 2)

class FlowMatchEulerDiscreteScheduler:
    """Scheduler Flow-Match Euler: x_{k+1} = x_k + (σ_{k+1} − σ_k) · v_θ.

Timestep đưa vào transformer là t = 1000 · sigma (quy ước FlowMatch).
    """

def __init__(self, sigmas):
        self.sigmas = np.asarray(sigmas, dtype=float)
        self.step_index = 0

def timestep(self):
        return 1000.0 * self.sigmas[self.step_index]

def sigma(self):
        return float(self.sigmas[self.step_index])

def step(self, x, velocity):
        s0 = float(self.sigmas[self.step_index])
        s1 = float(self.sigmas[self.step_index + 1])
        self.step_index += 1
        return x + (s1 - s0) * velocity

@property
    def finished(self):
        return self.step_index >= len(self.sigmas) - 1

class QwenImage21Transformer2DModel:
    """DiT single-stream 2 lớp trên chuỗi joint 11 token (7 prefix + 4 target).

Mode "extract": tính cả 11 token, lưu prefix K/V vào cache.
    Mode "cached":  chỉ tính 4 token target (joint_hidden_states[7:]), prefix
    K/V lấy từ cache — bằng chính xác nhờ prefix luôn đọc t = 0.0 (mã [3B]).
    """

def forward(self, target_std, t_step, cache=None, mode="extract",
                causal_condition=True, capture=None, perturb=None):
        assert mode in ("extract", "cached")
        table = modulation_table(t_step)                                # (2, 48)
        target_tokens = target_std @ W_IMG_IN                           # (4, 8)

if mode == "extract":
            joint = np.concatenate([PREFIX_TOKENS, target_tokens], axis=0)  # (11, 8)
            if perturb:
                for (idx, dim, delta) in perturb:
                    joint[idx][dim] += delta
            if capture is not None:
                capture["joint_input"] = joint.copy()
            # Mỗi token đọc ĐÚNG 1 hàng của bảng [t_step, 0.0]:
            prefix_row = 0 if not causal_condition else 1
            mod_rows = np.concatenate([
                np.repeat(table[prefix_row:prefix_row + 1], PREFIX_LEN, axis=0),
                np.repeat(table[0:1], TARGET_LEN, axis=0)], axis=0)         # (11, 48)
            mod_tok = chunk_mod(mod_rows)                                   # (11, 6, 8)
            h = joint
            for l in range(N_LAYERS):
                h, k_rot, v, cap_l = self._layer_full(l, h, mod_tok,
                                                      capture if l == 0 else None)
                if cache is not None and cache.mode == "extract":
                    cache.store(l, k_rot, v)
                if cache is not None and capture is not None and l == 0:
                    capture["k_full_step"] = k_rot.transpose(1, 0, 2).copy()
            joint_final = h
        else:
            mod_t = chunk_mod(table[0:1].repeat(TARGET_LEN, axis=0))        # (4, 6, 8)
            h = target_tokens.copy()
            for l in range(N_LAYERS):
                h, cap_l = self._layer_cached(l, h, mod_t, cache)
                if capture is not None and l == 0:
                    capture["cached_rep"] = cap_l
            joint_final = None

# [6B] Final AdaLN SCALE-ONLY: out = LN(h) · (1 + scale_final).
        h_t = h[PREFIX_LEN:] if mode == "extract" else h
        scale_final = timestep_emb(t_step) @ W_OUT_MOD + SCALE_OUT_BASE     # (8,)
        out = layernorm(h_t) * (1.0 + scale_final[None, :])
        velocity = out @ W_UNPATCH                                          # (4, 4)
        if capture is not None:
            capture["scale_final"] = scale_final.copy()
            capture["out_after_adaln"] = out.copy()
        return {"velocity": velocity, "joint_final": joint_final,
                "scale_final": scale_final, "table": table}

def _layer_full(self, l, h, mod_tok, capture):
        """MỘT lớp DiT trên TẤT CẢ 11 token (mã [4A] + [5A])."""
        shift1, scale1, gate1, shift2, scale2, gate2 = [mod_tok[:, i, :] for i in range(6)]
        normed = layernorm(h)
        x_mod = normed * (1.0 + scale1) + shift1          # [4A] modulated LN
        q = (x_mod @ W_Q[l]).reshape(S_JOINT, N_HEADS, HEAD_DIM).transpose(1, 0, 2)
        k = (x_mod @ W_K[l]).reshape(S_JOINT, N_HEADS, HEAD_DIM).transpose(1, 0, 2)
        v = (x_mod @ W_V[l]).reshape(S_JOINT, N_HEADS, HEAD_DIM).transpose(1, 0, 2)
        q_rms = rmsnorm(q, RMS_QW)                        # QK-RMSNORM per-head
        k_rms = rmsnorm(k, RMS_KW)
        q_rot = rope_rotate_heads(q_rms, ANGLE_PAIR)
        k_rot = rope_rotate_heads(k_rms, ANGLE_PAIR)
        scores = (q_rot @ k_rot.transpose(0, 2, 1)) / np.sqrt(HEAD_DIM) + np.where(
            BLOCK_MASK[None, :, :], 0.0, MASK_NEG)
        weights = softmax(scores, axis=-1)                # (2, 11, 11)
        attn = weights @ v                                # (2, 11, 4)
        attn_flat = attn.transpose(1, 0, 2).reshape(S_JOINT, DIM)
        out = attn_flat @ W_O[l]
        h = h + np.tanh(gate1) * out                      # [4A] cổng tanh
        normed2 = layernorm(h)
        ffn_in = normed2 * (1.0 + scale2) + shift2        # [5A] modulated LN
        g = ffn_in @ W_SWIGLU_GATE[l]
        u = ffn_in @ W_SWIGLU_UP[l]
        hidden = silu(g) * u
        ffn = hidden @ W_SWIGLU_DOWN[l]
        h = h + np.tanh(gate2) * ffn                      # [5A] cổng tanh
        if capture is not None:
            capture.update({"x_mod": x_mod.copy(), "q_raw": q.copy(),
                            "q_rms": q_rms.copy(), "k_rms": k_rms.copy(),
                            "q_rot": q_rot.copy(), "k_rot": k_rot.copy(),
                            "v": v.copy(), "weights": weights.copy(),
                            "ffn_in": ffn_in.copy(), "ffn_g": g.copy(),
                            "ffn_u": u.copy(), "ffn_hidden": hidden.copy(),
                            "gate1": gate1.copy(), "gate2": gate2.copy(),
                            "attn_flat": attn_flat.copy()})
        return h, k_rot, v, {}

def _layer_cached(self, l, h, mod_t, cache):
        """MỘT lớp DiT trên 4 token TARGET, prefix đọc từ cache (mã [6A])."""
        shift1, scale1, gate1, shift2, scale2, gate2 = [mod_t[:, i, :] for i in range(6)]
        normed = layernorm(h)
        x_mod = normed * (1.0 + scale1) + shift1
        q = (x_mod @ W_Q[l]).reshape(TARGET_LEN, N_HEADS, HEAD_DIM).transpose(1, 0, 2)
        k = (x_mod @ W_K[l]).reshape(TARGET_LEN, N_HEADS, HEAD_DIM).transpose(1, 0, 2)
        v = (x_mod @ W_V[l]).reshape(TARGET_LEN, N_HEADS, HEAD_DIM).transpose(1, 0, 2)
        q = rmsnorm(q, RMS_QW)
        k_t = rmsnorm(k, RMS_KW)
        q = rope_rotate_heads(q, ANGLE_PAIR[:, PREFIX_LEN:])
        k_t = rope_rotate_heads(k_t, ANGLE_PAIR[:, PREFIX_LEN:])
        # cat([cached_k (7,2,4), target_k (4,2,4)]) → k_full (11, 2, 4) S-major.
        k_full = np.concatenate([cache.k[l], k_t.transpose(1, 0, 2)], axis=0)
        v_full = np.concatenate([cache.v[l], v.transpose(1, 0, 2)], axis=0)
        k_full_h = k_full.transpose(1, 0, 2)              # (2, 11, 4)
        v_full_h = v_full.transpose(1, 0, 2)
        scores = (q @ k_full_h.transpose(0, 2, 1)) / np.sqrt(HEAD_DIM) + np.where(
            BLOCK_MASK[PREFIX_LEN:][None, :, :], 0.0, MASK_NEG)   # (2, 4, 11)
        weights = softmax(scores, axis=-1)
        attn = weights @ v_full_h
        attn_flat = attn.transpose(1, 0, 2).reshape(TARGET_LEN, DIM)
        out = attn_flat @ W_O[l]
        h = h + np.tanh(gate1) * out
        normed2 = layernorm(h)
        ffn_in = normed2 * (1.0 + scale2) + shift2
        hidden = silu(ffn_in @ W_SWIGLU_GATE[l]) * (ffn_in @ W_SWIGLU_UP[l])
        ffn = hidden @ W_SWIGLU_DOWN[l]
        h = h + np.tanh(gate2) * ffn
        rep = {"k_cached_shape": cache.k[l].shape, "k_t_shape": k_t.transpose(
            1, 0, 2).shape, "k_full_shape": k_full.shape}
        return h, rep

def txt_in(ctx):
    """[1B] txt_in của ngữ cảnh Qwen3-VL: RMSNorm rồi MLP GELU (micro)."""
    hidden = gelu(rmsnorm(ctx, RMS_CTXW) @ W_TXT1)
    return hidden @ W_TXT2

# ── Trạng thái cố định của chuỗi joint (prefix KHÔNG phụ thuộc bước) ────────
vae = AutoencoderKLQwenImage21()
Z_COND = vae.encode(IMG_DEMO)                            # (4, 4) latent điều kiện
Z_COND_STD = standardize(Z_COND)
COND_TOKENS = Z_COND_STD @ W_IMG_IN + CONTEXT[2]         # img_in + slot VLM
TXT_TOKENS = np.stack([txt_in(CONTEXT[i]) for i in (0, 1, 3)])
PREFIX_TOKENS = np.concatenate([TXT_TOKENS[:2], COND_TOKENS,
                                TXT_TOKENS[2:]], axis=0)  # (7, 8)
NOISE = _n((TARGET_LEN, LATENT_CHANNELS), 1.0)   # noise target ngày sigma=1
SCALE_OUT_BASE = _n((DIM,), 0.15)                # scale-not-học cho final AdaLN

REPEATS = [IMG_TOKENS_PER_SLOT if s else 1 for s in VLM_IMG_MASK]  # [1,1,4,1]
IMAGE_IDS = [-1, -1] + [0] * 4 + [-1] + [1] * 4
TARGET_TOKEN_MASK = [False] * PREFIX_LEN + [True] * TARGET_LEN

def run_pipeline(use_cache):
    """Chạy 3 bước Euler. use_cache=True: extract → cached; False: recompute.

Trả về velocities (3), x sau mỗi bước, RGBA cuối, cache và capture từng bước.
    """
    cache = QwenImage21KVCache(N_LAYERS)
    sched = FlowMatchEulerDiscreteScheduler(TURBO_SIGMAS)
    x = NOISE.copy()
    records = []
    while not sched.finished:
        t_step = sched.timestep()
        sigma = sched.sigma()
        if use_cache:
            mode = "extract" if sched.step_index == 0 else "cached"
        else:
            mode = "extract"          # recompute: tính cả 11 token mọi bước
            cache = QwenImage21KVCache(N_LAYERS) if False else cache
        cap = {}
        out = dtr.forward(x, t_step, cache if mode == "extract" else cache,
                          mode, capture=cap)
        records.append({"sigma": sigma, "t_step": t_step, "mode": mode,
                        "velocity": out["velocity"], "x_next": sched.step(
                            x, out["velocity"]), "capture": cap})
        x = records[-1]["x_next"]
    rgba = vae.decode(unstandardize(x))
    return {"records": records, "rgba": rgba, "cache": cache,
            "cache_snap": {"k": [c.k.shape for c in [None] * 0] or
                           [cache.k[l].shape for l in range(N_LAYERS)],
                           "v": [cache.v[l].shape for l in range(N_LAYERS)]}}

def fmt_vec(v, nd=3):
    return "[" + ", ".join(f"{x:.{nd}f}" for x in np.asarray(v, dtype=float)) + "]"

def say(code, title, lines):
    print(f"[{code}] {title}")
    for ln in lines:
        print("   " + ln)
    print()

def print_header():
    print("=== LAB QWEN-IMAGE-2.1: TOÀN TRÌNH 1A → 6B BẰNG NUMPY ===")
    print("MỤC TIÊU: chạy một lượt denoise 3 bước Euler từ noise đến ảnh RGBA,")
    print("đọc được từng khối 1A..6B và kiểm chứng prefix KV cache + block-causal.")
    print(f"ĐẦU VÀO: ảnh điều kiện {IMG_SHAPES[0]} (2×2, {RGBA_CHANNELS} kênh RGBA),")
    print(f"  noise target 2×2×{LATENT_CHANNELS}, chuỗi VLM (4 slot, slot ảnh mở rộng 4×),")
    print(f"  TURBO_SIGMAS = (1.0, 0.85, 0.45, 0.0).")
    print(f"CẤU HÌNH THU NHỎ: DIM={DIM} (prod 4096), heads={N_HEADS}×{HEAD_DIM},")
    print(f"  RoPE trục {AXES_DIMS_ROPE} (prod (16, 56, 56)), FFN ratio {MLP_RATIO}x "
          f"(MLP_DIM={MLP_DIM}),")
    print(f"  layers={N_LAYERS} (prod 32), latent {LATENT_CHANNELS} kênh (prod 64),")
    print(f"  S_joint={S_JOINT} (prefix {PREFIX_LEN} + target {TARGET_LEN}).")
    print(f"Weights chưa học, seed {SEED}. Các số kiểm chứng CƠ CHẾ — KHÔNG chứng minh")
    print("chất lượng ảnh hay tốc độ production.")
    print()

def print_sections(step1_cap, pipeline, recompute):
    # ── [1A] ────────────────────────────────────────────────────────────────
    say("1A", "VAE RGBA 64 KÊNH & UNPATCHED LATENTS (AutoencoderKLQwenImage21)", [
        f"Ảnh điều kiện demo: {RGBA_CHANNELS} pixel × R, G, B, A (giá trị cố định):",
        "   pixel 0: " + fmt_vec(IMG_DEMO.reshape(4, RGBA_CHANNELS)[0]),
        "   pixel 1: " + fmt_vec(IMG_DEMO.reshape(4, RGBA_CHANNELS)[1]),
        "   pixel 2: " + fmt_vec(IMG_DEMO.reshape(4, RGBA_CHANNELS)[2]),
        "   pixel 3: " + fmt_vec(IMG_DEMO.reshape(4, RGBA_CHANNELS)[3]),
        f"encode → latent (4, {LATENT_CHANNELS}) — prod dùng 64 kênh, lab dùng "
        f"{LATENT_CHANNELS} kênh",
        "  để giữ đúng vai trò mà không cần mạng conv thật.",
        f"Thống kê chuẩn hóa (mỗi kênh một giá trị): mean = {fmt_vec(LATENT_MEAN)},",
        f"  std = {fmt_vec(LATENT_STD)}  → chuẩn hóa: z = (z - mean) / std.",
        "latent sau chuẩn hóa (4 token × 4 kênh):",
    ] + [f"   token {i}: " + fmt_vec(Z_COND_STD[i]) for i in range(4)] + [
        f"img_in với PATCH_SIZE={PATCH_SIZE} (không chia ô): mỗi token (4,) → ({DIM},).",
        "   token 0 sau img_in: " + fmt_vec(Z_COND_STD[0] @ W_IMG_IN),
        "Kết luận: latent LUÔN chuẩn hóa trước img_in theo thống kê của VAE",
        "(prod đọc mean/std từ tokenizer 64 kênh); denoise chạy trong không",
        "gian chuẩn hóa — decode VAE chỉ xảy ra ở [6B] sau un-standardize.",
    ])

# ── [1B] ────────────────────────────────────────────────────────────────
    ctx0 = CONTEXT[0]
    ms0 = float(np.mean(ctx0 ** 2))
    say("1B", "NGỮ CẢNH QWEN3-VL & RMSNORM ZERO-CENTER (w+1)", [
        f"CONTEXT (4 token VLM × {DIM}): [text, text, COND_IMG-slot, text].",
        "   slot 0: " + fmt_vec(CONTEXT[0]),
        "   slot 1: " + fmt_vec(CONTEXT[1]),
        "   slot 2 (ảnh): " + fmt_vec(CONTEXT[2]),
        "   slot 3: " + fmt_vec(CONTEXT[3]),
        f"RMSNORM zero-center của slot 0: mean(x²) = {ms0:.4f},",
        f"  rsqrt(ms + {EPS_NORM:g}) = {1.0 / np.sqrt(ms0 + EPS_NORM):.4f};",
        "  (w + 1) = " + fmt_vec(1.0 + RMS_CTXW),
        "  kết quả: giữ hướng vector, chuẩn hóa RMS về 1 rồi nhân (w+1):",
        "   norm(slot 0): " + fmt_vec(rmsnorm(ctx0, RMS_CTXW)),
        "txt_in GELU MLP: hidden = GELU( RMSNorm(ctx)·W1 ) rồi ·W2:",
        "   slot 0 sau txt_in: " + fmt_vec(TXT_TOKENS[0]),
        "   slot 3 sau txt_in: " + fmt_vec(TXT_TOKENS[2]),
        "Kết luận: 3 token văn bản đi qua txt_in (MLP GELU); vector slot ảnh",
        "giữ nguyên để nhân bản 4× khi ghép với latents ở [2A].",
    ])

# ── [2A] ────────────────────────────────────────────────────────────────
    slot_emb = CONTEXT[2]
    capstep1 = pipeline["records"][0]["capture"]
    say("2A", "MỞ RỘNG SLOT ẢNH 4× & ĐÓNG GÓI CHUỖI JOINT 11 TOKEN", [
        f"repeats = {REPEATS} — token ảnh VLM (slot 2) được PHỔ ĐẠI 4 lần",
        f"  (1 slot → {IMG_TOKENS_PER_SLOT} token = lưới latent 2×2; 1 text slot giữ 1).",
        f"image_ids = [{', '.join(str(int(i)) for i in IMAGE_IDS)}] — bản đồ sau mở rộng:",
        "   -1 = văn bản, 0 = ảnh điều kiện, 1 = ảnh target; S_joint = 11.",
        f"prefix_len = {PREFIX_LEN} (2 text + 4 cond_img + 1 text), "
        f"target_len = {TARGET_LEN} (2×2).",
        f"Slot ảnh mở rộng CỘNG slot VLM nhân bản (slot_emb = {fmt_vec(slot_emb, 2)}):",
        "   cond token 0 = img_in(latent 0) + slot_emb = "
        + fmt_vec(COND_TOKENS[0]),
        f"target_token_mask = [{', '.join('T' if m else 'F' for m in TARGET_TOKEN_MASK)}]"
        f" — {sum(TARGET_TOKEN_MASK)}/{S_JOINT} token là ảnh target.",
        f"joint (11, {DIM}) = prefix (7, {DIM}) + target (4, {DIM}):",
        "   joint[0] (text): " + fmt_vec(capstep1["joint_input"][0]),
        "   joint[2] (cond img): " + fmt_vec(capstep1["joint_input"][2]),
        "   joint[7] (target 0): " + fmt_vec(capstep1["joint_input"][7]),
        "Kết luận: mở rộng slot KHÔNG đổi ý nghĩa — ảnh bung đúng bằng số",
        "token lưới latent; khối target thêm SAU prefix nên prefix không xô lệch.",
    ])

# ── [3A] ────────────────────────────────────────────────────────────────
    role_names = ["text 0", "text 1", "cond img", "cond img", "cond img",
                  "cond img", "text 2", "target img", "target img", "target img",
                  "target img"]
    lines = ["Tọa độ (frame, height, width) zero-centered & 2 cặp góc xoay:"]
    for i in range(S_JOINT):
        f, hh, ww = ROPE_COORDS[i]
        lines.append(f"   pos {i:2d} ({role_names[i]:11s}): "
                     f"(f={f:+.0f}, h={hh:+.0f}, w={ww:+.0f})  "
                     f"cặp0 = {ANGLE_PAIR[0, i]:.4f}, cặp1 = {ANGLE_PAIR[1, i]:.4f}")
    lines += [
        "Text TIẾN DẦN cả 3 trục: (0,0,0) → (1,1,1) → (2,2,2) cho 3 token chữ.",
        "Hai khối ảnh đều GHIM frame = 0, (height, width) căn giữa:",
        "   [(-1,-1), (-1,0), (0,-1), (0,0)] — lưới nguyên thu nhỏ của prod",
        "  (prod lấy lưới nửa chiều arange(d) - (d-1)/2 với d 56 và 56).",
        f"Micro chia sẻ cặp RoPE: cặp 0 cho frame ({ROPE_TF:g}), cặp 1 cho height",
        f"  ({ROPE_TH:.0e}) + width ({ROPE_TW:.0e}) cộng dồn; prod đủ 64 cặp như",
        "  bài qwen21-3d-axial-rope đã mô tả.",
        "Kết luận: RoPE 3D cho attention phân biệt văn bản (tiến dần theo trục)",
        "với ảnh (căn giữa lưới) — không cần embedding cứng khác.",
    ]
    say("3A", "ROPE TRỤC 3D (FRAME, HEIGHT, WIDTH) & LƯỚI ZERO-CENTERED", lines)

# ── [3B] ────────────────────────────────────────────────────────────────
    tab1 = modulation_table(1000.0 * TURBO_SIGMAS[0])
    tab2 = modulation_table(1000.0 * TURBO_SIGMAS[1])
    say("3B", "ĐIỀU KIỆN NHÂN QUẢ THEO t (BẢNG [t_step, 0.0]) & SHARED MODULATION", [
        f"Bảng modulation shape (2, {6 * DIM}); mỗi hàng 6 cột × DIM là phần đầu:",
        "Step 1 (t = 1000·1.0 = 1000) — 6 cột đầu của 2 hàng:",
        "   hàng 0 (t_step): " + fmt_vec(tab1[0][:6]),
        "   hàng 1 (t=0.0):  " + fmt_vec(tab1[1][:6]),
        "Step 2 (t = 850) — cùng 6 cột đầu:",
        "   hàng 0 (t_step): " + fmt_vec(tab2[0][:6]),
        "   hàng 1 (t=0.0):  " + fmt_vec(tab2[1][:6]),
        f"Hàng 1 (t=0.0) GIỐNG Y HỆT giữa hai bước: lệch tối đa "
        f"{np.max(np.abs(tab1[1] - tab2[1])):.2e}; hàng 0 đổi theo t:",
        f"  lệch chuẩn hóa tối đa {float(np.max(np.abs(tab1[0] - tab2[0]))):.4f}.",
        "Ai đọc hàng nào?",
        "   7 token prefix (0..6): hàng t = 0.0 (2 text + 4 cond_img + 1 text)",
        "   4 token target (7..10): hàng t_step của bước đang chạy.",
        "Cột trong 6 nhóm: shift1, scale1, gate1, shift2, scale2, gate2 —",
        "  dùng bởi các AdaLN của [4A] và [5A] (mod SHARED cho cả chuỗi).",
        "Kết luận: prefix nguyên ngày nhờ t=0.0 — chính là giấy phép cho KV",
        "cache ở [6A]: K/V prefix tính MỘT LẦN, dùng mọi bước sigma.",
    ])

# ── [3C] ────────────────────────────────────────────────────────────────
    lines = ["Mask block-causal 11×11 (1 = được phép, . = bị chặn):"]
    for q in range(S_JOINT):
        row = " ".join("1" if BLOCK_MASK[q][kv] else "." for kv in range(S_JOINT))
        lines.append(f"   q={q:2d} | {row} | {int(BLOCK_MASK[q].sum())} key")
    lines += [
        "Quy tắc: (q >= kv) HOẶC cùng khối ảnh; text thì nhân quả thuần.",
        "Text: q=0 chỉ thấy 1 key; q=6 thấy 0..6 (7 key) — không thấy tương lai.",
        "Khối ảnh 2..5: SONG HƯỚNG — nhìn thấy nhau bất kể thứ tự (khối 4×4)",
        "  gồm cả 2 token text phía trước (nhờ q >= kv) = 6 key mỗi hàng.",
        "Khối target 7..10: q >= kv cho cả prefix lẫn trong khối 7..10 = 11 key.",
        "Kết luận: ảnh 'đọc trọn khối' của mình, văn bản nhân quả, target thấy",
        "prefix — nền tảng cho prefix KV cache và kiểm chứng cuối.",
    ]
    say("3C", "MASK ATTENTION BLOCK-CAUSAL 11×11", lines)

# ── [4A] ────────────────────────────────────────────────────────────────
    cap = step1_cap
    say("4A", "ATTENTION DiT SINGLE-STREAM & QK-RMSNORM (lớp 0, mẫu q=7 head 0)", [
        "Modulated LayerNorm: x_mod = LN(x)·(1 + scale1) + shift1 — mỗi token",
        "  nhân HÀNG modulation của chính nó (target 7 đọc hàng t_step).",
        "   x_mod[7]: " + fmt_vec(cap["x_mod"][7]),
        f"Chiếu Q/K/V: ({DIM},) → ({N_HEADS}, {HEAD_DIM}) mỗi token (single stream",
        "  CHUNG cho text và ảnh — khác double-stream của thế hệ trước).",
        "Q token 7 head 0: chưa RMSNorm    : " + fmt_vec(cap["q_raw"][0, 7]),
        "                  sau QK-RMSNORM  : " + fmt_vec(cap["q_rms"][0, 7]),
        "                  sau RoPE 3D     : " + fmt_vec(cap["q_rot"][0, 7]),
        "K token 2 (khối ảnh) head 0: sau QK-RMSNORM " + fmt_vec(cap["k_rms"][0, 2]),
        "   sau RoPE: " + fmt_vec(cap["k_rot"][0, 2]),
        "Attention weights token 7 head 0 trên 11 key (mask cho phép cả 11):",
        "   " + fmt_vec(cap["weights"][0, 7]),
        "Residual: x = x + tanh(gate1) · W_O(attn) — gate1 học sẵn, tanh (giữ",
        f"  biên độ ≤ 1); gate1 lớp 0: " + fmt_vec(cap["gate1"][0], 2),
        "Kết luận: chuỗi 11 token đi một dòng (single-stream); QK-RMSNORM giữ",
        "biên độ, RoPE 3D gán vị trí theo trục, cổng tanh bơm vào residual.",
    ])

def fmt_row(v, n=6, nd=3):
    return fmt_vec(np.asarray(v, dtype=float)[:n], nd)

def print_sections_b(step1_cap, pipeline, recompute):
    cap = step1_cap  # capture lớp 0 của Step 1
    # ── [5A] (tiếp) ─────────────────────────────────────────────────────────
    say("5A", f"FFN SWIGLU ĐIỀU BIẾN (RATIO {MLP_RATIO}×) & CỔNG TANH (token 7, lớp 0)", [
        "Trước FFN: ffn_in = LN(x)·(1 + scale2) + shift2 — cột 3..5 của hàng mod.",
        "   ffn_in[7]: " + fmt_vec(cap["ffn_in"][7]),
        f"W_gate, W_up: ({DIM},) → ({MLP_DIM},) — MLP_RATIO = {MLP_RATIO}× chiều input.",
        "SiLU(W_gate·x) 6 cột đầu: " + fmt_row(silu(cap["ffn_g"][7])),
        "W_up·x         6 cột đầu: " + fmt_row(cap["ffn_u"][7]),
        "hidden = silu(gate)·up (SwiGLU) 6 cột đầu: " + fmt_row(cap["ffn_hidden"][7]),
        f"W_down: ({MLP_DIM},) → ({DIM},) rồi residual: x = x + tanh(gate2)·MLP.",
        f"gate2 lớp 0: " + fmt_vec(cap["gate2"][0], 2) + " → tanh = "
        + fmt_vec(np.tanh(cap["gate2"][0]), 2),
        f"‖MLP token 7‖ = {np.linalg.norm(cap['ffn_hidden'][7] @ W_SWIGLU_DOWN[0]):.3f}",
        "Kết luận: 2 lớp micro × 2 cổng tanh; đúng cơ chế modulated SwiGLU đã",
        "học ở bài kiến trúc — ở FFN không hề có softmax.",
    ])

# ── [6A] ────────────────────────────────────────────────────────────────
    cached_cap = pipeline["records"][1]["capture"]["cached_rep"]
    mode_lines = [
        "Bảng bước: sigma 1.00 → t=1000, mode=extract | sigma 0.85 → t=850, mode=cached",
        "             sigma 0.45 → t=450, mode=cached | sigma 0.00 — kết thúc.",
    ]
    say("6A", "PREFIX KV CACHE (STEP 1 EXTRACT → STEPS 2-3 CACHED)", mode_lines + [
        "Step 1 (sigma=1.0) kv_cache_mode='extract': tính cả 11 token, GHI:",
    ] + [f"   lớp {l}: cached_k shape "
         f"{pipeline['cache_snap']['k'][l]} · cached_v shape "
         f"{pipeline['cache_snap']['v'][l]} — (7 prefix, {N_HEADS} head, {HEAD_DIM})"
         for l in range(N_LAYERS)] + [
        "Steps 2-3 (sigma=0.85, 0.45) kv_cache_mode='cached': chiếu Q/K/V CHỈ cho",
        f"  4 token target (joint_hidden_states[7:]): Q/K/V (4, {N_HEADS}, {HEAD_DIM}).",
        f"  cat([cached_k (7, 2, 4), target_k ({TARGET_LEN}, {N_HEADS}, {HEAD_DIM})])",
        f"  → k_full {cached_cap['k_full_shape']} — cách tập hợp tương tự cho V.",
        "Recompute toàn trình (đối chiếu ở [KIỂM TRA & THỬ NGHIỆM] (1)): 3 bước",
        "  đều chạy full 11 token với mask block-causal — kết quả như cached.",
        "Kết luận: prefix K/V tính 1 lần dùng 3 lần; mỗi bước chỉ phải tính",
        "target token — tiền đề của denoise batching nói trong bài Turbo.",
    ])

# ── [6B] ────────────────────────────────────────────────────────────────
    rec3 = pipeline["records"][2]
    lines = [
        f"Final AdaLN SCALE-ONLY: out = LN(h) · (1 + scale_final), bỏ shift —",
        f"  scale_final của Step 3 (t=450): {fmt_vec(rec3['capture']['scale_final'])}",
        "Velocity v_theta (4 token × 4 kênh chuẩn hóa) sau W_unpatch, Step 3:",
    ]
    for i in range(TARGET_LEN):
        lines.append(f"   token {i}: " + fmt_vec(rec3["velocity"][i]))
    lines.append("Vòng Euler FlowMatch: x_{k+1} = x_k + (sigma_{k+1} − sigma_k)·v θ")
    lines.append("   sigma (1.0, 0.85, 0.45, 0.0) → delta (−0.15, −0.40, −0.45); x đầу = noise.")
    for k in range(3):
        xk = pipeline["records"][k]["x_next"]
        lines.append(f"   sau Step {k + 1} (sigma {TURBO_SIGMAS[k + 1]}): "
                     f"‖x‖ = {np.linalg.norm(xk):.4f}, mode="
                     f"{pipeline['records'][k]['mode']}")
    x_last = pipeline["records"][2]["x_next"]
    z_raw = unstandardize(x_last)
    lines += [
        "un-standardize: z = x · std + mean (4 token × 4 kênh):",
    ] + [f"   token {i}: " + fmt_vec(z_raw[i]) for i in range(TARGET_LEN)] + [
        "decode VAE (tanh về [-1, 1]) — ảnh RGBA 2×2 × (R, G, B, Alpha):",
        "   hàng 0: R " + fmt_vec(pipeline["rgba"][0, :, 0]) + " | G "
        + fmt_vec(pipeline["rgba"][0, :, 1]),
        "           B " + fmt_vec(pipeline["rgba"][0, :, 2]) + " | Alpha "
        + fmt_vec(pipeline["rgba"][0, :, 3]),
        "   hàng 1: R " + fmt_vec(pipeline["rgba"][1, :, 0]) + " | G "
        + fmt_vec(pipeline["rgba"][1, :, 1]),
        "           B " + fmt_vec(pipeline["rgba"][1, :, 2]) + " | Alpha "
        + fmt_vec(pipeline["rgba"][1, :, 3]),
        "Kết luận: Euler đi trong KHÔNG GIAN CHUẨN HÓA; decode chỉ là bước ra.",
        "Weights chưa học → ảnh không có ý nghĩa trực quan, chỉ kiểm wiring.",
    ]
    say("6B", "FINAL ADALN SCALE-ONLY, VELOCITY & VÒNG EULER → RGBA", lines)

# ── [KIỂM TRA & THỬ NGHIỆM] ────────────────────────────────────────────────
def print_checks(pipeline, recompute):
    # (1) Prefix KV cache: cached và recompute phải cho velocity & RGBA như nhau.
    print("[KIỂM TRA & THỬ NGHIỆM]")
    for step in (1, 2):
        diff = float(np.max(np.abs(pipeline["records"][step]["velocity"]
                                   - recompute["records"][step]["velocity"])))
        print(f"  (1) Velocity bước {step + 1} (sigma {TURBO_SIGMAS[step + 1]}): cached vs"
              f" recompute lệch {diff:.2e} (ngưỡng 1e-9).")
        assert diff < 1e-9
    rgba_diff = float(np.max(np.abs(pipeline["rgba"] - recompute["rgba"])))
    print(f"      RGBA cuối: lệch {rgba_diff:.2e} (ngưỡng 1e-9).")
    assert rgba_diff < 1e-9
    k_cache = pipeline["cache"].k[0]                          # ghi ở Step 1
    k_recomp = recompute["records"][1]["capture"]["k_full_step"]
    k_prefix_diff = float(np.max(np.abs(k_cache - k_recomp[:PREFIX_LEN])))
    print(f"      Prefix K lớp 0: cache Step 1 vs recompute Step 2 lệch "
          f"{k_prefix_diff:.2e} (ngưỡng 1e-9).")
    assert k_prefix_diff < 1e-9
    print("      → PASS: cache prefix K/V (7, 2, 4) tái tạo chính xác attention.")

# (2) Vì sao cần causal_condition = True (t = 0.0 trên prefix)?
    cache_bad = QwenImage21KVCache(N_LAYERS)
    step1_bad = dtr.forward(NOISE, 1000.0, cache_bad, "extract", causal_condition=False)
    x1_bad = NOISE + (float(TURBO_SIGMAS[1]) - float(TURBO_SIGMAS[0])) * step1_bad["velocity"]
    cache_bad2 = QwenImage21KVCache(N_LAYERS)
    dtr.forward(x1_bad, 850.0, cache_bad2, "extract", causal_condition=False)
    for l in range(N_LAYERS):
        drift = float(np.max(np.abs(cache_bad.k[l] - cache_bad2.k[l])))
        print(f"  (2) Nếu prefix ĐỌC t_step: cached K lớp {l} của Step 1 (t=1000) vs"
              f" recomputed K prefix Step 2 (t=850) lệch {drift:.2e} — DRIFT.")
        assert drift > 1e-9
    print("      → PASS (đối chiếu ngược): t = 0.0 trên prefix là điều kiện bắt")
    print("         buộc để cache không 'trôi' giữa các bước sigma.")

# (3) Block-causal isolation: đổi token target không rỉ sang prefix.
    base = dtr.forward(NOISE, 1000.0, None, "extract")
    base_h = base["joint_final"]
    per_tok = dtr.forward(NOISE, 1000.0, None, "extract",
                          perturb=[(PREFIX_LEN, 0, 0.5)])        # token 7 (target 0)
    d_target = per_tok["joint_final"] - base_h
    pref_leak = float(np.max(np.abs(d_target[:PREFIX_LEN])))
    print(f"  (3) Perturb token target 7: 7 token prefix lệch {pref_leak:.2e} "
          f"(ngưỡng==0); 4 token target lệch tới "
          f"{float(np.max(np.abs(d_target[PREFIX_LEN:]))):.2e}.")
    assert pref_leak == 0.0
    per_tok2 = dtr.forward(NOISE, 1000.0, None, "extract", perturb=[(5, 0, 0.5)])
    d2 = np.abs(per_tok2["joint_final"] - base_h).max(axis=1)
    print(f"      Perturb ảnh điều kiện idx 5: khối ảnh 2..5 lệch "
          f"[{d2[2]:.2e}, {d2[3]:.2e}, {d2[4]:.2e}, {d2[5]:.2e}] — tất cả > 0 (song hướng).")
    print(f"      Token văn bản 0, 1 đứng tuyệt đối: {d2[0]:.2e}, {d2[1]:.2e} "
          f"(nhân quả chặn).")
    assert np.all(d2[2:6] > 0)
    assert d2[0] == 0.0 and d2[1] == 0.0
    print("      → PASS: khối ảnh song hướng, văn bản nhân quả, prefix không bao")
    print("         giờ bị target 'làm bẩn' — nền móng KV cache prefix.")
    print("PASS: toàn trình 1A→6B chạy đúng; các số trên mới kiểm chứng CƠ CHẾ")
    print("đã khai báo, không chứng minh chất lượng ảnh hay tốc độ production.")

# ── Demo chạy khi thực thi (trình duyệt/CLI/test đều dùng chung) ────────────
print_header()

dtr = QwenImage21Transformer2DModel()

# 1) Step 1 (extract): capture toàn bộ chi tiết dùng cho in 1A..5A.
step1_cap = {}
dtr.forward(NOISE, 1000.0, QwenImage21KVCache(N_LAYERS), "extract", capture=step1_cap)

# 2) Pipeline CACHED (Step 1 extract → Steps 2-3 cached) và pipeline RECOMPUTE.
pipeline = run_pipeline(use_cache=True)
recompute = run_pipeline(use_cache=False)
assert pipeline["cache"].k[0].shape == (7, N_HEADS, HEAD_DIM)
assert [r["mode"] for r in pipeline["records"]] == ["extract", "cached", "cached"]

# 3) In toàn bộ câu chuyện 1A..6B bằng số thật.
print_sections(step1_cap, pipeline, recompute)
print_sections_b(step1_cap, pipeline, recompute)

# 4) Kiểm chứng số học.
print_checks(pipeline, recompute)
print()
print("TỔNG KẾT:")
print(f"1. Latent {LATENT_CHANNELS} kênh (prod 64) chuẩn hóa (z - mean)/std trước img_in (1A).")
print(f"2. Slot ảnh VLM nhân bản {IMG_TOKENS_PER_SLOT} lần; chuỗi joint {S_JOINT} token (2A).")
print("3. RoPE 3D trục; prefix đọc hàng t=0.0 của bảng modulation — đổi t không xô cache (3A/3B).")
print("4. Block-causal 11×11: text nhân quả, khối ảnh song hướng (3C).")
print(f"5. KV cache prefix (7, {N_HEADS}, {HEAD_DIM}) tái tạo chính xác; Euler 3 bước cho RGBA 2×2 (6A/6B).")

Python outputSẵn sàng.

Nguồn công bố & mã nguồn tham chiếu

Qwen/Qwen-Image-2.1 & Qwen-Image-2.1-Turbo (Hugging Face) ↗ Model card, transformer/config.json, vae/config.json & 8-step Turbo sigmas transformer_qwenimage21.py (HuggingFace Diffusers) ↗ Mã nguồn QwenImage21Transformer2DModel, QwenImage21Rope, BlockMask & Prefix KV Cache QwenLM/Qwen-Image-2.1 (GitHub) ↗ Kho mã nguồn chính thức, hướng dẫn Text-to-Image, Multi-Image Editing & RGBA

Danh sách bài học (11)