Two streams: why one diffusion pass couldn't do virtual try-on

Mask the torso, prompt for a red saree, run Flux Fill. You get a red saree, and it's usually a decent one. It's also not the saree the brand shipped us, and for a catalogue that gap is the entire product. This is about the second conditioning stream we added to close it, and why the obvious single-stream version kept failing in ways that looked fine right up until someone on the brand side opened the file at full size.

Why one stream wasn't enough

Inpainting alone. Flux Fill takes an image, a mask, and a prompt, and fills the masked region with something that agrees with the surrounding pixels. That "something that agrees" is exactly the problem. The model is doing a plausibility job: given a torso, a pose, and the words "red saree", it produces a red saree. Nothing obliges it to produce the red saree, the one with a specific zari border and a specific print repeat, because nothing in the conditioning path is carrying that information. Text is a bad channel for texture. You can write two hundred words describing a weave and the model will still pull it toward whatever weave dominated its training data.

We ended up calling this Style Drift. The output matches the reference in colour and rough silhouette and loses everything else. A knit sweater comes back as a flat red shirt. A hand-blocked print comes back as a generic floral. Colour survives because colour is low frequency and cheap to get right; structure is where the money is and structure is what goes first.

Structural guidance alone. The mirror image fails too. If you condition strongly on the reference garment and let the model regenerate the region with a weak spatial anchor, you get a beautiful picture of the garment and a bad picture of a person wearing it. The drape ignores where the shoulder actually is. The hem lands at a plausible hem height rather than this model's hem height. Hands that were resting on a hip get repainted as fabric, which is the failure we named the Amputated Hand, and it's the single fastest way to get an image rejected by a brand.

So the two streams have a clean division of labour. Flux Fill owns where: the body, the pose, the lighting on the skin, the boundary between garment and not-garment. Flux Redux owns what: the weave, the print, the structure of the reference. Neither one can be asked to do the other's job, and I spent a couple of weeks trying before accepting that.

The pipeline

1. SAM2 and the bleeding zone

The mask is upstream of everything, and a bad mask poisons both streams at once. Threshold-based or parsing-based masks were not precise enough at the garment boundary, so we moved to SAM2 with a hierarchical prompt set and then dilated the result.

# mask_garment.py
"""Produce the inpainting mask for a source image.
 
SAM2 gives a tight garment boundary; morphological dilation turns that
boundary into a "bleeding zone" the sampler is allowed to repaint, so the
new fabric can meet skin without a visible seam.
"""
 
import cv2
import numpy as np
import torch
from sam2.build_sam import build_sam2
from sam2.sam2_image_predictor import SAM2ImagePredictor
 
GARMENT_PROMPTS = ["shirt", "tshirt", "top", "saree", "dress"]
 
# Bleeding zone radius in pixels, measured at the crop resolution the sampler
# will actually see. 5 for tight-fitting tops, up to 15 for loose drape.
DILATION_PX = 9
 
# Sanity floor: a mask covering less than this fraction of the frame usually
# means SAM2 latched onto a sleeve or a shadow rather than the garment.
MIN_MASK_FRACTION = 0.04
 
 
def build_predictor(checkpoint: str, config: str) -> SAM2ImagePredictor:
    """Load SAM2 once per worker; the encoder pass is the expensive part."""
    model = build_sam2(config, checkpoint, device="cuda")
    return SAM2ImagePredictor(model)
 
 
def segment_garment(predictor, image_bgr: np.ndarray, prompts=GARMENT_PROMPTS):
    """Run the hierarchical prompt set and keep the highest-scoring mask."""
    image_rgb = cv2.cvtColor(image_bgr, cv2.COLOR_BGR2RGB)
    predictor.set_image(image_rgb)
 
    best_mask, best_score, best_prompt = None, -1.0, None
    for prompt in prompts:
        masks, scores, _ = predictor.predict(text=prompt, multimask_output=True)
        top = int(np.argmax(scores))
        if scores[top] > best_score:
            best_mask, best_score, best_prompt = masks[top], float(scores[top]), prompt
 
    return best_mask.astype(np.uint8), best_score, best_prompt
 
 
def dilate(mask: np.ndarray, radius_px: int = DILATION_PX) -> np.ndarray:
    """Grow the mask by radius_px with an elliptical structuring element."""
    k = 2 * radius_px + 1
    element = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (k, k))
    return cv2.dilate(mask, element, iterations=1)
 
 
