nodes/Nynxz-ComfyUI-AsymFLUX2
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Support image editing (FLUX-Kontext-style reference conditioning)

Browse Source 
- surgery.py: in the wrapped forward, divide ref_latents by `s` (the
  AsymFlow scale_buffer, ~2.4375) before handing them to comfy's flux2
  `_forward`. This matches upstream's `condition_hidden_states / s` step
  in `_AsymFlux2Transformer2DModel.forward` — without this, reference
  tokens reach `x_embedder` at ~2.4x the magnitude the model trained
  with and editing produces jank.
- example_workflows/asymflux2_edit.json: end-to-end edit workflow using
  comfy's stock `ReferenceLatent` node. Encoded source image attaches to
  positive conditioning; KSampler.latent_image stays as empty pixel
  latent (sampling from pure noise, not img2img).
- README: editing section with the wiring diagram and the two common
  mistakes (don't pipe encoded latent into KSampler.latent_image; don't
  attach reference to the negative branch).

Comfy's flux2 `_forward` already supports ref_latents and uses
`ref_index_scale=10.0` which matches upstream's
`_prepare_condition_latent_ids(scale=10)` for position IDs — no other
wiring needed on our side.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
This commit is contained in:
Nynxz
2026-05-14 22:29:04 +10:00
co-authored by Claude Opus 4.7
parent 789526f05e
commit 922550be01
3 changed files with 280 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files
README.md
+23 -2
View File
@@ -20,8 +20,29 @@ CLIPLoader (Qwen3 8B, type=flux2) ─► CLIPTextEncode (pos / neg) ───
AsymFLUX2 Empty Pixel Latent ───────────────────────┘
```
A working workflow lives at
[`example_workflows/asymflux2_t2i.json`](example_workflows/asymflux2_t2i.json).
Working workflows:
- [`example_workflows/asymflux2_t2i.json`](example_workflows/asymflux2_t2i.json) — text-to-image.
- [`example_workflows/asymflux2_edit.json`](example_workflows/asymflux2_edit.json) — image editing (FLUX-Kontext-style: a reference image conditions the generation, sampling still starts from pure noise).
## Image editing
```
LoadImage ─► AsymFLUX2 Oklab Encode ─► ReferenceLatent ─► (positive) CLIPTextEncode ─► KSampler.positive
▲
│ (positive CONDITIONING from CLIPTextEncode also goes in here)
(negative) CLIPTextEncode ────────────────────────────────────────► KSampler.negative
AsymFLUX2 Empty Pixel Latent ───────────────────────────────────── ► KSampler.latent_image
```
Two things to know:
- **Use `ReferenceLatent` (ships with comfy** at `advanced/conditioning/edit_models/`**).** Your encoded image goes into the `latent` socket; your `CLIPTextEncode` positive output goes into the `conditioning` socket; the result is a CONDITIONING with the reference attached. Wire that into KSampler.positive.
- **`KSampler.latent_image` stays as the empty pixel latent.** This is *not* img2img — sampling starts from pure noise, conditioned on (text + reference image). The reference image flows in through the conditioning, never through the latent_image socket.
- **Don't put the reference on the negative branch.** Only positive.
Behind the scenes our forward wrapper divides the reference tokens by ``s`` (the AsymFlow scale_buffer, ~2.4375) before they reach `x_embedder`, matching the upstream `_AsymFlux2Transformer2DModel.forward`. ComfyUI's stock flux2 path then handles patchifying + index-offset position IDs (`ref_index_scale=10` matches upstream's `_prepare_condition_latent_ids(scale=10)` for free).
## Nodes
asymflux2/model/surgery.py
+16 -1
View File
@@ -124,9 +124,24 @@ def _make_asymflux2_forward(
k_x = cal.k.reshape(-1, 1, 1, 1).to(dtype=x.dtype, device=x.device)
x_scaled = x * k_x
cal_t = cal.timestep.to(dtype=timestep.dtype, device=timestep.device)
# Image-editing path: upstream divides reference tokens by `s`
# before they go through x_embedder. We do the equivalent in image
# space (linearity of patchify makes pre-scaling commutative with
# the rearrange). Comfy's flux2 `_forward` will then patchify and
# concat ref tokens at index-offset positions (`ref_index_scale=10`
# matches upstream's `_prepare_condition_latent_ids(scale=10)`).
if ref_latents:
s_val = float(s_dev.detach().to(torch.float32).item())
ref_latents_scaled = [
r.to(device=x.device, dtype=x.dtype) / s_val for r in ref_latents
]
else:
ref_latents_scaled = ref_latents
u_a = original_forward(
x_scaled, cal_t, context,
y=y, guidance=guidance, ref_latents=ref_latents,
y=y, guidance=guidance, ref_latents=ref_latents_scaled,
control=control, transformer_options=transformer_options, **kwargs,
)
v = asymflow_velocity(
example_workflows/asymflux2_edit.json
+241
View File
@@ -0,0 +1,241 @@
{
"last_node_id": 12,
"last_link_id": 13,
"nodes": [
{
"id": 1,
"type": "UNETLoader",
"pos": [40, 40],
"size": [320, 90],
"flags": {},
"order": 0,
"mode": 0,
"inputs": [],
"outputs": [
{ "name": "MODEL", "type": "MODEL", "links": [1] }
],
"properties": { "Node name for S&R": "UNETLoader" },
"widgets_values": ["flux2-klein-base-9b.safetensors", "default"]
},
{
"id": 2,
"type": "asymflux2.ApplyAdapter",
"pos": [400, 40],
"size": [380, 200],
"flags": {},
"order": 4,
"mode": 0,
"inputs": [
{ "name": "model", "type": "MODEL", "link": 1 }
],
"outputs": [
{ "name": "model", "type": "MODEL", "links": [2] }
],
"properties": {},
"widgets_values": ["asymflux2-klein-9b.safetensors", 17.0, 1.0, 1.0]
},
{
"id": 3,
"type": "CLIPLoader",
"pos": [40, 180],
"size": [320, 110],
"flags": {},
"order": 1,
"mode": 0,
"inputs": [],
"outputs": [
{ "name": "CLIP", "type": "CLIP", "links": [3, 8] }
],
"properties": {},
"widgets_values": ["qwen3_8b.safetensors", "flux2", "default"]
},
{
"id": 10,
"type": "LoadImage",
"pos": [40, 340],
"size": [320, 320],
"flags": {},
"order": 2,
"mode": 0,
"inputs": [],
"outputs": [
{ "name": "IMAGE", "type": "IMAGE", "links": [10] },
{ "name": "MASK", "type": "MASK", "links": null }
],
"properties": {},
"widgets_values": ["source.png", "image"]
},
{
"id": 11,
"type": "asymflux2.OklabEncode",
"pos": [400, 280],
"size": [320, 200],
"flags": {},
"order": 5,
"mode": 0,
"inputs": [
{ "name": "image", "type": "IMAGE", "link": 10 }
],
"outputs": [
{ "name": "latent", "type": "LATENT", "links": [11] }
],
"properties": {},
"widgets_values": [0.56, 0.0, 0.01, 0.16]
},
{
"id": 4,
"type": "CLIPTextEncode",
"pos": [400, 510],
"size": [380, 200],
"flags": {},
"order": 6,
"mode": 0,
"inputs": [
{ "name": "clip", "type": "CLIP", "link": 3 }
],
"outputs": [
{ "name": "CONDITIONING", "type": "CONDITIONING", "links": [12] }
],
"properties": {},
"widgets_values": ["The same scene at golden hour, warm orange sunset light, long shadows, cinematic atmosphere."]
},
{
"id": 12,
"type": "ReferenceLatent",
"pos": [800, 510],
"size": [320, 90],
"flags": {},
"order": 7,
"mode": 0,
"inputs": [
{ "name": "conditioning", "type": "CONDITIONING", "link": 12 },
{ "name": "latent", "type": "LATENT", "link": 11 }
],
"outputs": [
{ "name": "CONDITIONING", "type": "CONDITIONING", "links": [4] }
],
"properties": {}
},
{
"id": 5,
"type": "CLIPTextEncode",
"pos": [400, 730],
"size": [380, 130],
"flags": {},
"order": 8,
"mode": 0,
"inputs": [
{ "name": "clip", "type": "CLIP", "link": 8 }
],
"outputs": [
{ "name": "CONDITIONING", "type": "CONDITIONING", "links": [5] }
],
"properties": {},
"widgets_values": ["Low quality, worst quality, blurry, deformed, bad anatomy, unclear text"]
},
{
"id": 6,
"type": "asymflux2.EmptyPixelLatent",
"pos": [800, 640],
"size": [320, 110],
"flags": {},
"order": 3,
"mode": 0,
"inputs": [],
"outputs": [
{ "name": "latent", "type": "LATENT", "links": [6] }
],
"properties": {},
"widgets_values": [960, 1280, 1]
},
{
"id": 7,
"type": "KSampler",
"pos": [1160, 40],
"size": [330, 480],
"flags": {},
"order": 9,
"mode": 0,
"inputs": [
{ "name": "model", "type": "MODEL", "link": 2 },
{ "name": "positive", "type": "CONDITIONING", "link": 4 },
{ "name": "negative", "type": "CONDITIONING", "link": 5 },
{ "name": "latent_image", "type": "LATENT", "link": 6 }
],
"outputs": [
{ "name": "LATENT", "type": "LATENT", "links": [7] }
],
"properties": {},
"widgets_values": [42, "fixed", 38, 4.0, "dpmpp_2m", "simple", 1.0]
},
{
"id": 8,
"type": "asymflux2.OklabDecode",
"pos": [1510, 40],
"size": [320, 200],
"flags": {},
"order": 10,
"mode": 0,
"inputs": [
{ "name": "latent", "type": "LATENT", "link": 7 }
],
"outputs": [
{ "name": "image", "type": "IMAGE", "links": [9] }
],
"properties": {},
"widgets_values": [0.56, 0.0, 0.01, 0.16]
},
{
"id": 9,
"type": "SaveImage",
"pos": [1850, 40],
"size": [320, 290],
"flags": {},
"order": 11,
"mode": 0,
"inputs": [
{ "name": "images", "type": "IMAGE", "link": 9 }
],
"outputs": [],
"properties": {},
"widgets_values": ["AsymFLUX2_edit"]
}
],
"links": [
[1, 1, 0, 2, 0, "MODEL"],
[2, 2, 0, 7, 0, "MODEL"],
[3, 3, 0, 4, 0, "CLIP"],
[4, 12, 0, 7, 1, "CONDITIONING"],
[5, 5, 0, 7, 2, "CONDITIONING"],
[6, 6, 0, 7, 3, "LATENT"],
[7, 7, 0, 8, 0, "LATENT"],
[8, 3, 0, 5, 0, "CLIP"],
[9, 8, 0, 9, 0, "IMAGE"],
[10, 10, 0, 11, 0, "IMAGE"],
[11, 11, 0, 12, 1, "LATENT"],
[12, 4, 0, 12, 0, "CONDITIONING"]
],
"groups": [
{
"title": "Model + adapter",
"bounding": [20, 0, 780, 250],
"color": "#3f789e",
"font_size": 16
},
{
"title": "Reference image (the photo to edit)",
"bounding": [380, 260, 760, 230],
"color": "#a1309b",
"font_size": 16
},
{
"title": "Conditioning",
"bounding": [380, 490, 760, 380],
"color": "#88a850",
"font_size": 16
}
],
"config": {},
"extra": {},
"version": 0.4
}