fix(diffusion): re-inject clean first-frame image during OmniDreams context-noise KV cache write

[Cerdore]

Problem

The SGLang OmniDreams port produces blurry video where only the first frame
is sharp and every generated frame after it is blurry. The VAE decoder is
fine -- the bug is in the autoregressive denoising rollout.

Root Cause

At AR step 0, the context-noise forward (the final self.transformer() call
that writes the authoritative K/V into the rolling KV cache) feeds a
noise-corrupted first-frame latent to the network. After denoising ->
postprocess_clean_latent (which re-injects the clean VAE-encoded first-frame
image at frame 0), the code calls scheduler.add_noise(clean, 128) which
corrupts ALL frames at sigma ~0.13. This noisy latent is then fed directly
to the transformer for the KV-cache write -- the first frame's K/V in the
cache are noise-corrupted.

In the FlashDreams reference, _maybe_inject_image() fires INSIDE
_predict_branch during finalize_kv_cache -> predict_flow, re-injecting the
clean first-frame VAE latent AFTER add_noise but BEFORE the network forward
that writes the final KV cache. This means FlashDreams stores CLEAN frame-0
K/V, while SGLang stores NOISY frame-0 K/V.

When subsequent AR chunks attend to the first frame's K/V, they see
corrupted keys/values -- a distribution shift that the model was not
trained to handle, which accumulates into progressively blurry video
across AR chunks.

Fix

One if pin: block (2 lines) inserted after scheduler.add_noise() and before
self.transformer() in the context-noise KV cache write path.

File: python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/omnidreams.py

Diff:

         ctx_latent = scheduler.add_noise(
             clean,
             ctx_noise_t,
             rng=gen,
         )

•        if pin:
  

•            ctx_latent = ctx_latent * (1.0 - inject_mask) + image_full * inject_mask
         self.transformer(
             hidden_states=ctx_latent,
  

For AR step >= 1, pin is False, so the new block is a no-op -- matching
FlashDreams where _maybe_inject_image short-circuits on
cache.autoregressive_index != 0.

---

CI States

Latest PR Test (Base): ❌ Run #27461916594
Latest PR Test (Extra): ❌ Run #27461916363

[Cerdore]

以下是 AgainstEntropy 6 条意见在当前代码上的状态：

意见
当前状态
是否需要处理
1
清理
SKILL.md/CLAUDE.md
等文件
.claude/ 目录下的文件已正式提交到仓库，是项目标准结构
不需要 — 但如果要合入 main，可能仍需确认是否要包含
2
复用 LayerNormScaleShift
FinalLayer 中仍在手动做 layer_norm(x) * (1+scale) + shift，而其他 Block 已经用了 LayerNormScaleShift
仍存在 — 建议统一
3
用 MergedColumnParallelLinear 合并 Q/K/V
仍然是 3 个独立的 _tp_col_linear 调用
仍存在 — 可优化
4
用 overlay 替代重写 loader
自定义 loader 是因为 OmniDreams 用非 Diffusers 格式的 flat .pt checkpoint，与 Hunyuan3D 做法一致
不需要 — 有充分理由
5
去掉冗余 TP helpers
仍有 _use_tp() 等带 try-except 的防御性函数
仍存在 — 如果分布式环境保证已初始化，可简化
6
复用已有 RoPE 实现
omnidreams_rope.py 实现的是 3D RoPE（T:H:W=44:42:42, NeoX-style），已复用底层 _apply_rotary_emb 原语，上层逻辑是 OmniDreams 特有的
不需要 — 分离合理
7
wanvae.py:112 “Is it validated?”
是一个已文档化的 TODO，说明当前用的是 Wan 2.1 默认值作为 fallback，待 GPU 验证
低优 — 需要 GPU 才能验证