AMSample
A whole second Flux sampler hiding in a custom node
- pipe
- attention_mask
- prompt_batch
- IMAGE
- IMAGE
AMSample is where the actual pixels happen in this pack - and it's the node that reminds you this isn't a normal ComfyUI workflow. Instead of handing off to a KSampler, it runs the entire Flux (or SD3) diffusers pipeline internally, feeding your attention mask into the transformer's joint attention. You're effectively inside a second, self-contained sampler that ComfyUI's scheduler and sampler ecosystem can't touch. That's a real limitation, not a quirk: whatever sampler choices you've grown to trust in native Flux nodes, they don't apply here.
Mechanically it's simple. You feed it the pipe from AMModelLoader, plus the matched attention_mask and prompt_batch from AttentionMask. It calls the pipeline with joint_attention_kwargs={'attention_mask': ...}, which is the whole point of the pack - the mask tells the model to ignore the padding that fills out long prompts, rather than treating it as meaningful text. On SD3 it also injects a blank negative prompt to match how AttentionMask built the batch.
Inputs a beginner actually sets:
- pipe, attention_mask, prompt_batch - the wiring triple. Get them all from the other two nodes in this pack; there's no other source for these types.
- guidance_scale - defaults to 3.5, which is Flux's classic guidance setting (a distilled-model guidance scale, not CFG). If you're generating people, the community consensus is to drop it to roughly 1.8–2.0; it's also the knob most likely to give you that same-face Flux look at the default.
- width / height - defaults 1360×768. Flux is happiest in the two-megapixel range, and going larger also helps offset the adherence loss you get at low guidance.
- steps - 50 default, and with the Simple scheduler that's plenty.
- seed - plain integer, 0 default.
- batch_size - up to 64, but on a 12B model batched sampling is how you run out of VRAM fastest.
Output: IMAGE. One small trap: the node declares two IMAGE outputs but the code returns a single image tensor, so treat the first socket as the real one and ignore the second. It's the kind of unfinished edge that tells you this is a research demo.
Install is the same shared pack install, and the README's clone line is stale - use this one:
cd ComfyUI/custom_nodes
git clone https://github.com/leeguandong/ComfyUI_FluxAttentionMask
Restart ComfyUI. No requirements.txt, so it relies on the diffusers already in your environment - and that's the biggest gotcha here. Mainline diffusers accepts attention_mask in joint_attention_kwargs and silently ignores it; the mask only takes effect with a diffusers build that includes the still-unmerged PR #10044. The tell is if your output looks identical whether you pad to 64 or 512 tokens. Other things to expect: FLUX.1-dev is gated and needs an HF login plus accepted license, and with CPU offload on, generation is slow enough that you'll appreciate the steps count.
Realistically, this pack exists to test one hypothesis - that a text-side attention mask meaningfully improves semantic alignment on long Flux prompts. That's a fine experiment to run once. But the wider community evidence runs the other way: past roughly 100 effective tokens, these models start to drift, so the "more tokens = better" premise is exactly what you should be skeptical of. Reach for AMSample if you want to test that for yourself. If you just want Flux images, the native sampler you already have is better maintained and faster.
Inputs (9)
| Name | Type | Default | Description |
|---|---|---|---|
| pipe | MODEL_PIPE | — | |
| attention_mask | ATTENTION_MASK | — | |
| prompt_batch | PROMPT_BATCH | — | |
| guidance_scale | FLOAT | 3.501–20 | — |
| width | INT | 1360 | — |
| height | INT | 768 | — |
| steps | INT | 50 | — |
| seed | INT | 00–2147483647 | — |
| batch_size | INT | 11–64 | — |
Outputs (2)
| Name | Type | Description |
|---|---|---|
| IMAGE | IMAGE | — |
| IMAGE | IMAGE | — |