def prepare_mask(predictor, image_bgr, radius_px=DILATION_PX):
    raw, score, prompt = segment_garment(predictor, image_bgr)
    coverage = float(raw.mean())
    if coverage < MIN_MASK_FRACTION:
        raise ValueError(f"mask too small for '{prompt}': coverage={coverage:.3f}")
 
    final = dilate(raw, radius_px)
    grown = float(final.mean()) - coverage
    print(f"prompt={prompt} sam2_score={score:.3f} coverage={coverage:.3f} "
          f"radius={radius_px}px grown_by={grown:.3f}")
    return final

The dilation radius looks like a nuisance parameter and isn't. Set it too tight and the sampler never sees a skin pixel adjacent to the garment, so it has no evidence about the lighting on the body and paints an edge that reads as a cutout pasted on. Set it too wide and you erase the evidence you actually need: the shoulder line, the neck, the exact point where the arm leaves the torso. The model will happily invent a new shoulder, and an invented shoulder is worse than a hard seam because nobody notices it until the image is on a product page.

Two things make the radius annoying to pin down. It's in pixels, not in a fraction of the body, so it has to be set at whatever resolution the crop is handed to the sampler rather than at the resolution of the original file. And the right value depends on the garment: something tight against the body wants a narrow zone because there's very little skin to blend into, while something with real drape wants a wide one because the true boundary moves with the fabric and SAM2's tight boundary is only one plausible version of it. Five to fifteen pixels covered everything we shot. I never found a principled way to pick inside that range, and it stayed a per-category constant that somebody set by eye.

2. The reference stream

Redux is where the specific garment enters. The distinction I care about is against a plain CLIP vision encoder, which compresses the reference into something close to a caption in vector form. "A photo of a red shirt" is enough to steer colour and category and nowhere near enough to reproduce a print repeat. Redux keeps structural information through the bottleneck, which is why it can carry thread pattern and logo placement.

# reference_stream.py
"""Encode the reference garment into the structural conditioning vector.
 
C_style = ReduxEncoder(SigCLIP(I_ref))
 
The garment is cut out before encoding. Background pixels in a reference
shot are conditioning noise: a studio backdrop pushes the encoder toward
"studio backdrop" and steals capacity from the weave.
"""
 
from dataclasses import dataclass
 
import torch
from PIL import Image
 
 
@dataclass
class ReferenceStreamConfig:
    redux_strength: float = 0.85   # scales C_style before cross-attention
    cutout_feather_px: int = 3     # soften the RMBG alpha edge
    square_pad: bool = True        # pad rather than crop; aspect distortion
                                   # rotates print repeats and it shows
 
 
class ReferenceStream:
    def __init__(self, siglip, redux, rmbg, config: ReferenceStreamConfig):
        self.siglip = siglip
        self.redux = redux
        self.rmbg = rmbg
        self.config = config
 
    def cutout(self, ref: Image.Image) -> Image.Image:
        """Isolate the garment on a transparent channel with RMBG."""
        alpha = self.rmbg(ref)
        garment = ref.convert("RGBA")
        garment.putalpha(alpha.filter_edge(self.config.cutout_feather_px))
        if self.config.square_pad:
            garment = pad_to_square(garment)
        return garment
 
    @torch.inference_mode()
    def encode(self, ref: Image.Image) -> torch.Tensor:
        """Return C_style, the conditioning injected into cross-attention."""
        garment = self.cutout(ref)
        pixel_values = self.siglip.preprocess(garment).to("cuda", torch.bfloat16)
 
        # SigCLIP produces the visual features; Redux turns them into tokens
        # the Flux transformer can attend to alongside the text embedding.
        features = self.siglip(pixel_values)
        c_style = self.redux(features)
 
        c_style = c_style * self.config.redux_strength
        print(f"c_style tokens={tuple(c_style.shape)} "
              f"strength={self.config.redux_strength}")
        return c_style

The redux_strength scalar turned into the most-touched knob in the whole pipeline. Push it up and fidelity to the reference improves until it starts fighting the inpainting stream, at which point the garment stops respecting the pose: you see the reference's own drape stamped onto a body standing differently. Push it down and Style Drift comes back. There's a usable band, it's garment-dependent, and the failure on the high side is much easier to miss than the failure on the low side because a high-strength output looks crisp and detailed while being subtly wrong about the body.

