sgl-project/sglang · error · RuntimeError

SANA-WM refiner text encoder must return per-layer hidden_st

Error message

SANA-WM refiner text encoder must return per-layer hidden_states.

What it means

_encode_prompt calls the text encoder with output_hidden_states=True and then reads outputs.hidden_states. If the encoder's returned object lacks per-layer hidden_states (attribute None or absent), the refiner cannot build the stacked (B, L, D, n_layers) embedding tensor the SANA-WM refiner consumes, so it raises this RuntimeError.

Source

Thrown at python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/sana_wm/refiner.py:520

        )
        input_ids = text_inputs.input_ids.to(device)
        attention_mask = text_inputs.attention_mask.to(device)

        # Diffusers-backed official path loads HF Gemma3ForConditionalGeneration.
        # NVlabs encodes through `.model`; the fallback SGLang-native encoder is
        # still callable directly, so keep both surfaces.
        with self.use_declared_component(
            component_name="text_encoder_2", module=self.text_encoder
        ):
            text_backbone = getattr(self.text_encoder, "model", self.text_encoder)
            outputs = text_backbone(
                input_ids=input_ids,
                attention_mask=attention_mask,
                output_hidden_states=True,
            )
        per_layer_hidden = getattr(outputs, "hidden_states", None)
        if per_layer_hidden is None:
            raise RuntimeError(
                "SANA-WM refiner text encoder must return per-layer hidden_states."
            )
        stacked = torch.stack(per_layer_hidden, dim=-1)  # (B, L, D, n_layers)
        seq_lengths = attention_mask.sum(dim=-1)
        log_sana_wm_tensor_stats("refiner.text_hidden_states_stacked", stacked)
        prompt_embeds = _pack_text_embeds(
            stacked,
            seq_lengths,
            padding_side=tokenizer.padding_side,
        ).to(dtype=self.dtype)
        log_sana_wm_tensor_stats("refiner.prompt_embeds_packed", prompt_embeds)

        with self.use_declared_component(
            component_name="connectors", module=self.connectors
        ):
            video_text_embedding, _, video_attention_mask = self.connectors(
                prompt_embeds, attention_mask
            )

View on GitHub (pinned to 0132848349)

Solutions

1. Use the supported text encoder class for the SANA-WM pipeline
2. If wrapping the encoder, return an object with .hidden_states populated (pass output_hidden_states=True through and forward the tuple)
3. Check the wrapper returns the encoder's native output object rather than a reduced one

Example fix

# before
class Wrapper(nn.Module):
    def forward(self, **kw):
        out = self.enc(**kw)
        return out.last_hidden_state  # hidden_states lost
# after
class Wrapper(nn.Module):
    def forward(self, **kw):
        kw["output_hidden_states"] = True
        return self.enc(**kw)  # exposes .hidden_states

out = encoder(input_ids=ids, attention_mask=mask, output_hidden_states=True)
assert getattr(out, "hidden_states", None) is not None, "encoder must expose per-layer hidden_states"

def exposes_hidden_states(enc) -> bool:
    import inspect
    out = enc(input_ids=torch.zeros(1, 4, dtype=torch.long), attention_mask=torch.ones(1, 4, dtype=torch.long), output_hidden_states=True)
    return getattr(out, "hidden_states", None) is not None

try:
    return stage._encode_prompt(ids, mask)
except RuntimeError as e:
    if "hidden_states" in str(e):
        raise TypeError(f"encoder {type(encoder).__name__} incompatible: {e}")
    raise

When it happens

Trigger: Configuring the refiner with a text encoder whose forward signature/return type does not expose hidden_states (custom encoder wrapper, wrong model class, diffusers version returning a different output dataclass, or a wrapper that drops the field).

Common situations: Swapping in a different text encoder via component overrides; upgrading/downgrading transformers where model output objects changed; a custom wrapper that returns only last_hidden_state.

Related errors

• Unsupported text encoder output: expected `hidden_states`.
• SANA-WM stage-1 expects exactly one Gemma-2 text encoder.
• missing captured KV on {attr}
• SANA-WM height/width must be divisible by the LTX-2 spatial
• SD3 CLIP postprocessing requires hidden_states from encoder

AI-assisted analysis of sgl-project/sglang@0132848349 (2026-08-28). Data as JSON: /api/errors/3977e9a48ca539c5. Report an issue: GitHub.