The RMBG cutout is not cosmetic. Reference garment shots arrive as flat lays, on hangers, or on a different model entirely, and whatever is behind the garment gets encoded along with it unless you remove it.

3. Both streams into one sampler

# run_vton.py
"""Single try-on pass: Fill conditioning on the masked body, Redux
conditioning on the reference garment, one sampler, one refinement pass."""
 
SAMPLER = {
    "sampler_name": "euler_ancestral",
    "scheduler": "beta",
    "cfg": 3.5,
}
 
PRIMARY = {**SAMPLER, "steps": 30, "denoise": 1.0}
REFINE = {**SAMPLER, "steps": 10, "denoise": 0.28}
 
# The VAE runs in FP16 regardless of transformer precision. Lower precision
# decoding produced colour blotching on large flat fabric areas, which is
# the worst possible artefact for a solid-colour garment.
VAE_DTYPE = "fp16"
 
 
def try_on(pipe, source, mask, reference, ref_stream, prompt: str):
    """Run the two-stream pass and return the composited result."""
    c_style = ref_stream.encode(reference)
 
    latents = pipe.fill(
        image=source,
        mask_image=mask,
        prompt=prompt,
        style_conditioning=c_style,
        **PRIMARY,
    )
 
    # Second pass at low denoise. This is not a quality upscale; it exists to
    # settle the mask boundary and the high-frequency texture that the first
    # pass leaves slightly soft.
    latents = pipe.fill(
        image=latents,
        mask_image=mask,
        prompt=prompt,
        style_conditioning=c_style,
        **REFINE,
    )
 
    image = pipe.decode(latents, vae_dtype=VAE_DTYPE)
    print(f"primary_steps={PRIMARY['steps']} refine_steps={REFINE['steps']} "
          f"sampler={SAMPLER['sampler_name']}/{SAMPLER['scheduler']}")
    return image

Euler Ancestral with the Beta scheduler was picked by comparison rather than by theory. The ancestral noise helps with fabric texture, which is a high-frequency, slightly stochastic thing that deterministic samplers tend to render too smoothly. Thirty steps for the primary pass and ten for refinement is where quality stopped improving in a way I could see. The refinement pass at low denoise is doing boundary work more than anything else; drop it and the seam at the mask edge starts showing up again often enough to notice.

ControlNet was a baseline, not a component

Worth being clear about this, because the paper cites ControlNet and T2I-Adapter and people assume they're in the stack. They aren't, and they never were. They're related work, and SDXL with a ControlNet was one of the two baselines we measured against. Spatial control over pose or depth is real, but it doesn't solve texture transfer, which is the actual VTON problem. Reference-only ControlNets do attempt texture and in our testing they bled colour and degraded the style. The gap in the benchmark is mostly a global-garment gap: SSIM on Global Traditional garments came in at 0.45 for SDXL with ControlNet against 0.85 for the two-stream pipeline, while the Western Casual numbers were much closer together. Pose control was never the thing standing between us and a shippable saree.

What broke

Sheer fabrics are still bad. Lace and chiffon need the model to blend skin tone with fabric rather than paint fabric over skin, and the pipeline renders them close to opaque. I don't have a fix, only an awareness that we routed those SKUs away from the automated path.

Extreme poses degrade badly. Standing and seated are well covered, anything acrobatic isn't, and the failure is a warp rather than an obvious break, so it doesn't trip any check.

The two knobs interact. Redux strength and dilation radius are not independent: a wider bleeding zone gives the model more freedom, which makes a high Redux strength more likely to run away with the drape. When someone reported a regression, the cause was usually one of those two moving for a different garment category and nobody noticing the other one now needed to move too.

What this taught me

Conditioning channels have shapes, and text is the wrong shape for texture. Most of the early time went into prompt engineering that could never have worked, because the information I was trying to send simply doesn't survive the text encoder.

Splitting responsibility beats tuning one component harder. The single-stream version wasn't undertuned, it was underdetermined, and no amount of sampler work was going to tell it which saree to produce.

The parameters that look boring are the ones that break production. Dilation radius is a one-line morphological op that decides whether a brand accepts the image.

Being explicit about what's a baseline saves a lot of confusion later. Half the questions I got about this architecture were people assuming the related-work section described our pipeline.

The next thing to try is making the dilation radius a function of the segmentation rather than a per-category constant: SAM2 already gives a confidence signal at the boundary, and the places where it's least certain are exactly the places where the fabric moves. Deriving the radius from that would remove the last hand-set number in the masking stage.