DEV Community: Wlad Radchenko

Growing an image past its border without a visible seam

Wlad Radchenko — Wed, 05 Aug 2026 13:38:00 +0000

You have a photo and you want more of it. More sky above the head, more room to the left of the subject, a wider frame than the camera gave you. Outpainting is the name for inventing those new pixels. The model part is well covered everywhere. The part nobody talks about is the few dozen lines that run before the model, and that part is where the seam is won or lost.

If you get the prep wrong, the new region meets the old photo at a hard line. The colors step. The texture changes mid-wall. You can point at exactly where the real photo ended. This post is a walkthrough of the prep code in Wunjo Make that stops that from happening. It lives in one file, visual_generation/generation/outpaint.py, and it is all NumPy and OpenCV. No network in sight.

Outpainting is asking for canvas that was never photographed. The trick is making the new canvas agree with the old one at the join. Photo: Unsplash.

The shape of the problem

Say you want to extend an image to the left. You need a bigger canvas, the old pixels parked on the right side of it, and an empty strip on the left for the new content. Three things have to be true at the end:

The known pixels survive untouched in the middle and far side.
The model is told, precisely, which strip is allowed to change.
The old edge does not become a wall the model paints up against.

The whole function is built around those three. Here are the defaults it runs with:

def process_image(image, fill_color=(0, 0, 0), mask_offset=50, blur_radius=500,
                  expand_pixels=256, direction="left",
                  inpaint_mask_color=50, max_size=1024):

expand_pixels=256 is how much new canvas you grow. mask_offset=50 is the overlap that hides the seam, and it is the most important number in the file. blur_radius=500 feathers the mask edge. max_size=1024 is the hard cap on the output dimension. Keep those four in mind.

Step 1. Size the new canvas, and crop if it would blow the cap

First it computes the new size by adding expand_pixels on the axis you are growing:

new_height = height + (expand_pixels if direction in ["top", "bottom"] else 0)
new_width  = width  + (expand_pixels if direction in ["left", "right"] else 0)

Then the part people forget. If growing the image would push it past max_size, it does not let the canvas grow without limit. It crops the same number of pixels off the opposite side first, so the output stays at max_size:

if new_width > max_size:
    if direction == "left":
        image = image[:, :max_size]            # drop the far-right columns
    elif direction == "right":
        image = image[:, new_width - max_size:] # drop the far-left columns
    new_width = max_size

This is a practical decision, not a cosmetic one. The model downstream has a size it is happy with. Letting the canvas creep past 1024 on every pass would either run out of memory or force a resize that softens the whole image. So the canvas slides: you gain 256 new pixels on one edge and lose the equivalent off the other, holding the frame at a fixed size.

Step 2. Three buffers

After the size is settled, it builds three arrays the same size as the new canvas:

new_image    = np.full((new_height, new_width, 3), fill_color, dtype=np.uint8)  # the canvas
mask         = np.full_like(new_image, 255, dtype=np.uint8)  # soft guidance mask
inpaint_mask = np.full_like(new_image, 0,   dtype=np.uint8)  # hard fill region

mask         = cv2.cvtColor(mask, cv2.COLOR_BGR2GRAY)
inpaint_mask = cv2.cvtColor(inpaint_mask, cv2.COLOR_BGR2GRAY)

Two different masks, and the difference matters.

mask is the soft one that the generation step uses. It starts all white (255), and white here means "keep this." inpaint_mask is the hard one, a quick scratch buffer used a moment later to pre-fill the empty strip. It starts all black (0).

Step 3. Place the pixels and carve the strips

This is the heart of it. For the left direction:

new_image[:, expand_pixels:] = image[:, : max_size - expand_pixels]
mask[:, : expand_pixels + mask_offset] = inpaint_mask_color   # 50
inpaint_mask[:, :expand_pixels] = 255

Line one drops the real pixels into the canvas, shifted right by expand_pixels so the left strip is left empty for the new content.

Line three marks the hard fill region. Exactly the expand_pixels columns on the left, set to 255. That is the genuinely empty part, the 256 columns that were never photographed.

Line two is the trick. It paints the soft mask to a low value (inpaint_mask_color, which is 50) across expand_pixels + mask_offset columns. Not 256. 306. It reaches 50 pixels past the new region, into the real photo.

That 50-pixel overlap is the whole game. If the editable region stopped exactly at the old border, the model would have to match the real photo perfectly at a single column, and it never will, so you would see a line. By letting the model also repaint a thin band of real pixels right at the join, you give it room to ramp the new content into the old. It blends across the edge instead of butting up against it.

Note the mask value is 50, not 0. The region is not marked "fully replace" but "mostly free," a gray instruction rather than a hard switch. Softer guidance near the seam, which is exactly what you want there.

Step 4. Father the mask edge

A sharp transition in the mask still prints a sharp transition in the result. So the mask gets a big Gaussian blur:

if blur_radius % 2 == 0:
    blur_radius += 1                 # OpenCV needs an odd kernel
mask = cv2.GaussianBlur(mask, (blur_radius, blur_radius), 0)

blur_radius=500 becomes a 501x501 kernel. That is enormous, and on purpose. It smears the boundary between the kept region and the editable region across hundreds of pixels, so there is no single row or column where guidance flips. The model gets a gradient of freedom, most free in the new strip, gently constrained as it moves back into real territory. The odd-number guard is there because OpenCV rejects an even kernel size.

Step 5. Pre-fill the empty strip with Telea

Here is the last and most underrated step. The empty strip is still flat fill_color (black, by default). Handing a model a black bar to start from is a bad init. So before returning, it fills that strip with a cheap classical inpaint:

_, mask_np = cv2.threshold(inpaint_mask, 128, 255, cv2.THRESH_BINARY | cv2.THRESH_OTSU)
inpaint = cv2.inpaint(new_image, mask_np, 3, cv2.INPAINT_TELEA)

cv2.inpaint with INPAINT_TELEA (the Fast Marching Method) walks inward from the border of the hole and paints each new pixel from its neighbors. The 3 is the inpaint radius, how far it looks for those neighbors. It is not smart. It will not invent a tree or a window. What it does is bleed the edge colors of the real photo into the empty strip, so the strip starts as a soft smear of the right colors instead of a black void.

Why bother, if a generator is about to repaint it anyway? Because the generator does better work when it starts from a plausible color field than from flat black. The smear carries the lighting and palette of the old photo into the new region before generation even begins. The black-bar init pulls the result toward dark, muddy borders. The Telea smear pulls it toward the photo it is extending.

The function returns two things: inpaint, the canvas with real pixels plus a softly pre-filled strip, and mask, the feathered soft mask. That pair is the init image and the mask you hand to an inpainting or diffusion model. A conditioning setup like this is the same idea behind ControlNet (Zhang, Rao, Agrawala, ICCV 2023): you do not just ask the model to generate, you hand it a structured starting point and a map of where it is allowed to work.

Doing it on both sides: the centering pass

decode_image is the wrapper that turns a single landscape-or-portrait image into a fixed canvas (default 576x1024). It figures out which axis is short, then calls process_image for both sides of that axis and stitches them:

expand_pixels_half_x = int((max_width - width) / 2)
expand_pixels_half_y = int((max_height - height) / 2)

if expand_pixels_half_x > 0:
    expand_pixels_half = expand_pixels_half_x
    directions = ["left", "right"]
else:
    expand_pixels_half = expand_pixels_half_y
    directions = ["top", "bottom"]

Splitting the needed growth in half and growing both edges keeps the original subject centered, instead of shoving it to one side. merge_images then pastes the two halves back together into one canvas.

The wrapper also returns its own combined binary mask, and it pulls in by padding=20 on the long axis:

mask_combined[pad_y_start:pad_y_end, pad_x_start:pad_x_end] = 0

That 20-pixel inset is the same instinct as the 50-pixel offset upstream: do not trust the boundary exactly, let the editable zone eat slightly into the known region so the join has somewhere to blend.

About the author. I'm Wlad Radchenko, a software engineer. The code in this article comes from Wunjo Make (open source), local software for video makers, and Wunjo Design, an offline PWA for designers. Get in touch to find more on GitHub and LinkedIn.

Gotchas you will actually hit

A few things that bit me, in case you build something similar.

The mask polarity is easy to get backwards. In this code white (255) means keep and the low value (50) means free to change. If your downstream model expects the opposite, you will outpaint the wrong region and destroy the photo. Check the convention before anything else.

The 50-pixel overlap is tunable, but do not drop it to zero. Zero overlap brings the hard seam right back. Too large, and the model starts rewriting parts of the photo you wanted kept. Fifty pixels on a 1024 canvas is a reasonable middle.

The Gaussian kernel must be odd. The code guards for this with the % 2 check, and if you change blur_radius, keep that guard or OpenCV will throw.

The opposite-side crop in step 1 is silent. If you feed in an image that is already near max_size and keep extending, you are quietly losing content off the far edge each pass. That is correct for a fixed-canvas product, but surprising if you expected pure growth.

Takeaway

The model invents the new pixels. Everything that makes them join cleanly happens in classical code before the model runs: pad the canvas, place the known pixels, mark the empty strip, overlap the editable mask 50 pixels into the real photo so it can blend across the old edge, feather that mask with a 501-pixel blur, and pre-fill the empty strip with a Telea smear so the model starts from the right colors instead of black. No generation in any of it.

The full file is visual_generation/generation/outpaint.py in the Wunjo Make repo. If you are fighting a visible seam in your own outpainting, the mask_offset overlap and the Telea pre-fill are the two places I would look first.

References

Zhang, Rao, Agrawala. "Adding Conditional Control to Text-to-Image Diffusion Models." ICCV 2023. arXiv:2302.05543
Telea. "An Image Inpainting Technique Based on the Fast Marching Method." Journal of Graphics Tools, 2004. (The INPAINT_TELEA method in OpenCV.)

The file is called encryption.py. It is actually an invisible watermark!

Wlad Radchenko — Tue, 04 Aug 2026 11:37:00 +0000

There is a file in my app called encryption.py. The class is EncryptionEncoder. The methods on the video manager are encrypted() and decrypted(). Every label says crypto.

It is not crypto. There is no cipher in that file, no key, no AES, no Fernet. What it actually does is hide a short text string inside the pixels of a frame so you cannot see it but a decoder can read it back. That is a watermark. The naming is just unfortunate, and I am going to use the honest word for the rest of this post.

I want to walk through the code, because the mechanism is small and clever, and then be straight about what a thing like this can and cannot do when it ships inside a fully local app that the user owns.

The name on the file says encryption. The code says watermark. Photo: Unsplash.

What problem this solves

The app generates synthetic media. Face swaps, lip-sync, retargeted portraits. At some point you want to answer one question about a file: did this come out of my tool, or is it an untouched original?

The honest way to answer that without a visible logo is to embed a marker the eye does not notice and a program can find. So when the app finishes a video, it writes a hidden string into a corner of each frame. Later, an analysis step reads the corners and checks for that string.

Here is the read side, from inference.py:

is_original = VideoManipulation.decrypted(source_file)
if is_original:
    lprint(msg="Content is original.", ...)
else:
    lprint(msg="Content is Wunjo fake.", ...)

So the whole point is a yes/no flag: is this our synthetic output. The marker string in the code is literally name="fake". That is the payload it stamps and the payload it looks for.

How the marker gets into a frame

The payload is just bytes. EncryptionEncoder turns a UTF-8 string into a flat list of bits:

seq = np.array([n for n in content], dtype=np.uint8)
self._encryptions = list(np.unpackbits(seq))   # e.g. b"fake" -> 32 bits
self._wmLen = len(self._encryptions)

Four ASCII characters give 32 bits. That 32 matters later, the decoder is hard-coded to expect 32-bit payloads.

The default embedding method in the app is dwtDctSvd. Three transforms stacked, and each one earns its place. Walk it from the outside in.

Step one, color space. Convert the patch from BGR to YUV and only touch the Y (luma) and U channels:

yuv = cv2.cvtColor(bgr, cv2.COLOR_BGR2YUV)
for channel in range(2):   # only Y and U
    ...

Luma and chroma are where you can nudge values without the change jumping out at a viewer. The V channel is left alone.

Step two, wavelet. Run a single-level Haar dwt2 on the channel. This splits the image into a coarse approximation ca1 and three detail bands:

ca1, (h1, v1, d1) = pywt.dwt2(yuv[:row//4*4, :col//4*4, channel], 'haar')

Everything gets embedded into ca1, the coarse band. Coarse coefficients survive compression and resizing far better than fine detail, which is the first reason this watermark is hard to wipe by accident.

Step three, the actual bit. Chop ca1 into 4x4 blocks. For each block, take a DCT, then an SVD, and quantize the top singular value:

def diffuse_dct_svd(self, block, wmBit, scale):
    u, s, v = np.linalg.svd(cv2.dct(block))
    s[0] = (s[0] // scale + 0.25 + 0.5 * wmBit) * scale
    return cv2.idct(np.dot(u, np.dot(np.diag(s), v)))

That one line is the trick. Read it slowly.

scale is 36. s[0] is the largest singular value of the block, the bulk of its energy. s[0] // scale snaps it down to a multiple of 36. Then you add back a fixed offset that depends only on the bit you want to store: 0.25 * scale for a 0, 0.75 * scale for a 1. So after this, every block sits at either "a quarter of the way into its 36-wide bin" or "three quarters of the way in," and which one it is encodes the bit.

Decoding just asks where in the bin the value landed:

score = int((s[0] % scale) > scale * 0.5)

Below the halfway mark reads 0, above reads 1. The 0.25 and 0.75 placement is deliberate, it parks each bit as far from the decision boundary as possible so JPEG noise has to be large to flip it.

The part that makes it survive: repeat and average

A 32-bit payload does not need many blocks. A 256x256 patch gives you hundreds of 4x4 blocks. So the encoder writes the same 32 bits over and over:

wmBit = self._encryptions[(num % self._wmLen)]   # cycle through the 32 bits

Block 0 holds bit 0, block 32 holds bit 0 again, and so on. Every bit is stamped dozens of times across the patch.

The decoder collects all the readings for each bit position and takes the mean before thresholding:

scores = [[] for i in range(self._wmLen)]
# ... append a score per block to scores[num % wmLen] ...
avgScores = list(map(lambda l: np.array(l).mean(), scores))
bits = (np.array(avgScores) * 255 > 127)

This is the cheap, effective version of error correction. Compression might flip a handful of blocks, but the average over a hundred copies of the same bit barely moves. That is why the marker survives a save-and-reopen.

Where it gets written, and the skip check

The app does not watermark the whole frame. It picks a 256x256 box, usually a corner or the center, and stamps only that:

positions = [
    (0, box_size, 0, box_size),                       # top-left
    (0, box_size, int(width - box_size), width),      # top-right
    (int(height - box_size), height, 0, box_size),    # bottom-left
    (int(height - box_size), height, int(width - box_size), width),
    # ...center, plus some (None, None, None, None) no-op slots
]
y1, y2, x1, x2 = random.sample(positions, 1)[0]

Note the None slots. On video, some frames get no watermark at all, picked at random. That spreads the marker across frames instead of branding every single one, which keeps any per-frame artifact from accumulating into something visible.

Before stamping, the writer decodes first and skips if the marker is already there:

decoded_text = decoder.decode(bgr, 'dwtDctSvd').decode('utf-8')
if decoded_text != name and width >= box_size and height >= box_size:
    # encode

So re-running the tool on its own output does not double-stamp. And the 256x256 minimum is enforced in the encoder too, it raises if r*c < 256*256. Below that you do not have enough blocks for the repeat-and-average to be reliable.

On the read side, decrypted() checks the four corners of each frame and returns False (not original) the moment any corner decodes to the marker. One hit is enough.

Gotchas you will actually hit

A few things bit me, and will bit you if you reuse this code:

The payload length is wired to 32 bits. EncryptionDecoder('bytes', 32) everywhere. Your marker string has to be exactly four bytes, or you change the length in both encoder and decoder. set_by_ipv4 even asserts self._wmLen == 32.
Lossless transcode matters. The video is reassembled with ffmpeg ... -c:v copy, a stream copy, no re-encode. If you instead re-encode at a low bitrate, you are testing how much the watermark survives, and the answer depends on scale. A bigger scale survives harder compression but distorts the patch more.
The rivaGan method needs ONNX models. The third method loads encoder.onnx and decoder.onnx from the module directory and even runs chmod/icacls on them to fix file permissions. The app does not use this path by default; dwtDctSvd is what videoio.py calls.
There is dead code. infer_dct_svd has a return score followed by an unreachable if score >= 0.5. Do not be confused by it, the function returns the int and exits.

The honest part: this is obfuscation, not protection

Now the uncomfortable bit. This ships inside a fully local desktop app. Everything is on the user's machine: the encoder, the decoder, and the marker string fake, all in plain Python in an open-source repo.

So be clear about what that means:

There is no secret. Anyone who reads the repo knows the algorithm, the scale, and the payload. They can read the marker, strip it, or forge it onto an untouched file.
It is not DRM. It does not stop anyone from using, copying, or editing the output. It does not lock anything.
It is not a cipher, despite the filename. A cipher needs a key the attacker does not have. A local app has no place to hide a key from its own user, so this scheme does not try.

What it actually buys you is provenance under cooperative conditions. If a file passes through normal handling, a re-save, a crop that leaves a corner intact, a platform transcode, the marker is likely still readable, and the app can honestly say "this came out of me." That is genuinely useful for a tool that generates synthetic media and wants to flag its own output. It is a fragile-but-invisible label, not a lock.

The threat it addresses is accidental loss of provenance, not a motivated adversary. Knowing which threat you are defending against is the whole game. Calling it encryption oversells it. Calling it a watermark, and being clear that the watermark is removable by anyone who cares, is the truth.

Wrap-up

The mechanism is one quantized singular value per 4x4 block, the same 32 bits repeated across a 256x256 patch, decoded by averaging. Small, hard to wipe by ordinary handling, and completely open. If you want to read the whole thing, it is portable/src/visual_processing/utils/encryption.py in the repo, with the call site in videoio.py.

If you take one thing away: name your modules for what they do. A watermark called encryption.py will make every future reader, including you, reason about a threat model that the code never had.

LivePortrait without the puppet look: stitching and retargeting in a 126-number MLP

Wlad Radchenko — Wed, 29 Jul 2026 12:36:00 +0000

You animate a still portrait with LivePortrait. The face moves, the mouth talks, the eyes blink. And then you notice it: the head floats. It drifts a little off the neck, the jaw slides past the shoulders, and the whole thing reads like a paper mask on a stick. The puppet look.

I hit this in Wunjo Make while wiring up the portrait animation path. The fix is not a bigger generator or a fancier warp. It is three small multi-layer perceptrons that predict tiny corrections to the face keypoints. One keeps the head attached to the body. One controls the eyes. One controls the lips. They are so small you could miss them in the repo, and they are the reason the output stops looking like a puppet.

This post walks through what those networks take in, what they output, and how a handful of numbers gets added back to the keypoints to fix the artifact. All of it from the actual code in the repo.

The input is one still photo. Everything after is keypoints and deltas. Photo: Unsplash

The one network class, used three ways

Here is the whole network definition. It lives in portable/src/visual_processing/portrait_animation/modules/stitching_retargeting_network.py:

class StitchingRetargetingNetwork(nn.Module):
    def __init__(self, input_size, hidden_sizes, output_size):
        super(StitchingRetargetingNetwork, self).__init__()
        layers = []
        for i in range(len(hidden_sizes)):
            if i == 0:
                layers.append(nn.Linear(input_size, hidden_sizes[i]))
            else:
                layers.append(nn.Linear(hidden_sizes[i - 1], hidden_sizes[i]))
            layers.append(nn.ReLU(inplace=True))
        layers.append(nn.Linear(hidden_sizes[-1], output_size))
        self.mlp = nn.Sequential(*layers)

    def forward(self, x):
        return self.mlp(x)

That is it. Linear, ReLU, Linear, ReLU, Linear. No convolutions, no attention, no recurrence. A plain stack of fully connected layers. The interesting part is that this same class is instantiated three times with different sizes, and each instance solves a different problem.

The sizes come straight from the model config in live_portrait_wrapper.py:

'stitching_retargeting_module_params': {
    'stitching': {
        'input_size': 126, 'hidden_sizes': [128, 128, 64], 'output_size': 65
    },
    'lip': {
        'input_size': 65, 'hidden_sizes': [128, 128, 64], 'output_size': 63
    },
    'eye': {
        'input_size': 66, 'hidden_sizes': [256, 256, 128, 128, 64], 'output_size': 63
    }
}

Those numbers are the whole story. To read them you need one fact: LivePortrait represents a face as 21 implicit 3D keypoints. The config sets num_kp: 21. Twenty one points, each with an x, y, z. So one full set of keypoints is 21 * 3 = 63 numbers. Hold onto 63 and 21, because every size below is built from them.

Stitching: keep the head on the shoulders

Think about what causes the floating head. The animation pipeline takes the source keypoints (the still photo) and the driving keypoints (the motion you want), warps the source feature volume from one to the other, and decodes an image. The warp is good at the face. It does not know or care that the face has to connect to a neck and shoulders that are not moving. So the face can shift as a block and detach from the body.

Stitching is a learned correction that nudges the driving keypoints so the face lines up with where the body already is. Plain version: before we render, we ask a small network "given where the face started and where you want it to go, what small shift puts it back on the shoulders?" Then we add that shift.

The input is 126 numbers. That is 63 for the source keypoints plus 63 for the driving keypoints, flattened and concatenated. Here is the concat, from utils/helper.py:

def concat_feat(kp_source: torch.Tensor, kp_driving: torch.Tensor) -> torch.Tensor:
    # kp_source: (bs, k, 3), kp_driving: (bs, k, 3) -> Return: (bs, 2k*3)
    feat = torch.cat([kp_source.view(bs_src, -1), kp_driving.view(bs_dri, -1)], dim=1)
    return feat

So the stitching MLP sees both poses at once: where the head is anchored, and where the animation wants it. Now the output, 65 numbers. Not 63. The extra two matter. Here is how the wrapper unpacks them in stitching():

def stitching(self, kp_source, kp_driving):
    bs, num_kp = kp_source.shape[:2]
    kp_driving_new = kp_driving.clone()
    delta = self.stitch(kp_source, kp_driving_new)

    delta_exp = delta[..., :3*num_kp].reshape(bs, num_kp, 3)     # 63 -> (21, 3)
    delta_tx_ty = delta[..., 3*num_kp:3*num_kp+2].reshape(bs, 1, 2)  # the last 2

    kp_driving_new += delta_exp
    kp_driving_new[..., :2] += delta_tx_ty
    return kp_driving_new

Read it slowly. The first 63 outputs are a per-keypoint 3D delta. Each of the 21 points gets its own small x, y, z correction. Those get added to every keypoint. The last two outputs are a single global translation in x and y, added to all keypoints at once. So the network can both reshape the face locally and slide the whole face as one piece to meet the shoulders. That global tx, ty is the literal "stop floating" knob.

The correction is small by design. Note that the network has an initialize_weights_to_zero method that zeros every weight and bias. Trained from a zero start, the natural resting state of the delta is "change nothing," and the network only learns to add a nudge where a nudge is needed. A correction network, not a generator.

Eye and lip retargeting: control the blink and the mouth

The other two instances solve a different problem. When you drive one person's face with another's, the eyes and lips do not always transfer cleanly. A person with small eyes driving a person with large eyes leaves the large eyes half open when they should be shut. Lips that should close stay slightly parted. The driving motion is in the right direction but the wrong amount.

Retargeting fixes the amount. The trick: instead of feeding the network a pose, you feed it a target. You tell it "I want the eyes this open" or "I want the lips this closed," and it returns the keypoint delta that gets you there.

That target is a ratio computed straight from facial landmarks, in utils/retargeting_utils.py:

def calculate_distance_ratio(lmk, idx1, idx2, idx3, idx4, eps=1e-6):
    return (np.linalg.norm(lmk[:, idx1] - lmk[:, idx2], axis=1, keepdims=True) /
            (np.linalg.norm(lmk[:, idx3] - lmk[:, idx4], axis=1, keepdims=True) + eps))

def calc_eye_close_ratio(lmk, target_eye_ratio=None):
    lefteye_close_ratio  = calculate_distance_ratio(lmk, 6, 18, 0, 12)
    righteye_close_ratio = calculate_distance_ratio(lmk, 30, 42, 24, 36)
    ...

def calc_lip_close_ratio(lmk):
    return calculate_distance_ratio(lmk, 90, 102, 48, 66)

A close ratio is vertical opening over horizontal width. For an eye: distance between the top and bottom lids divided by the eye's corner-to-corner width. A wide open eye has a high ratio, a shut eye a low one. Dividing by width makes it scale-free, so it works whether the face is big or small in the frame. The eps = 1e-6 is there so a degenerate landmark set never divides by zero.

Now the sizes click into place. The lip network takes input_size 65: that is 63 for the source keypoints plus a lip ratio of 2 numbers. From the wrapper:

def retarget_lip(self, kp_source, lip_close_ratio):
    # kp_source: BxNx3, lip_close_ratio: Bx2
    feat_lip = concat_feat(kp_source, lip_close_ratio)
    delta = self.stitching_retargeting_module['lip'](feat_lip)
    return delta.reshape(-1, kp_source.shape[1], 3)   # -> (B, 21, 3)

The eye network takes input_size 66: 63 keypoints plus an eye ratio of 3 numbers (left eye, right eye, and a target). Both output 63, a per-keypoint 3D delta that you add to the source keypoints. The eye network has the deeper stack, [256, 256, 128, 128, 64], because opening and closing eyes the right amount across identities is the harder regression. Same class, more layers.

The pattern is identical to stitching: feed in the current keypoints plus a small control signal, get back a keypoint delta. The only difference is what the control signal is. For stitching it is the driving pose. For retargeting it is a ratio you can dial yourself.

How it actually runs

These three calls compose. In the wrapper you can see the full chain when you control everything by hand:

eyes_delta = self.retarget_eye(x_s_user, combined_eye_ratio_tensor)
lip_delta  = self.retarget_lip(x_s_user, combined_lip_ratio_tensor)
x_d_new    = x_s_user + eyes_delta + lip_delta
x_d_new    = self.stitching(x_s_user, x_d_new)

Deltas just add. You take the source keypoints, add the eye correction, add the lip correction, then run stitching last so the final pose is glued back to the body after you have changed the eyes and lips. Order matters: stitch after retargeting, not before, so the global tx, ty accounts for everything you moved.

Every call runs under torch.no_grad(). These are inference-only correction heads. They are also cheap. A 126-input MLP with three hidden layers is a rounding error next to the appearance extractor and the dense motion network, which is why LivePortrait can offer this control without a real frame-rate cost.

Gotchas worth knowing

A few things I learned the slow way.

Stitching is on by default, retargeting is not. In the config, flag_stitching defaults to True and the eye and lip retargeting flags default to False. The comments in argument_config.py flag eye and lip retargeting as "not recommend to be True, WIP." Stitching is the safe, always-want-it correction. Eye and lip retargeting are an expressive extra that can introduce jitter, which is why in the Wunjo path the per-frame deltas get run through a Kalman filter before they are applied. If your blink flickers frame to frame, that smoothing is the lever.

Stitching does not load by default for animals. The same comment notes you want stitching off for large head movement or for animal sources. The networks were trained on the human face keypoint layout, so the corrections only make sense there.

The whole module is optional. Look at the wrapper init: if checkpoint_S is missing, self.stitching_retargeting_module = None, and stitching() returns the driving keypoints untouched. The pipeline still runs. It just runs with the floating head. That is the cleanest way to see what these tiny networks buy you: delete them and watch the head detach.

Wrap-up

The lesson I keep relearning: the artifact you cannot stand is often fixed by the smallest part of the system. Here the floating head, the puppet look, the half-open eyes, all come down to three ReLU MLPs that predict a few keypoint deltas and add them back. 126 numbers in, 65 out for stitching. 65 in, 63 out for lips. 66 in, 63 out for eyes. Everything else in the pipeline is huge by comparison, and none of it would look right without these.

If you want to read the code yourself, start at stitching_retargeting_network.py for the class, then live_portrait_wrapper.py for stitching, retarget_eye, and retarget_lip. The full repo is open source on GitHub.

References

Guo, Zhang, Liu, et al. "LivePortrait: Efficient Portrait Animation with Stitching and Retargeting Control." 2024. arXiv:2407.03168. https://arxiv.org/abs/2407.03168

A safety gate inside a face-swap tool or how the NSFW check is wired into the pipeline

Wlad Radchenko — Tue, 28 Jul 2026 13:33:00 +0000

A face-swap tool can be pointed at the wrong content. That is not a hypothetical, it is the first thing some users try. So before the swap runs, there is a check: a small classifier looks at the target frame, and if it sees explicit content, the swap is skipped. The original frame goes through untouched.

This is a walkthrough of that check in an open-source repo, down to the labels and the number that decides. It is short code. The interesting part is where it sits in the pipeline and what it does when it fires.

The check sits in front of the swap, not after it. Nothing explicit ever gets a new face. Photo: Unsplash.

What the check actually is

The classifier is NudeNet, run as an ONNX model called filter.onnx. It is a single-stage object detector, the same family as YOLO: you give it an image, it gives you back boxes, each box with a class label and a confidence score. It has no academic paper, so I will describe it only from what the code does with it.

The model knows 18 labels. Here they are, straight from the file:

__labels = [
    "FEMALE_GENITALIA_COVERED", "FACE_FEMALE", "BUTTOCKS_EXPOSED",
    "FEMALE_BREAST_EXPOSED", "FEMALE_GENITALIA_EXPOSED", "MALE_BREAST_EXPOSED",
    "ANUS_EXPOSED", "FEET_EXPOSED", "BELLY_COVERED", "FEET_COVERED",
    "ARMPITS_COVERED", "ARMPITS_EXPOSED", "FACE_MALE", "BELLY_EXPOSED",
    "MALE_GENITALIA_EXPOSED", "ANUS_COVERED", "FEMALE_BREAST_COVERED",
    "BUTTOCKS_COVERED",
]

Notice the split: most labels come in COVERED and EXPOSED pairs. The model is not a single yes/no nudity switch. It localizes specific body parts and tells you whether each one is covered or exposed. That distinction is the whole basis of the gate, because the gate only cares about the exposed ones.

Preprocessing: letterbox to 320x320

The input side is plain object-detector plumbing, but it matters because the boxes have to map back to the original image. The model wants a square 320x320 input. An arbitrary photo is not square, so the image is resized to fit one dimension, then padded with black bars to fill the rest:

def _read_image(img, target_size=320):
    img_height, img_width = img.shape[:2]
    img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
    aspect = img_width / img_height

    if img_height > img_width:
        new_height = target_size
        new_width = int(round(target_size * aspect))
    else:
        new_width = target_size
        new_height = int(round(target_size / aspect))

    resize_factor = math.sqrt(
        (img_width**2 + img_height**2) / (new_width**2 + new_height**2)
    )
    img = cv2.resize(img, (new_width, new_height))
    # ... pad with cv2.copyMakeBorder to reach target_size ...
    image_data = img.astype("float32") / 255.0
    image_data = np.transpose(image_data, (2, 0, 1))  # HWC -> CHW
    image_data = np.expand_dims(image_data, axis=0)   # add batch dim
    return image_data, resize_factor, pad_left, pad_top

The three values returned alongside the tensor, resize_factor, pad_left, and pad_top, are the undo instructions. The detector finds boxes in the padded 320x320 space, and these let postprocessing put them back where they belong in the real image. Pixels are scaled to 0..1 and the layout is flipped from height-width-channels to channels-first, which is what the ONNX graph expects.

Postprocessing: two thresholds, then NMS

The model outputs one row per candidate box. Each row is four box coordinates followed by the per-class scores. Postprocessing keeps a candidate only if its best class scores at least 0.2, then maps the box back to image space:

for i in range(rows):
    classes_scores = outputs[i][4:]
    max_score = np.amax(classes_scores)
    if max_score >= 0.2:
        class_id = np.argmax(classes_scores)
        x, y, w, h = outputs[i][0], outputs[i][1], outputs[i][2], outputs[i][3]
        left = int(round((x - w * 0.5 - pad_left) * resize_factor))
        top  = int(round((y - h * 0.5 - pad_top)  * resize_factor))
        # ... collect boxes, scores, class_ids ...

indices = cv2.dnn.NMSBoxes(boxes, scores, 0.25, 0.45)

That 0.2 is a loose first pass: keep almost anything the model is not sure is nothing. The pad_left and pad_top subtractions strip off the black bars, and multiplying by resize_factor scales the box back up to the original resolution. Then non-maximum suppression (NMSBoxes with a 0.25 score cutoff and 0.45 IoU) throws away overlapping duplicate boxes so each region is reported once.

The output of detect() is a list of dicts, each {"class", "score", "box"}. That is the raw material. The gate is built on top of it.

The gate itself: status()

This is the method the pipeline calls. It runs detection, keeps only the labels it cares about above a strict score, and returns a single boolean:

def status(self, img, classes=None):
    detections = self.detect(img)
    if classes is None:
        classes = [
            "BUTTOCKS_EXPOSED",
            "FEMALE_BREAST_EXPOSED",
            "FEMALE_GENITALIA_EXPOSED",
            #"MALE_BREAST_EXPOSED",
            "ANUS_EXPOSED",
            "MALE_GENITALIA_EXPOSED",
        ]
    detections = [d for d in detections
                  if d["class"] in classes and d["score"] > 0.7]
    if len(detections) > 0:
        return False
    return True

Read the return values carefully, because the polarity is easy to get backwards. True means clean, go ahead. False means a blocking detection was found, stop.

Three design choices live in those lines:

Only EXPOSED classes block. The default list is five exposed body parts. The COVERED variants are not in it. A swimsuit or underwear photo, which the model would tag as COVERED, does not trip the gate. The tool is filtering explicit nudity, not skin.
The threshold jumps from 0.2 to 0.7. Detection keeps candidates at 0.2 so nothing is missed early. The gate then demands 0.7 confidence before it will block. That gap is a deliberate bias toward not falsely rejecting clean content. A low-confidence guess does not stop a user's legitimate edit.
MALE_BREAST_EXPOSED is commented out. A bare male chest is an everyday photo, so it was pulled from the block list. That one commented line is a small, honest record of a real product call: where to draw the line between explicit and ordinary.

There is also a sibling method, censor(), that uses the same class list and threshold but instead of returning a boolean it paints black rectangles over the detected boxes. The pipeline does not use censor() for the swap gate; it uses status(). Same detector, two policies, and the swap path chose the hard stop.

Where it is wired in

A detector that nobody calls is decoration. The point of this article is the wiring, so here is every place the swap pipeline consults the gate. It is constructed once, in the processor's __init__:

self.filter_model = NudeDetector(
    providers=self.access_providers if self.device != "cpu" else None
)

Then it guards every swap path. For a single image:

def swap_image(self, target_frame, source_face, face_fields, save_dir, multiface=False):
    save_file = os.path.join(save_dir, "swapped_image.png")
    if self.filter_model.status(target_frame):
        # ... detect faces, run face_swap_model.get(...) ...
    else:
        print("Face swap cannot be applied to nude or explicit content. "
              "Please upload images with appropriate content.")
        cv2.imwrite(save_file, target_frame)   # save the ORIGINAL, unswapped
        return target_frame

For video on the threaded CPU path, the same guard sits inside the per-frame worker:

def process_frame(self, args):
    frame, face_det_result, source_face, ... = args
    tmp_frame = frame.copy()
    if self.filter_model.status(tmp_frame):
        for face in face_det_result:
            if face is None:
                break
            tmp_frame = self.face_swap_model.get(tmp_frame, face, source_face, paste_back=True)
    else:
        print("Face swap cannot be applied to nude or explicit content. ...")
    cv2.imwrite(filename, tmp_frame)   # explicit frame is written unmodified
    return filename

And the CUDA video path has the identical check before its swap call. Three entry points, one rule: status() runs on the target frame first, and face_swap_model.get(...) only runs if it returns True.

Two details worth pulling out:

The check runs on the target, per frame. The target is the content you are swapping a face into. For a video, that means the classifier runs on every frame independently. A clip that is clean for ten seconds and explicit for two will get faces swapped only on the clean frames. The block is granular, not all-or-nothing for the file.
A blocked frame is not dropped, it is passed through. Look at the else branches: the original frame still gets written to disk and into the output video. The gate removes the swap, not the content. The user gets their video back, just without a new face on the frames that tripped the filter.

The face side, for context

The swap that the gate protects is built on InsightFace. Recognition uses the buffalo_l pack, which is an ArcFace embedding model (Deng et al., CVPR 2019). ArcFace turns a face into a vector where same-person faces sit close together, which is how the pipeline finds and aligns the right face to replace. The gate is upstream of all of it. The recognition and the swap model never even see a frame that status() rejected.

Gotchas if you reuse this

A few things that will bite you if you lift this pattern into your own pipeline.

Get the boolean polarity right. status() returns True for clean. I have seen people wire if status(): block() and silently invert the whole safety behavior. Test both cases with a known-clean and a known-explicit image before you trust it.

It runs on CPU by default, and per frame. In the constructor, the detector is given the CUDA provider only when a GPU device is set; otherwise it is CPU. On a long video the gate is real per-frame work. It is cheap next to the swap itself, but it is not free, and it is in the hot loop.

A detector is not a guarantee. This is a 320x320 model with a 0.7 threshold tuned to avoid false blocks. It will miss things, and it can be fooled. It is a reasonable first line, not a proof. If you ship a tool like this, the filter is one layer; clear terms of use and the open-source license are others.

The threshold and class list are policy, not magic numbers. The 0.7 cutoff and the five-item list are where "what counts as explicit" is actually decided in code. If you change them, you are changing the product's stance, not tuning performance. The commented-out male chest line is the example to study.

Takeaway

The safety gate is about thirty lines of real logic. A YOLO-style ONNX model detects body parts, status() keeps only the five exposed classes above 0.7 confidence and returns a boolean, and three swap entry points refuse to run the swap when that boolean is False, passing the original frame through instead. The check runs on the target content, per frame, before recognition or swapping touches it.

The numbers are the policy. The 0.7 threshold biases toward not blocking clean work, the exposed-only class list filters nudity rather than skin, and the commented MALE_BREAST_EXPOSED line marks one place the line was drawn by hand. The code is in visual_processing/face_swap/nudenet.py and swap.py in the Wunjo Make repo.

References

Deng, Guo, Xue, Zafeiriou. "ArcFace: Additive Angular Margin Loss for Deep Face Recognition." CVPR 2019. arXiv:1801.07698 (InsightFace / buffalo_l recognition.)

Filling a hole in video with the next frame, not a guess

Wlad Radchenko — Wed, 22 Jul 2026 12:32:00 +0000

You paint a mask over a person walking across a shot and ask the app to remove them. Behind that person, in this exact frame, there is nothing. The pixels do not exist. But two frames later the person has moved, and the patch of wall they were standing in front of is right there, fully visible.

So why would you ask a network to imagine that wall when the camera already filmed it?

That is the whole idea behind the object remover in Wunjo Make. Most of the hole does not get hallucinated. It gets carried in from frames where the same spot was not covered. The code does this with optical flow, and the interesting part is the ordering: flow first, pixels second, and a transformer only for the leftovers. This is a walkthrough of that pipeline as it runs in portable/src/visual_processing/remover/.

Across a strip of frames, a covered spot in one frame is usually open in another. Flow is the map that connects them. Photo: Unsplash.

The plan in plain words

Think of flow as a set of tiny arrows, one per pixel, that say "this pixel moved here in the next frame." If you have those arrows, you can follow a pixel backward and forward through time. A spot hidden by the object in frame 10 might be visible in frame 7. Follow the arrows from 7 to 10 and you can drop frame 7's real pixel into frame 10's hole.

There is a catch. The arrows themselves are missing inside the mask. Optical flow is computed from image content, and inside the hole the content is the object you want gone. So the motion of the background under the object is unknown. The pipeline solves this in a specific order:

Estimate flow between every pair of adjacent frames (RAFT).
Fill in the flow inside the masked region, so the arrows exist even where the background was hidden (recurrent flow completion).
Use that completed flow to warp real pixels from neighbor frames into the hole, and shrink the mask wherever this works.
Whatever is still masked after all that, because it was never visible in any frame, gets painted by a transformer.

Four stages, and only the last one invents anything.

Stage 1. The arrows (RAFT)

The flow estimator is RAFT (Teed and Deng, ECCV 2020). In the pipeline it runs on short clips and produces both directions at once:

flows_f, flows_b = self.fix_raft(frames[:, f:end_f], iters=raft_iter)

flows_f is forward flow (frame t to t+1), flows_b is backward (t to t-1). raft_iter defaults to 20. RAFT is iterative: it starts from a zero flow field and refines it a fixed number of times, each step looking up a correlation volume and adding a small delta. You can see that loop in raft.py:

for itr in range(iters):
    coords1 = coords1.detach()
    corr = corr_fn(coords1)              # index correlation volume
    flow = coords1 - coords0
    net, up_mask, delta_flow = self.update_block(net, inp, corr, flow)
    coords1 = coords1 + delta_flow       # F(t+1) = F(t) + delta

The clip length is tied to resolution. The code drops it as frames get bigger:

if frames.size(-1) <= 640:
    short_clip_len = 12
elif frames.size(-1) <= 720:
    short_clip_len = 8
elif frames.size(-1) <= 1280:
    short_clip_len = 4
else:
    short_clip_len = 2

This is a memory budget, not an accuracy choice. RAFT's correlation volume is large, so wider frames mean fewer of them in VRAM at once. The result of this stage is gt_flows_bi, a clean flow field everywhere except inside the mask, where the arrows are garbage because they were computed over the object.

Stage 2. Completing the flow inside the hole

This is the part people skip past, and it is the one that makes the whole thing work. Before you can move pixels along the arrows, you need arrows inside the hole. The background moved under that object even though you could not see it, and the model has to guess that motion.

Here is the key intuition: guessing motion is far easier than guessing pixels. A wall's appearance can be anything. A wall's motion is almost always smooth and continuous with the motion just outside the mask. If the visible background is panning left, the hidden background behind it is panning left too. Flow is a slowly-varying field, so you can extend it inward from the edges of the hole with real confidence.

The completion network starts by deleting the unreliable flow inside the mask, on purpose:

masked_flows_forward  = masked_flows_bi[0] * (1 - masks_forward)
masked_flows_backward = masked_flows_bi[1] * (1 - masks_backward)

Multiplying by (1 - mask) zeroes the flow wherever the mask is 1. It does not try to salvage the bad arrows. It throws them out and rebuilds them. The masked flow plus the mask itself go into a 3D convolutional encoder, a dilated middle block, a bidirectional propagation module, and a decoder, all in recurrent_flow_completion.py:

inputs = torch.cat((masked_flows, masks), dim=1)
x = self.downsample(inputs)
feat_e1 = self.encoder1(x)
feat_e2 = self.encoder2(feat_e1)
feat_mid = self.mid_dilation(feat_e2)
...
feat_prop = self.feat_prop_module(feat_mid)   # bidirectional, second-order
...
flow = self.upsample(feat_d1)

The propagation runs both directions, forward and backward, so a frame's flow is informed by its neighbors on both sides. The backward pass literally flips the clip in time, completes it, and flips it back:

masked_flows_backward = torch.flip(masked_flows_backward, dims=[1])
masks_backward        = torch.flip(masks_backward, dims=[1])
pred_flows_backward, _ = self.forward(masked_flows_backward, masks_backward)
pred_flows_backward    = torch.flip(pred_flows_backward, dims=[1])

Then the last step, and it is a small line that matters a lot. combine_flow keeps the original measured flow everywhere outside the mask and only uses the predicted flow inside it:

pred_flows_forward  = pred_flows_bi[0] * masks_forward  + masked_flows_bi[0] * (1 - masks_forward)
pred_flows_backward = pred_flows_bi[1] * masks_backward + masked_flows_bi[1] * (1 - masks_backward)

The network never gets to "improve" the flow it did not need to touch. Outside the hole, RAFT's measurement wins. Inside the hole, the prediction fills in. You end up with one continuous flow field with no seam at the mask boundary.

There is an edgeDetector in this network too, but check the forward pass: it only runs under self.training. At inference it is skipped. So in the shipped app it contributes nothing, it was a training-time loss helper.

Stage 3. Moving real pixels along the completed flow

Now the arrows exist everywhere. Time to follow them and carry pixels in. First the object is blanked out of the frames:

masked_frames = frames * (1 - masks_dilated)

Then image propagation runs with the completed flow:

prop_imgs, updated_local_masks = model.img_propagation(
    masked_frames, pred_flows_bi, masks_dilated, 'nearest'
)

Inside, this is a BidirectionalPropagation module built with learnable=False. That flag matters. The feature-level propagation later in the pipeline is a trained network. This one is not. It is plain warping and bookkeeping, no weights, because here we are moving actual RGB pixels and we do not want a network reinterpreting them. The non-learnable branch in propainter.py is where the real work happens:

mask_prop_valid = flow_warp(mask_prop, flow_prop.permute(0, 2, 3, 1))
mask_prop_valid = self.binary_mask(mask_prop_valid)

union_vaild_mask = self.binary_mask(mask_current * flow_vaild_mask * (1 - mask_prop_valid))
feat_prop = union_vaild_mask * feat_warped + (1 - union_vaild_mask) * feat_current
mask_prop = self.binary_mask(mask_current * (1 - (flow_vaild_mask * (1 - mask_prop_valid))))

Read union_vaild_mask slowly, because it is the gate that decides "can I trust the pixel I just dragged in here." Three conditions, all multiplied:

mask_current is 1 where this frame still has a hole to fill.
flow_vaild_mask is 1 where the flow is trustworthy.
(1 - mask_prop_valid) is 1 where the source pixel I am pulling from was itself not a hole.

Only where all three hold does the warped pixel get written in. Everywhere else the frame keeps what it had. Then mask_prop updates: the hole shrinks wherever the fill succeeded.

That flow trust check is its own neat trick, fbConsistencyCheck. Follow a pixel forward with the forward flow, then backward with the backward flow. If you land where you started, the flow is consistent and probably correct. If you land somewhere else, something is occluded or wrong, so do not trust it:

flow_bw_warped = flow_warp(flow_bw, flow_fw.permute(0, 2, 3, 1))
flow_diff_fw   = flow_fw + flow_bw_warped
mag_sq_fw      = length_sq(flow_fw) + length_sq(flow_bw_warped)
occ_thresh_fw  = alpha1 * mag_sq_fw + alpha2
fb_valid_fw    = (length_sq(flow_diff_fw) < occ_thresh_fw).to(flow_fw)

flow_fw + flow_bw_warped should cancel to near zero for a good pixel. The threshold scales with the flow magnitude (alpha1 = 0.01) plus a constant floor (alpha2 = 0.5), so fast motion gets more slack than slow motion. This is what stops the propagation from smearing a pixel across an occlusion boundary.

After this stage the frames are updated and the mask is smaller. The bookkeeping line that writes it back:

updated_frames = frames * (1 - masks_dilated) + prop_imgs.view(b, t, 3, height, width) * masks_dilated
updated_masks  = updated_local_masks.view(b, t, 1, height, width)

For a lot of videos, a clean pan across a static background, this stage alone fills almost the entire hole with genuine filmed pixels. The output is real footage rearranged in time, not a generated guess.

Stage 4. The transformer fills only what was never seen

Some pixels were covered in every single frame. The object stood still over a spot the camera never caught. No amount of flow can carry in a pixel that does not exist anywhere. That is the only thing left for a generative model.

The inpainting generator (ProPainter, Zhou et al., ICCV 2023) runs on a sliding window of neighbor frames plus a few reference frames:

pred_img = model(selected_imgs.half(), selected_pred_flows_bi,
                 selected_masks, selected_update_masks, l_t)

Notice it takes two masks. selected_masks is the original hole, selected_update_masks is the smaller hole left after stage 3. The encoder gets the masked frame plus both masks stacked as channels, so it knows what was filled by propagation and what still needs inventing:

enc_feat = self.encoder(torch.cat([masked_frames.view(b * t, 3, ori_h, ori_w),
                                   masks_in.view(b * t, 1, ori_h, ori_w),
                                   masks_updated.view(b * t, 1, ori_h, ori_w)], dim=1))

Inside InpaintGenerator.forward, the flow is reused once more, this time to propagate features (the learnable feat_prop_module, downsampled flow at quarter resolution), and only then does the transformer attend across time:

_, _, local_feat, _ = self.feat_prop_module(local_feat, ds_flows_f, ds_flows_b, prop_mask_in, interpolation)
enc_feat = torch.cat((local_feat, ref_feat), dim=1)
trans_feat = self.ss(enc_feat.view(-1, c, h, w), b, fold_feat_size)
trans_feat = self.transformers(trans_feat, fold_feat_size, mask_pool_l, t_dilation=t_dilation)

The transformer is the expensive, creative part, and by the time it runs, most of the hole is already real footage. It only has to dream up the slivers that were genuinely never filmed. That is why the output rarely looks hallucinated: there is very little hallucination in it.

Gotchas you will actually hit

A few things bite when you run this:

The flow array is one shorter than the frames. Flow is between pairs, so N frames give N-1 flows. The code slices accordingly, neighbor_ids[:-1] when selecting flows. Off-by-one here corrupts the warps silently, the video does not crash, it just smears.
Dilate the mask. The pipeline uses masks_dilated, not the raw mask. A tight mask leaves a halo of the object's edge pixels just outside the hole, and propagation happily carries that halo into other frames. Grow the mask a few pixels so the object's fringe is inside it.
Flow completion before propagation, always. If you propagate with raw RAFT flow, the arrows inside the hole are nonsense and you drag object pixels around instead of background. The completion step is not optional polish, it is what makes the warp point at the right source.
The edge detector is dead weight at inference. Do not wire it into your inference path expecting an edge map. It returns None outside training.
fp16 is fine downstream, not for RAFT. The pipeline runs RAFT in higher precision and casts to half later (selected_imgs.half()), because the correlation lookup is sensitive to precision in a way the conv layers are not.

Why this ordering is the point

You could build a remover that hands the whole masked video to a transformer and says "figure it out." It would invent the background frame by frame, and it would flicker, because an inventor has no reason to invent the same wall twice. The flow-guided approach inverts that. It treats inpainting as a search problem first ("where else in this video is this pixel visible?") and a generation problem only for what the search cannot find. Real pixels are temporally consistent for free, because they are the same pixels.

The code is in portable/src/visual_processing/remover/. The three model files are raft.py (the arrows), recurrent_flow_completion.py (completing the arrows inside the hole), and propainter.py (warping pixels, then the transformer). The orchestration that runs them in order is in __init__.py, in process_video_with_mask.

References

Zhou, Li, Loy, et al. "ProPainter: Improving Propagation and Transformer for Video Inpainting." ICCV 2023. arXiv:2309.03897. https://arxiv.org/abs/2309.03897
Teed, Deng. "RAFT: Recurrent All-Pairs Field Transforms for Optical Flow." ECCV 2020. arXiv:2003.12039. https://arxiv.org/abs/2003.12039

Removing text without removing the wall behind it: EAST quads to an inpaint mask

Wlad Radchenko — Tue, 21 Jul 2026 11:29:00 +0000

You want to wipe a sign off a wall in a video. A subtitle, a watermark, a logo painted on brick. The hard part is not finding the text. The hard part is removing only the text and leaving the brick alone. Mask one pixel too wide and the inpainter starts inventing wall where there was wall already, and now the patch looks worse than the text did.

So the whole job comes down to one question: how tight is the mask. This is a walkthrough of how Wunjo Make builds that mask, straight from the code in visual_processing/segmentation/east/. I went in expecting one architecture and found a simpler one, so I will also tell you where the real code path differs from the version you might have read about.

The interesting part of text removal is not the detector. It is the few lines that turn boxes into a mask the inpainter can trust. Photo: Unsplash.

The shape of the problem

Think of the detector as a person with a highlighter. A bounding box is the highlighter held straight, so a slanted street sign gets a fat rectangle with a lot of wall caught in the corners. A rotated quadrilateral is the highlighter tilted to follow the text, so the marked area hugs the letters. EAST gives you the tilted highlighter. That tilt is the whole reason it is worth using here.

EAST stands for Efficient and Accurate Scene Text detector (Zhou et al., CVPR 2017). It is a single fully-convolutional network. No region proposals, no per-character steps. Feed it an image, get back two maps the same shape as a downscaled version of the image: a score map that says "text is here" and a geometry map that says "and the box around it looks like this."

Step 1. The network outputs a score and a geometry, not boxes

The model itself is small. A VGG16 backbone pulls features at several scales, a merge branch fuses them back up to a quarter of the input resolution, and an output head produces three things. Here is the head, from east/model.py:

def forward(self, x):
    score = self.sigmoid1(self.conv1(x))
    loc   = self.sigmoid2(self.conv2(x)) * self.scope          # 4 channels
    angle = (self.sigmoid3(self.conv3(x)) - 0.5) * math.pi     # 1 channel
    geo = torch.cat((loc, angle), 1)
    return score, geo

Read those three lines slowly, because they are the design.

score is one channel through a sigmoid, so every pixel gets a 0-to-1 number for "is this text." loc is four channels, each one a distance, scaled by self.scope which is 512. For a text pixel, those four numbers are the distance up, down, left, and right to the edges of its text box. angle is one channel mapped into the range minus pi-over-two to pi-over-two, the rotation of that box.

So the network never draws a box. It labels each text pixel with four distances and an angle. The box is implied. That is the trick that lets one convolution describe a rotated rectangle: a center pixel plus four reach values plus a twist.

Step 2. Rebuild rotated quads from the geometry

Now you have to turn "distance up, down, left, right, and an angle" back into four corner points. That is restore_polys in east/detect.py. For each pixel that scored high enough:

y_min = y - d[0, i]
y_max = y + d[1, i]
x_min = x - d[2, i]
x_max = x + d[3, i]
rotate_mat = get_rotate_mat(-angle[i])

temp_x = np.array([[x_min, x_max, x_max, x_min]]) - x
temp_y = np.array([[y_min, y_min, y_max, y_max]]) - y
coordidates = np.concatenate((temp_x, temp_y), axis=0)
res = np.dot(rotate_mat, coordidates)
res[0, :] += x
res[1, :] += y

In plain terms: build an axis-aligned rectangle from the four distances, move it so the text pixel is the origin, rotate it by the predicted angle, then move it back. Out come four corner points. The result row is the eight numbers of a quadrilateral, [x0,y0, x1,y1, x2,y2, x3,y3], the corners in order.

Two details that matter in practice. The score threshold and the candidate filter both live here:

xy_text = np.argwhere(score > score_thresh)   # score_thresh defaults to 0.9

A 0.9 threshold is high. It throws away anything the network is not very sure about. For text removal that is the right bias. A false positive means you erase a chunk of real image, which is much worse than missing a faint logo. And is_valid_poly quietly drops any quad that lands mostly outside the frame, so edge garbage does not become a mask.

Every high-scoring pixel proposes its own quad, so a single word produces hundreds of nearly identical quads. The code carries a confidence score in the ninth column of each box for exactly this reason. Worth knowing: in this code path the boxes go to the mask without a separate non-max-suppression merge, so you are relying on the high score threshold and the fact that overlapping filled polygons just paint the same white pixels twice. Painting white twice costs nothing.

Step 3. Rasterize the quads into a mask

Here is where the boxes stop being numbers and become a mask. plot_filled_mask:

def plot_filled_mask(img, boxes):
    mask = Image.new('RGB', img.size, (0, 0, 0))   # black canvas
    draw = ImageDraw.Draw(mask)
    if boxes is not None:
        for box in boxes:
            draw.polygon(
                [box[0], box[1], box[2], box[3],
                 box[4], box[5], box[6], box[7]],
                fill=(255, 255, 255))
    return mask

Black image the size of the frame, then every detected quad gets painted solid white. White means "inpaint this," black means "leave it." Because the quads are tilted, the white regions follow slanted and rotated text instead of boxing it in a wide rectangle. That is the wall-saving step. A straight bounding box around a diagonal sign would whitewash the triangles of brick in the corners, and the inpainter would happily repaint brick it did not need to touch.

The full detect call wraps the network and this rasterizer together, in SegmentText.detect_text:

def detect_text(self, img_path, max_resolution=1280):
    img = Image.open(img_path)
    w, h = img.size
    if max(w, h) > max_resolution:                # cap the long side
        scale = max_resolution / max(w, h)
        img = img.resize((int(w*scale), int(h*scale)))
    img, _, _ = resize_img(img)                   # make dims divisible by 32
    boxes = detect(img, self.model, self.device)
    filled_mask = plot_filled_mask(img, boxes)
    return filled_mask.resize((w, h))             # back to original size

The resize_img step is not cosmetic. The VGG backbone halves the resolution five times, so the input has to be divisible by 32 or the feature maps will not line up. The 1280 cap is there so a 4K frame does not blow up memory in the detector for no gain.

Step 4. The optional crop, for "only this corner"

If the user drew a box around just the subtitle area, you do not want to erase a logo elsewhere in the frame. SegmentText.crop handles that by keeping only the part of the mask inside the chosen rectangle:

@staticmethod
def crop(filled_mask, text_area: list = None):
    if text_area is not None:
        x1, y1, x2, y2 = text_area
        img = Image.new('L', filled_mask.size, color=0)    # all black
        region = filled_mask.crop((x1, y1, x2, y2))        # cut the user's box
        img.paste(region, (x1, y1))                        # paste it back, rest stays black
        return img
    return filled_mask

It does not crop the image down. It blacks out everything outside the user's rectangle and keeps the detected text inside it. So you still get tight quads, just restricted to the region you care about.

Step 5. Hand the mask to inpainting (and where SAM actually sits)

Now the part where I corrected my own mental model. I expected the EAST quads to become point or box prompts for SAM (Segment Anything, Kirillov et al., ICCV 2023), letting SAM tighten the mask further. That is a reasonable design. It is not this code.

In visual_processing/inference.py, the remove path is one function that handles two jobs. When the user clicks an object, it runs SAM. When the user asks to remove text, it runs EAST. They are siblings, not a chain. The text branch is gated behind use_remove_text, and SAM is explicitly skipped there:

if predictor is None and not use_remove_text:
    predictor = segmentation.init_vit(sam_vit_checkpoint, vit_model_type, device)
else:
    lprint(msg=f"[{method}] Model predictor will not be loaded "
               f"because of not need to use for text.", ...)

The text branch builds its own detector and writes the EAST mask straight to disk per frame:

segment_text = SegmentText(device=device, vgg16_path=vgg16_baseline_path, east_path=vgg16_east_path)
...
mask_text_frame = segment_text.detect_text(img_path=frame_file_path)
mask_textarea_frame = segment_text.crop(mask_text_frame, coord)
mask_textarea_frame.save(mask_text_frame_path)

So the real flow for text is EAST to mask to inpaint, no SAM in between. SAM is the tool for "remove this person," EAST is the tool for "remove this text," and they share plumbing because both of them produce a white-on-black mask that the same inpainter consumes.

The inpainter is chosen by source. Video on a capable GPU goes to ProPainter (Zhou et al., ICCV 2023), which propagates pixels across frames so the filled region stays consistent over time. Images, or anything on CPU, go to a single-frame inpainting model. Either way, the mask reaches it after one more cleanup. The white mask is dilated and binarized:

binary_mask = cv2.threshold(segment_mask, 128, 1, cv2.THRESH_BINARY)[1]
...
kernel = np.ones((kernel_size, kernel_size), np.uint8)   # kernel_size = 10
mask_to_save = cv2.dilate(mask_to_save, kernel, iterations=1)

The dilation by a 10-pixel kernel is deliberate. EAST hugs the glyphs tightly, but ink has a faint halo, anti-aliasing and shadow, just outside the letter shape. If the mask is exactly the glyph, the inpainter leaves a ghost outline. Growing the mask by a few pixels eats that halo. So the pipeline goes tight, then loosens by a known amount, instead of starting loose and hoping.

Finally the inpainter only writes inside the mask. From process_remover:

merge = generated * mask_raw_3ch + img_raw * (1 - mask_raw_3ch)

Generated pixels where the mask is white, original pixels everywhere else. That single line is the promise that the wall behind the text is the exact wall that was there. The model can hallucinate freely inside the white region and it will never touch a pixel outside it.

Gotchas you will actually hit

A few things I would warn a teammate about before they touch this path.

The 0.9 score threshold is conservative on purpose, and it will miss low-contrast or stylized text. If you lower it, you trade missed text for erased background. For removal, keep it high and let the user draw a text_area box for the stubborn cases instead of dropping the threshold globally.

Divisibility by 32 is silent. If you bypass resize_img and feed an arbitrary size, the merge branch upsamples will not match the skip connections and you get a shape error deep in the model. Always route through resize_img.

The mask is per frame, recomputed independently for each frame. There is no temporal smoothing on the text mask itself. For video, the consistency comes later, from ProPainter propagating the fill, not from the detector. If your text flickers in and out of detection, that flicker can show up before propagation evens it out.

Dilation cuts both ways. The 10-pixel grow is great for eating halos around small subtitles. On very large text it can swallow a sliver of background you wanted to keep. The kernel size is a constant in the code worth tuning for your content.

Takeaway

Removing text without removing the wall is a masking problem, and the mask quality is set in three short functions. The network emits a score map plus four distances and an angle per pixel. restore_polys rebuilds rotated quads from that geometry so the marked area follows slanted text. plot_filled_mask paints those quads white on black. Then a 10-pixel dilation eats the ink halo, and the inpainter writes only inside the white. SAM is in the same file but it is the object-removal sibling, not a stage in the text path.

The code is in visual_processing/segmentation/east/ in the Wunjo Make repo. If you are fighting ghost outlines after text removal, start at the dilation kernel and the score threshold. Those two constants decide whether the wall survives.

References

Zhou, Yao, Wen, et al. "EAST: An Efficient and Accurate Scene Text Detector." CVPR 2017. arXiv:1704.03155
Kirillov et al. "Segment Anything." ICCV 2023. arXiv:2304.02643
Zhou, Li, Loy, et al. "ProPainter: Improving Propagation and Transformer for Video Inpainting." ICCV 2023. arXiv:2309.03897

Why lip-sync blacks out the bottom half of the face before the model sees it

Wlad Radchenko — Wed, 15 Jul 2026 14:00:00 +0000

A lip-sync model has one job: draw a mouth that matches the audio. So the obvious input is a face and a sound. Feed in the picture, feed in the speech, get back a talking face.

That is not what the code feeds it. Before a single frame reaches the model, the bottom half of every face is set to zero. Pure black. The model never sees the real mouth it is supposed to reproduce.

The first time I read this in the data generator, I assumed it was a bug or some leftover debug code. It is neither. Blanking the lower face is the whole reason the model learns to lip-sync instead of learning to copy. This post walks through the few lines that build that input, in datagen, the generator that batches frames and audio for Wav2Lip.

Give the model the lit half and ask it to invent the rest. Photo: Unsplash.

The setup: frames in, mel chunks in

datagen is a generator. It takes a list of frame file paths, a list of mel-spectrogram chunks (the audio, sliced into short pieces), the target face size, and a batch size. It yields batches the model can run.

def datagen(self, frame_files, mels, img_size, wav2lip_batch_size, crop=None):
    img_batch, mel_batch, frame_batch, coords_batch = [], [], [], []
    face_det_results = self.face_detect_with_alignment_crop(frame_files, crop)

First it runs face detection once over all the frames, so it knows where every face sits. Then it walks the audio, not the video:

for i, m in enumerate(mels):
    idx = i % len(frame_files)

That idx = i % len(frame_files) is small but it sets the rhythm of the whole loop. The driver is the audio. There is one mel chunk per output frame, and each chunk needs a face to draw on. The modulo just wraps the frame index back to the start when the audio runs longer than the video. Twenty seconds of speech over a five-second clip? The clip loops. You never run out of faces, because the frame list is treated as circular.

Be precise about which list leads. If you have more audio than video, you get more output frames than input frames and the source clip repeats. If you have more video than audio, the loop stops when the mel chunks stop, and the tail of the clip is unused. The mel list is the clock.

Crop the face, resize to a square

For each audio chunk, grab the matching frame and the detected face box:

frame_to_save = cv2.imread(frame_files[idx])
dets = face_det_results[idx] if idx < len(face_det_results) else face_det_results[i % len(face_det_results)]
coords = dets[0].get("bbox") if len(dets) > 0 and dets[0] is not None and isinstance(dets[0], dict) else (0, 0, 0, 0)
x1, y1, x2, y2 = map(int, coords)

Clamp the box to the image, then either cut out the face and resize it to a square, or, if no face was found, resize the whole frame as a fallback:

if int(x2 - x1) == 0 and int(y2 - y1) == 0:
    img_batch.append(cv2.resize(frame_to_save, (img_size, img_size)))
else:
    face = frame_to_save[int(y1): int(y2), int(x1):int(x2)]
    img_batch.append(cv2.resize(face, (img_size, img_size)))

mel_batch.append(m)
frame_batch.append(frame_to_save)
coords_batch.append((y1, y2, x1, x2))

Notice it keeps two parallel things. img_batch holds the resized face crops the model will work on. frame_batch holds the original full frames, untouched, and coords_batch holds where each face was. The full frame and the coordinates are not for the model. They are for later, when the predicted mouth has to be pasted back into the real picture. The model only ever sees the square crop.

The trick: blank the lower half, then stack the original on top

Here is the part that matters. Once a batch is full, it gets assembled:

if len(img_batch) >= wav2lip_batch_size:
    img_batch, mel_batch = np.asarray(img_batch), np.asarray(mel_batch)

    img_masked = img_batch.copy()
    img_masked[:, img_size // 2:] = 0

    img_batch = np.concatenate((img_masked, img_batch), axis=3) / 255.
    mel_batch = np.reshape(mel_batch, [len(mel_batch), mel_batch.shape[1], mel_batch.shape[2], 1])

    yield img_batch, mel_batch, frame_batch, coords_batch
    img_batch, mel_batch, frame_batch, coords_batch = [], [], [], []

Three lines do the real work. Take them one at a time.

Copy the faces.

img_masked = img_batch.copy()

We need two versions of the same faces, so we copy first. The original stays intact.

Black out the bottom half of the copy.

img_masked[:, img_size // 2:] = 0

img_batch has shape (batch, height, width, channels). The slice [:, img_size//2:] is "every image, every row from the vertical middle downward." Setting it to 0 paints those rows black. The face is square and roughly centered, so the middle row falls around the nose. Everything below, the mouth, the jaw, the chin, goes dark. What survives is the top half: eyes, brow, hairline, head angle.

Think of it like a witness sketch. You give the artist the eyes and the shape of the head and say "the rest matched this voice, draw it." The model gets identity and pose from the top half, and it has to invent the mouth from the audio. It cannot cheat by reading the answer off the bottom of the picture, because the bottom of the picture is gone.

Stack the untouched face behind the masked one.

img_batch = np.concatenate((img_masked, img_batch), axis=3) / 255.

This is the line people miss. axis=3 is the channel axis. A normal image has 3 channels (BGR). Concatenating the masked face and the original face along that axis makes a 6-channel image: channels 0 to 2 are the masked face (top half only), channels 3 to 5 are the full, unmasked face.

So the model gets both at once, pixel-aligned:

The masked half tells it where to draw and whose face this is.
The full face is a reference. It is the same person, in this lighting, at this resolution, with this skin tone, just from a different moment in the clip.

Why hand over the full face if the point was to hide the mouth? Because there is a difference between hiding the answer for this frame and hiding what this person looks like. The masked input removes the mouth for the audio chunk being generated. The reference shows the model the subject's appearance so the painted mouth comes out the right color and sharpness and clearly belongs to that face. Wav2Lip (Prajwal et al., ACM Multimedia 2020) is built around this: the reference is a different timestep than the masked frame, so the network cannot just copy pixels straight across. It has to actually paint a new mouth in the right style. Same person, wrong moment. Learn to paint, not to copy.

The whole thing is divided by 255. at the end to put pixels in [0, 1].

The mel chunk gets a trailing channel dimension so it reads as a single-channel image too:

mel_batch = np.reshape(mel_batch, [len(mel_batch), mel_batch.shape[1], mel_batch.shape[2], 1])

Now both inputs are image-shaped tensors, and the batch is yielded.

The tail batch

The same masking and concatenation are repeated once more after the loop, for whatever is left over when the audio does not divide evenly into full batches:

if len(img_batch) > 0:
    img_batch, mel_batch = np.asarray(img_batch), np.asarray(mel_batch)
    img_masked = img_batch.copy()
    img_masked[:, img_size // 2:] = 0
    img_batch = np.concatenate((img_masked, img_batch), axis=3) / 255.
    ...
    yield img_batch, mel_batch, frame_batch, coords_batch

Duplicated logic, easy to overlook if you only read the main loop. If you ever change the masking, change it in both places, or your last partial batch will be built differently from every full one.

Gotchas you will actually hit

A few things bit me, or would have:

Six channels, not three. If you write your own model loader or swap checkpoints, the first conv layer expects 6 input channels. Hand it a 3-channel image and the shapes will not match. The doubling happens here, not in the model.
The masked half is exactly half, by a hardcoded fraction. img_size // 2 assumes the face is roughly centered and upright in the square crop. If your face detector returns a loose or tilted box, the cut can land on the lip line or the nose. The face detection and the masking are coupled even though they live in different functions.
Channel order. Crops are read with cv2.imread, so they are BGR. Both halves of the 6-channel tensor are BGR. Keep that consistent with whatever the checkpoint was trained on.
Audio length decides output length. Because the loop iterates mels and wraps the frame index with %, the number of output frames equals the number of mel chunks, not the number of input frames. Short video, long audio: the clip loops. That is by design, but it surprises people who expect output length to match the source video.
The full frame is not model input. frame_batch and coords_batch ride along only so the predicted mouth can be placed back later. Do not feed them to the network.

Wrap-up

The interesting decision in lip-sync input is not "give the model a face and some audio." It is "give the model a face with the mouth deleted, plus a clean reference of the same person, and make the audio fill the gap." The masking forces the model to generate from sound. The reference keeps the result looking like the right person. Two faces stacked into six channels, one of them half-erased.

The pasting-back and seam-hiding is a separate problem with its own tricks, covered in another post. This one stops at the input, because the input is where the lip-sync actually gets taught.

Code: portable/src/visual_processing/lip_sync/sync.py in Wunjo Make, the datagen method.

References

Prajwal, Mukhopadhyay, Namboodiri, Jawahar. "A Lip Sync Expert Is All You Need for Speech to Lip Generation In the Wild." ACM Multimedia 2020. arXiv:2008.10010. https://arxiv.org/abs/2008.10010
Chung, Zisserman. "Out of Time: Automated Lip Sync in the Wild." ACCV 2016 Workshops (SyncNet). https://www.robots.ox.ac.uk/~vgg/publications/2016/Chung16a/

Hiding the seam: making a lip-synced mouth stop looking pasted on

Wlad Radchenko — Tue, 14 Jul 2026 12:22:00 +0000

Lip-sync models like Wav2Lip give you a small, correct mouth. Then you paste that mouth back onto the face and the illusion collapses. You can see the rectangle. The skin tone is off by a shade, the lighting does not match, and the jaw flickers between frames. The network did its job. The compositing gave it away.

The interesting code in a lip-sync pipeline is not the model. It is everything that happens after the model, in the few dozen lines that make a generated patch belong to a real face. This is a walkthrough of that part.

The eye goes straight to the mouth on a talking face, which is exactly why a sloppy paste there is so obvious. Photo: Unsplash.

Where the seam comes from

Wav2Lip (Prajwal et al., ACM Multimedia 2020) takes a masked face plus an audio chunk and predicts the lower face for that audio. You get back a crop the size of the mouth box. The model never saw the full frame, the exact skin tone of the cheeks, or the lighting on the chin. So when you drop the crop in:

ff = f.copy()
ff[y1:y2, x1:x2] = p   # p is the predicted mouth, f is the original frame

you get a hard edge where the crop meets the face, a color step across that edge, and detail that does not match the resolution of the rest of the head. Three separate problems. The fix is three separate tools, stacked.

Tool 1: restore the detail (GFPGAN)

The raw prediction is a little soft. Before blending anything, the patched frame goes through GFPGAN (Wang et al., CVPR 2021), a face restoration GAN, to bring the mouth detail up to the same sharpness as the surrounding skin:

_, _, restored_img = self.gfpgan.enhance(
    ff, has_aligned=False, only_center_face=True, paste_back=True
)

only_center_face=True keeps it from wasting work on faces in the background. Now we have a sharp version of the whole face, but pasting that whole face back would change the eyes and cheeks too. We only want its mouth.

Tool 2: take only the mouth, with a soft mask

This is the step that separates a clean result from a visibly fake one. Instead of a rectangle, use a face-parsing network to cut out only the mouth and lips, then feather the edge.

# select the mouth / lip classes from the face parser
mm = [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 255, 255, 255, 0, 0, 0, 0, 0, 0]
mouth_mask = np.zeros_like(restored_img)
tmp_mask = self.face_enhancement.faceparser.process(restored_img[y1:y2, x1:x2], mm)[0]
tmp_mask_resized = cv2.resize(tmp_mask, (x2 - x1, y2 - y1))[:, :, np.newaxis] / 255.
mouth_mask[y1:y2, x1:x2] = tmp_mask_resized

That mm list is a class selector. The face parser labels every pixel (skin, eyes, nose, lips, and so on), and setting only a few entries to 255 keeps just the mouth region. The result is a mask shaped like an actual mouth, not a box. That shape is what lets the blend follow the lips instead of cutting across the cheek.

Tool 3: blend at every scale (Laplacian pyramid)

A single alpha blend along the mask edge still shows a line, because skin has detail at many scales and a hard cut mixes them all at once. The classic answer is a Laplacian pyramid blend (Burt and Adelson, 1983): split both images into frequency bands, blend each band with a correspondingly blurred mask, then add the bands back. Low frequencies merge over a wide area, high frequencies merge over a narrow one, and the seam disappears.

def laplacian_pyramid_blending_with_mask(A, B, m, num_levels=6):
    GA, GB, GM = A.copy(), B.copy(), m.copy()
    gpA, gpB, gpM = [GA], [GB], [GM]
    for _ in range(num_levels):                 # build Gaussian pyramids
        GA, GB, GM = cv2.pyrDown(GA), cv2.pyrDown(GB), cv2.pyrDown(GM)
        gpA.append(np.float32(GA)); gpB.append(np.float32(GB)); gpM.append(np.float32(GM))

    lpA, lpB, gpMr = [gpA[-1]], [gpB[-1]], [gpM[-1]]
    for i in range(num_levels, 0, -1):          # Laplacian = level minus upscaled next level
        lpA.append(np.subtract(gpA[i-1], cv2.pyrUp(gpA[i])))
        lpB.append(np.subtract(gpB[i-1], cv2.pyrUp(gpB[i])))
        gpMr.append(gpM[i-1])

    LS = [la * gm[:, :, None] + lb * (1 - gm[:, :, None])   # blend each band by the mask
          for la, lb, gm in zip(lpA, lpB, gpMr)]

    out = LS[0]                                  # reconstruct
    for i in range(1, num_levels):
        out = cv2.add(cv2.pyrUp(out), LS[i])
    return out

The pipeline runs this at num_levels=10 on 512x512 crops. Ten bands is more than a still photo needs, but video punishes any visible edge frame after frame, so it is worth the extra levels.

Tool 4: match the light (Poisson blending)

The pyramid blend fixes texture. It does not fix a global color or brightness difference between the patch and the face. For that the final step is Poisson blending (Pérez, Gangnet, Blake, SIGGRAPH 2003), which pastes by matching gradients instead of pixels. It keeps the detail of the mouth but reads the color and lighting from the surrounding face, so the patch takes on the skin tone around it.

f, _, _ = self.face_enhancement.process(pp, f, bbox=c, possion_blending=True)

After this, there is no patch. There is a face that happens to be saying something new.

The bonus: a running mean that kills jaw jitter

One more problem only shows up in motion. Face detection wobbles a pixel or two every frame, so the mouth box jumps, and a jumping box makes the chin vibrate even when the blend is perfect.

The fix is cheap. Keep a short history of recent mouth centers and reject any new detection that jumps too far from the running average:

center_current = np.array([(x1 + x2) / 2, (y1 + y2) / 2])
len_coords_mouth = 3      # history length
max_distance_mouth = 20   # pixels

if len(coords_mouth) < len_coords_mouth:
    coords_mouth.append(center_current)
    # ... do the full blend for this frame ...
else:
    distances = [calculate_distance(center_current, c) for c in coords_mouth]
    if np.mean(distances) < max_distance_mouth:
        # close enough: accept and blend
        ...
    coords_mouth.pop(0)               # slide the window
    coords_mouth.append(center_current)

Three frames of memory and a 20-pixel threshold. That is the entire anti-jitter system, and it removes most of the chin shake that makes cheap lip-sync look cheap.

How the model is even asked to invent a mouth

Worth one paragraph, because it explains why the blend has anything to blend. Before inference, the bottom half of the reference face is zeroed out and stacked with the unmasked face as extra channels:

img_masked = img_batch.copy()
img_masked[:, img_size // 2:] = 0
img_batch = np.concatenate((img_masked, img_batch), axis=3) / 255.

The model sees the top of the face for identity and pose, and a blank lower half it must fill from the audio. That is why the output is a plausible new mouth rather than a copy of the original one. It also deserves its own post, which is coming.

Takeaway

The model gives you a correct mouth. Making it a believable mouth is four steps of classical image work plus three frames of memory: restore detail, cut to a mouth-shaped soft mask, blend across scales with a Laplacian pyramid, match the light with Poisson, and stabilize the box over time. None of it is deep learning. All of it is the reason the result holds up.

The full function is in lip_sync/sync.py in the Wunjo Make repo. If you have been fighting a visible patch in your own pipeline, the mask shape and the Poisson step are where I would start.

References

Prajwal, Mukhopadhyay, Namboodiri, Jawahar. "A Lip Sync Expert Is All You Need for Speech to Lip Generation In the Wild." ACM MM 2020. arXiv:2008.10010
Wang, Li, Zhang, Shan. "Towards Real-World Blind Face Restoration with Generative Facial Prior." CVPR 2021. arXiv:2101.04061
Pérez, Gangnet, Blake. "Poisson Image Editing." ACM SIGGRAPH 2003.
Burt, Adelson. "The Laplacian Pyramid as a Compact Image Code." IEEE Trans. Communications, 1983.

Offline-first in the browser. How I keeping PWA and AI jobs alive across a refresh

Wlad Radchenko — Wed, 08 Jul 2026 11:24:00 +0000

A user types a prompt, clicks generate, and while the model is working they close the tab by accident. Or the wifi drops. Or they just hit refresh out of habit. In a lot of web apps, that is the end of the job and sometimes the end of their unsaved work too. The request was living in a variable in memory, and memory is gone.

The browser is a hostile place to keep state. Tabs close, networks flap, the page reloads on a whim. An editor that loses work on a refresh feels broken even when every other part is polished. This is how I think about durability in Wunjo Design, where the editor itself runs offline and only AI generation and cloud sync need the network. The snippets below are minimal and generic, just enough to show the shape.

The refresh is the cheapest test of whether your app respects the user's work. Photo: Unsplash.

Two things have to survive

When you say "offline-first," people hear "cache the assets." That is the small part. The real work is keeping two kinds of state alive across a reload:

The working document, so the canvas comes back exactly as the user left it. And in-flight AI jobs, so a generation that was running before the refresh is still running, or already done, when the page comes back. Lose either and the user feels the floor move.

The mental model I use is a kitchen ticket rail. When an order comes in, you clip the ticket to the rail before you start cooking. If the cook walks away, the ticket is still there for whoever comes back. The order does not live in the cook's head. It lives on the rail.

The rail is IndexedDB, not localStorage

localStorage is the wrong tool here: it is synchronous, small, and string-only, so it blocks the main thread and chokes on real documents. IndexedDB is the right one: asynchronous, large, and it stores structured data and blobs. A thin wrapper keeps it readable.

import { openDB } from 'idb'

const db = await openDB('editor', 1, {
  upgrade(db) {
    db.createObjectStore('docs')                 // the working document
    db.createObjectStore('jobs', { keyPath: 'id' }) // AI generation jobs
  },
})

Autosave the document on a throttle, not on every keystroke, so you are not hammering the disk:

const save = throttle(async () => {
  await db.put('docs', encodeDoc(), 'current')
}, 1000)

On startup, the first thing the app does is read docs/current back and rehydrate the canvas. The user never sees a blank page after a reload.

The job has to exist before the network call

This is the part people skip, and it is the whole point. If you call the AI endpoint first and plan to record the job when it answers, a refresh in between leaves you with a charge and no record. So you write the job down first, in IndexedDB, as a stub, and only then talk to the network.

// 1. record intent BEFORE the request
const job = { id: uuid(), prompt, status: 'queued', createdAt: Date.now() }
await db.put('jobs', job)

// 2. submit, then persist the remote handle
const { remoteId } = await api.submit(prompt)
await db.put('jobs', { ...job, status: 'running', remoteId })

Now the job is durable. The browser can die at any line and the worst case is a job stuck in queued or running, which is recoverable, instead of a job that silently vanished.

Resume is just reading the rail on boot

Because every job is a row in IndexedDB with a status, resuming after a reload is not special logic. It is a query.

async function resumeJobs() {
  const jobs = await db.getAll('jobs')
  for (const job of jobs) {
    if (job.status === 'running') pollUntilDone(job)   // pick up where we left off
    if (job.status === 'queued')  submitAndRun(job)    // never actually started
  }
}

One observer watches job state and updates the UI, so the canvas does not care how a job finishes, only that its status changed:

class JobObserver {
  constructor() { this.listeners = new Set() }
  subscribe(fn) { this.listeners.add(fn); return () => this.listeners.delete(fn) }
  async update(job) { await db.put('jobs', job); this.listeners.forEach(fn => fn(job)) }
}

That decoupling matters. "Is the job done" lives in one place; "what the canvas shows" lives in another; they talk through status changes. You can close the panel, reopen it, refresh the page, and the UI rebuilds itself from the persisted jobs every time.

Keep the main thread for drawing

Persisting, encoding, exporting, and diffing are not free, and the one thread you cannot afford to block is the one painting the canvas at 60 frames a second. Push the heavy work to a worker.

const worker = new Worker('./export.worker.js', { type: 'module' })
worker.postMessage({ type: 'export', doc: encodeDoc() })
worker.onmessage = (e) => downloadBlob(e.data.blob)

The editor stays responsive while a large export or a re-encode runs off to the side. Offline-first and smooth are not in tension; they both come from treating the main thread as sacred and everything else as background work.

About the author. I'm Wlad Radchenko, a software engineer. This article draws on building Wunjo Design, a browser-based AI vector editor that works offline, and the open-source Wunjo Make. Get in touch to find more on GitHub and LinkedIn.

Things that bite you in practice

Make resume idempotent. A generation costs money, so resuming a running job must re-attach to the existing remote job by its remoteId, never submit a fresh one. The "write the stub first" rule is what makes this possible: you always have the handle.

Respect the quota. IndexedDB has limits and the browser can evict your data under pressure. Store images as Blobs, not base64 data URLs, which inflate size by about a third, and prune finished jobs on a schedule.

Version your schema. The upgrade callback runs when you bump the database version. Plan migrations now, because shipping a change that silently drops a user's stored documents is the one bug you cannot apologize your way out of.

Treat the network as optional. Writes go to local storage first and sync when a connection exists. That is the local-first stance Kleppmann and colleagues argued for in 2019, and in a creative tool it is the difference between "my work is safe" and "I hope the save went through."

What this buys the product

The editor in Wunjo Design opens and works with no connection at all; only AI generation and cloud storage reach for the network. Your document survives a refresh, and a generation you kicked off is still there when you come back. That reliability is not a feature on the pricing page. It is the absence of a specific bad feeling, the one where you close a tab and your stomach drops. Building for the refresh is how you make sure that feeling never happens.

References

Kleppmann, Wiggins, van Hardenberg, McGranaghan. "Local-first software: you own your data, in spite of the cloud." ACM Onward! 2019. PDF
MDN Web Docs. "IndexedDB API" and "Using Web Workers." developer.mozilla.org

The sliding window that stops clicks in audio source separation

Wlad Radchenko — Tue, 07 Jul 2026 13:21:00 +0000

A vocal-separation model cannot eat a whole song at once. The network in Wunjo Make takes a fixed-size block of audio, runs an STFT, predicts a mask, runs an inverse STFT, and hands you back a block of the same size. So a four-minute track has to be cut into blocks, processed one at a time, and glued back together.

The naive way to glue blocks back is to lay them end to end. That clicks. Every seam becomes an audible tick, and across a song you get a rhythmic stutter that no listener will forgive.

This article is about the fix, which lives entirely inside one method, demix(), in portable/src/sound_processing/separator/mdx/model.py. The fix is a sliding window with overlap and a weighted average. I want to walk through it line by line with the real constants from the code, because the trick is small and the payoff is large.

A long take has to be cut into pieces before a model can touch it. Photo: Unsplash

Why cutting at hard edges clicks

Picture cutting a long ribbon into sections, processing each section, then taping them back in order. If the cuts are clean and the sections line up perfectly, fine. But a neural net does not return perfectly continuous audio at a block boundary. The sample at the very end of block 1 and the sample at the very start of block 2 were predicted by two separate forward passes that never saw each other. The waveform value can jump.

A jump in a waveform is a step. A step is a click. Your ear is very good at hearing steps, because a sudden discontinuity contains energy across the whole frequency spectrum. That is the tick you hear at every seam.

So you have two problems at once. The model output near a block edge is less reliable (the STFT has less context there), and butting two blocks together creates a discontinuity. Overlap-add solves both with the same move.

The constants, with real numbers

Before the loop, the model settings get computed in initialize_model_settings():

self.n_bins = self.n_fft // 2 + 1
self.trim = self.n_fft // 2
self.chunk_size = self.hop_length * (self.segment_size - 1)
self.gen_size = self.chunk_size - 2 * self.trim

The fixed knobs in the MDX separator are these:

self.segment_size = 256
self.overlap = 0.25
self.hop_length = 1024
self.batch_size = 1

Plug in the two that do not depend on the model file. The chunk size is hop_length * (segment_size - 1), which is 1024 * 255 = 261120 samples. At a 44100 Hz sample rate that is about 5.9 seconds of audio per block. So the model chews through the song roughly six seconds at a time.

trim is n_fft // 2, where n_fft comes from the model file (mdx_n_fft_scale_set). It is the padding the STFT needs on each side so the transform has context. gen_size is the useful interior of a chunk, chunk_size - 2 * trim, the part you actually keep after throwing away the padded edges.

Hold on to one number: chunk size is 261120 samples. Everything else is relative to it.

The window: a fade-in and fade-out in one array

Here is the plain version. Instead of taping the ribbon sections edge to edge, you let each section fade in at its start and fade out at its end, and you let neighboring sections overlap. Where two sections overlap, you add them. The fades are shaped so that the overlapping fades sum back to one. No seam, because there is no single point where one block stops and the next starts cold.

The shape that does this is a Hann window, np.hanning. It is a smooth bell that is zero at both ends and one in the middle.

window = None
if overlap != 0:
    window = np.hanning(chunk_size_actual)
    window = np.tile(window[None, None, :], (1, 2, 1))

chunk_size_actual is the length of the current chunk, which equals chunk_size for every chunk except the last one (more on that below). np.tile just copies the one-dimensional bell across both stereo channels so it lines up with the (1, 2, samples) shape of the audio.

The window is created fresh per chunk, keyed on the actual length, which matters for the short final chunk.

The step: how far the window slides

step = int((1 - overlap) * chunk_size)

With overlap = 0.25 and chunk_size = 261120, the step is int(0.75 * 261120) = 195840 samples. So each chunk starts 195840 samples after the previous one, but each chunk is 261120 samples long. The difference, 261120 - 195840 = 65280 samples, is the overlap region. At 44100 Hz that is about 1.5 seconds where two consecutive windowed chunks both contribute.

The loop walks the whole padded mixture in steps:

for i in tqdm(range(0, mixture.shape[-1], step)):
    start = i
    end = min(i + chunk_size, mixture.shape[-1])

start advances by step, end is start + chunk_size clamped to the array length. That clamp is what makes the last chunk shorter, and it is why the window is sized from chunk_size_actual = end - start rather than a constant.

The accumulators: result and divider

This is the heart of it. Two arrays, both zeroed and both the full length of the padded mixture:

result = np.zeros((1, 2, mixture.shape[-1]), dtype=np.float32)
divider = np.zeros((1, 2, mixture.shape[-1]), dtype=np.float32)

Think of result as the running sum of all the windowed, predicted audio, and divider as the running sum of all the window weights at each sample position. They grow together, chunk by chunk.

Inside the loop, after the model runs:

if window is not None:
    tar_waves[..., :chunk_size_actual] *= window
    divider[..., start:end] += window
else:
    divider[..., start:end] += 1

result[..., start:end] += tar_waves[..., : end - start]

So for each chunk: multiply the model output by the bell (fade in, fade out), add that into result at this position, and add the same bell into divider at this position. The fade makes the chunk contribute almost nothing at its edges and full weight in its middle. Where two chunks overlap, both their bells land in divider, so divider records exactly how much total weight piled up at every sample.

The else branch (divider += 1) is the no-overlap fallback. If overlap were zero, every sample gets weight one and you are back to plain concatenation. The code keeps that path, but the default overlap = 0.25 never takes it.

The normalization: divide by the weight you actually applied

After the loop:

epsilon = 1e-8  # A small constant to avoid division by zero
tar_waves = result / (divider + epsilon)

This one line is the whole point. You summed up weighted predictions in result. You summed up the weights in divider. Dividing gives you a weighted average at every single sample. A sample that only one chunk covered gets that chunk's value. A sample in the overlap region gets a blend of two chunks, weighted by how confident each window was there (the fading-out tail of the left chunk, the fading-in head of the right chunk).

Because the blend is continuous, the waveform is continuous. No step, no click.

The epsilon is a guard. At the very ends of the padded array the window can taper to zero, so divider can be zero there, and 0 / 0 is a NaN that would poison the audio. Adding 1e-8 first turns 0 / 0 into 0 / 1e-8 = 0, which is silence, which is exactly what you want at a sample no chunk really covered.

This pattern, sum weighted signal over sum of weights, is the classic weighted overlap-add (WOLA). It is the same family of idea that makes the STFT itself invertible. Here it is applied one level up, across model-sized blocks instead of FFT frames.

The padding and the trim, so the edges line up

Two more pieces make the bookkeeping exact. First, padding before the loop:

pad = gen_size + self.trim - ((mix.shape[-1]) % gen_size)
mixture = np.concatenate(
    (np.zeros((2, self.trim), dtype="float32"), mix, np.zeros((2, pad), dtype="float32")), 1)

Silence of length trim is prepended, and enough silence is appended to round the length up so the chunking covers everything. The leading trim of silence gives the very first real sample the same STFT context every interior sample gets, so the model is not guessing at the start.

Second, after normalization, the trim comes back off:

tar_waves_ = np.vstack(tar_waves_)[:, :, self.trim: -self.trim]
tar_waves = np.concatenate(tar_waves_, axis=-1)[:, : mix.shape[-1]]

[:, :, self.trim: -self.trim] drops the padded edges, and [:, : mix.shape[-1]] cuts the result back to the exact length of the input you handed in. You get out exactly as many samples as you put in, click-free.

Gotchas you will actually hit

A few things that bit me or would have:

The last chunk is shorter, and that is fine. Because end is clamped with min(...), the final chunk's length is whatever is left. The window is sized from chunk_size_actual = end - start, so the bell still fits. There is also a zero-pad just before the tensor conversion to make the model see a full-size block:

mix_part_ = mixture[:, start:end]
if end != i + chunk_size:
    pad_size = (i + chunk_size) - end
    mix_part_ = np.concatenate((mix_part_, np.zeros((2, pad_size), dtype="float32")), axis=-1)

but only the first end - start samples of the model output are accumulated (tar_waves[..., : end - start]), so the padding never leaks into the result.

Do not forget the epsilon. If you reimplement this and divide by divider directly, you will get NaNs at the tapered ends and your exported WAV will be full of garbage or silence-with-pops. The + 1e-8 is not decoration.

Overlap is a quality-versus-time dial, not free. With overlap = 0.25 the step is 75 percent of a chunk, so consecutive chunks share 25 percent. Push overlap to 0.5 and you process roughly twice as many chunks for smoother seams. The match-mix pass in the same method deliberately uses a tiny overlap = 0.02, because that pass is reconstructing the mixture, not the delicate vocal, and does not need the expensive smoothing:

if is_match_mix:
    chunk_size = self.hop_length * (self.segment_size - 1)
    overlap = 0.02

The window must match the chunk slice exactly. tar_waves[..., :chunk_size_actual] *= window and divider[..., start:end] += window both use the same length. If those two ever disagree, the divide stops being a true average and the seams come back. Keeping both keyed on chunk_size_actual is what keeps it honest.

Try it on your own audio

The whole mechanism is one loop and one divide. If you want to see it work, separate a track in Wunjo Make and listen to a long sustained note that crosses a chunk boundary. With overlap on it is smooth. Set overlap = 0 in your own copy and you will hear the ticks come back roughly every 4.4 seconds (the step length), which is the clearest proof that the window is doing the work.

The code is in portable/src/sound_processing/separator/mdx/model.py in the Wunjo Make repo. The MDX separation approach this builds on comes from KUIELab-MDX-Net.

References

Kim, Lee, Lee, Lee. "KUIELab-MDX-Net: A Two-Stream Neural Network for Music Demixing." MDX Workshop / ISMIR 2021. arXiv:2111.12203. https://arxiv.org/abs/2111.12203
Mitsufuji et al. "Music Demixing Challenge 2021." Frontiers in Signal Processing, 2022. arXiv:2108.13559. https://arxiv.org/abs/2108.13559
Défossez. "Hybrid Spectrogram and Waveform Source Separation." (Demucs v3.) 2021. arXiv:2111.03600. https://arxiv.org/abs/2111.03600

One model, two stems: how a vocal remover gets the instrumental for free

Wlad Radchenko — Wed, 01 Jul 2026 11:10:00 +0000

You feed a song in. You get two files back: one with only the singer, one with only the band. The obvious way to build that is to train two models, one per stem. The audio separator in my open-source app does something cheaper and, in practice, cleaner: it runs one model, then gets the second stem by subtraction.

This post walks through the actual code that does it. It is short, it is specific, and once you see the subtraction trick you will use it in other places too.

A mixing desk keeps every instrument on its own channel. Source separation tries to recover those channels after they have already been bounced down to one. Photo: Unsplash.

The mental model: a song is a stack of transparent sheets

Picture the final track as a stack of clear plastic sheets pressed together. One sheet has the vocals printed on it, the others have drums, bass, guitar. When you hear an MP3, you are looking down through the whole stack at once.

Separation tries to lift one sheet off. Here is the part people miss: if you can lift the vocals sheet cleanly, you do not need a separate model to read the instrumental. Whatever is left on the table after the vocals are gone is the instrumental. Lift one, subtract, keep both.

That is the whole idea. The rest is making "lift the vocals" and "subtract" precise enough to sound good.

Step 1. Stop working with audio, start working with a picture

You cannot un-mix a waveform by staring at the wiggly line. So the first move is a Short-Time Fourier Transform (STFT), which turns the audio into a spectrogram: a heatmap where one axis is time, the other is pitch, and brightness is how loud that pitch is at that moment.

The useful framing: a spectrogram is an editable image. A vocal sits in a recognizable band with recognizable shapes. A model that edits images can learn to erase it.

class STFT:
    def __init__(self, n_fft, hop_length, dim_f, device):
        self.n_fft = n_fft            # 6144 here: window length
        self.hop_length = hop_length  # 1024: how far the window slides
        self.dim_f = dim_f            # 3072: how many frequency bins we keep
        self.hann_window = torch.hann_window(window_length=n_fft, periodic=True)

    def __call__(self, x):
        spec = torch.stft(x, n_fft=self.n_fft, hop_length=self.hop_length,
                          window=self.hann_window.to(x.device),
                          center=True, return_complex=True)
        spec = torch.view_as_real(spec)            # split into real + imaginary
        spec = spec.permute(...).reshape(...)      # pack as image-like channels
        return spec[..., : self.dim_f, :]          # keep only the first dim_f bins

Two numbers there matter more than they look.

n_fft = 6144 sets how many pitch bins exist: n_fft // 2 + 1 = 3073. hop_length = 1024 is the stride, the same idea as stride in a CNN, just along time. Smaller hop means more overlap and a smoother picture at the cost of more compute.

The frequency cut-off (the part that is not obvious)

Look at the last line: return spec[..., : self.dim_f, :]. The STFT produces 3073 frequency bins, but the model only ever sees 3072 of them. The top of the spectrum is thrown away before inference and padded back with zeros on the way out.

This is not a bug, it is a deliberate trick from the KUIELab-MDX-Net paper (Kim et al., ISMIR 2021). You cut the frequency range to the band where the target source actually lives, which lets you spend the model's resolution where it counts instead of on bins that carry almost no energy. The inverse transform pads the missing bins back so the shapes line up:

def pad_frequency_dimension(self, x, ...):
    freq_padding = torch.zeros([..., num_freq_bins - freq_dim, time_dim]).to(x.device)
    return torch.cat([x, freq_padding], -2)   # put the cut bins back as zeros

One more small thing the code does right before inference:

spek = self.stft(mix)
spek[:, :, :3, :] *= 0   # zero the lowest 3 bins

That kills the bottom three frequency bins, the sub-bass rumble and DC offset that add nothing to a vocal and only confuse the model. Three lines, real quality difference.

Step 2. Run the model on the picture

The model is shipped as ONNX, so inference is one call. It takes the spectrogram and returns a predicted spectrogram for the primary stem.

def run_model(self, mix):
    spek = self.stft(mix)
    spek[:, :, :3, :] *= 0
    spec_pred = self.model_run(spek)                 # ONNX forward pass
    return self.stft.inverse(spec_pred)              # back to a waveform

self.model_run is just the ONNX session wrapped in a lambda, which is the entire reason this runs fast on a CPU when no GPU is around:

session = ort.InferenceSession(model_path, providers=providers, sess_options=opts)
self.model_run = lambda spek: session.run(None, {"input": spek.cpu().numpy()})[0]

Step 3. The subtraction that hands you the second stem

Now the payoff. The model gave us one stem. The other one is mixture minus that stem. The code offers two ways to do it.

The plain way is waveform subtraction:

# source = the stem the model predicted
self.secondary_source = mix.T - source.T

The careful way is spectral subtraction, used when you want to suppress bleed more aggressively. It subtracts in the frequency domain and keeps the phase of the original mix:

def invert_audio(specs, invert_p=True):
    X_mag = np.abs(specs[0])                  # mixture magnitude
    y_mag = np.abs(specs[1])                  # predicted stem magnitude
    max_mag = np.where(X_mag >= y_mag, X_mag, y_mag)
    return specs[1] - max_mag * np.exp(1.0j * np.angle(specs[0]))

Both work. Waveform subtraction is the default because it is exact and fast.

The subtlety that makes subtraction sound clean

Here is the detail I am proud of, and it is easy to get wrong. If you subtract the model's clean output from the raw input, the reconstruction artifacts from the model's own STFT and inverse STFT do not cancel. You hear them as a faint warble in the second stem.

The fix is to subtract through the same lossy pipeline. Before subtracting, the code re-runs the mix through the transform in a pass-through mode that applies STFT and inverse STFT but skips the model:

raw_mix = self.demix(mix, is_match_mix=True)   # STFT -> ISTFT, no model

if self.invert_using_spec:
    self.secondary_source = invert_stem(raw_mix, source)
else:
    self.secondary_source = mix.T - source.T

In that pass-through mode, run_model just returns the spectrogram untouched:

if is_match_mix:
    spec_pred = spek.cpu().numpy()   # identity: no model, only transform round-trip

So both sides of the subtraction carry the same transform artifacts, and the artifacts cancel instead of leaking into the result. Small idea, audible difference.

The 1.022 fudge factor

One last constant worth knowing about:

if not is_match_mix:
    source *= self.compensate   # compensate = 1.022

The model very slightly under-predicts magnitude, so the output is multiplied by 1.022 before it is written. It is a calibration constant that ships with the model weights, not a magic number I picked. If your separated stem sounds a hair quiet, this is where it lives.

Gotchas if you build this yourself

A few things that cost me time:

The audio must be stereo before it hits the model. Mono gets duplicated into two channels (np.asfortranarray([mix, mix])), otherwise the channel math downstream breaks.
Normalize before and after. The code normalizes to a peak of 0.9 so the int16 export does not clip.
The ONNX path here assumes segment_size == dim_t (256 in this model). If they differ you have to fall back to an onnx2torch path, which the code guards against with an explicit error rather than failing silently.

I left out one piece on purpose: how the song is cut into overlapping chunks and stitched back without clicks at the seams. That sliding-window-and-Hann-window dance deserves its own walkthrough, and it is next in this series.

Takeaway

You do not need one model per stem. Train one to lift the hardest source, get the rest by subtraction, and spend your effort on making the subtraction honest: same transform on both sides, phase preserved, magnitude calibrated. The full file is separator/mdx/model.py in the repo if you want to read past the snippets.

If you try this on your own audio, tell me what the second stem sounds like before and after the match-mix trick. That difference is the whole point.

References

Kim, Lee, Lee, Lee. "KUIELab-MDX-Net: A Two-Stream Neural Network for Music Demixing." ISMIR 2021. arXiv:2111.12203
Mitsufuji et al. "Music Demixing Challenge 2021." Frontiers in Signal Processing, 2022. arXiv:2108.13559
Défossez. "Hybrid Spectrogram and Waveform Source Separation." (Demucs v3.) 2021. arXiv:2111.03600

Version history and live collaboration in a browser editor, without a server lock

Wlad Radchenko — Tue, 30 Jun 2026 12:15:00 +0000

Two people open the same canvas. One drags a rectangle, the other recolors a circle, both while a third person is offline on a train editing the same file. When everyone syncs, the document has to agree on one result, nobody's work can vanish, and you still want a history you can scrub back through. The lazy answer is a server lock: one editor at a time, everyone else read-only. It is easy to build and miserable to use.

I want to walk through the data model that makes the hard version work, the one behind real-time collaboration, offline editing, and version history in Wunjo Design. None of this is framework magic. It is one idea applied carefully, and you can reason about all of it from first principles. I will keep the code to small, generic snippets so the shape is clear.

Multiplayer editing looks like a UI feature. It is really a question about how two histories merge into one. Photo: Unsplash.

The real problem: merging two edits with no referee

Picture the document as a plain object: shapes keyed by id. Now two clients change it at the same time, each against their own copy. You receive both changes out of order, maybe minutes apart. What is the final state?

If you just replay "last write wins" on whole objects, the recolor clobbers the move, or the other way around, depending on which packet arrived last. Work disappears. If you funnel every edit through a central server that serializes them, you are back to a lock with extra steps, and offline editing is impossible because there is no server to ask.

The question underneath multiplayer is not "how do I send changes over a socket." It is "given two divergent versions of the same document, how do I merge them so that every client lands on the identical result, no matter what order the changes arrive in."

Two families of answer

There are two well-known approaches.

Operational Transformation (OT) is what classic Google Docs uses. Each edit is an operation, and when operations conflict you transform one against the other so they still compose. It works, but the transformation functions are notoriously fiddly, and most designs lean on a central server to order operations.

Conflict-free Replicated Data Types (CRDTs) take a different route. You design the data structure so that merging is commutative, associative, and idempotent by construction. Apply the same set of changes in any order, any number of times, and every replica converges to the same state. Shapiro and colleagues formalized this as strong eventual consistency in 2011. The practical upshot: no central referee is required. Any peer can merge any other peer's changes and they all end up equal.

I went with the CRDT route because it gives offline editing and history for free instead of as extra subsystems. The algorithm behind the JavaScript CRDT libraries in this space (Yjs, via the YATA algorithm) is built exactly for editing arbitrary structured data in a browser.

Modeling a canvas as shared types

The trick is to stop thinking of the document as one big blob and model it as nested CRDT containers: a map of shapes, where each shape is itself a map of properties. A minimal sketch with the Yjs API:

import * as Y from 'yjs'

const doc = new Y.Doc()
const shapes = doc.getMap('shapes')         // the whole scene

function addRect(id) {
  const rect = new Y.Map()
  rect.set('type', 'rect')
  rect.set('x', 10); rect.set('y', 20)
  rect.set('w', 100); rect.set('h', 60)
  rect.set('fill', '#3366ff')
  shapes.set(id, rect)
}

Granularity is the whole design decision here. Because each shape is its own map, two people editing different shapes never touch the same data, so there is nothing to conflict. And because each property is its own key, one person moving a rectangle (x, y) while another recolors it (fill) merges cleanly: the move and the recolor are independent keys, both survive. You only get a real conflict when two people set the same property of the same shape at the same instant, and there the CRDT picks a deterministic winner that every client agrees on.

Sync is just bytes you can apply in any order

Once the document is a CRDT, "sync" stops being a protocol you design and becomes two function calls. A client serializes its changes to a compact binary update and ships it. Any peer applies it. Order does not matter.

// on the client that made changes
const update = Y.encodeStateAsUpdate(doc)   // a small binary diff

// on any other client, whenever it arrives
Y.applyUpdate(otherDoc, update)             // merge; commutative and idempotent

That applyUpdate is the entire merge engine. Apply the same update twice, it is a no-op. Apply two clients' updates in opposite orders on two machines, both machines match. This is the property that makes the network boring, which is what you want a network to be.

Offline falls out for free

Here is the part I find satisfying. Offline editing is not a separate feature you bolt on. When a client is offline, it keeps producing updates against its local document. When it reconnects, those queued updates merge with everyone else's through the same applyUpdate call. There is no special "sync conflict resolution" screen, because the data type cannot produce an unresolvable conflict. This is the local-first idea that Kleppmann and colleagues argued for in 2019: the network is an enhancement, not a requirement, and your data lives on your device first.

History without saving fifty copies

Version history sounds like it needs a pile of saved files. It does not. A CRDT document already carries the information needed to reconstruct earlier states, because it is fundamentally a log of changes plus enough metadata to order them. You take lightweight snapshots at points worth remembering, and you can reconstruct or restore the document as of any snapshot.

const checkpoint = Y.snapshot(doc)   // a cheap marker, not a full copy
// later: rebuild the document state as of `checkpoint` for preview or restore

So "version history with rollback" becomes: keep a list of snapshots with labels, render a preview from any of them, and to restore, apply the delta that takes you back. You are storing operations and markers, not duplicate documents, which is why the history can be deep without bloating storage.

About the author. I'm Wlad Radchenko, a software engineer. This article draws on building Wunjo Design, a browser-based AI vector editor with collaboration and version history, and the open-source Wunjo Make. Get in touch to find more on GitHub and LinkedIn.

Things that bite you in practice

A few lessons that cost me time.

Keep transient state out of the document. Cursors, selections, and who-is-online change many times a second and mean nothing tomorrow. Put them on an ephemeral presence channel (Yjs calls it awareness), not in the CRDT, or your history fills with mouse twitches.

Do not store images inside the CRDT. Big binary blobs make every sync and every snapshot heavier. Store the blob separately, keyed by id, and keep only the reference in the document.

Undo has to be per-user and merge-aware. In a shared document, your undo must reverse your last action, not your collaborator's. A naive global undo stack will happily undo someone else's work. The CRDT libraries track the origin of each change so undo can be scoped, and that is its own topic, which I will take apart next.

Compact the history. An append-only log grows. Periodic snapshots plus garbage collection of superseded operations keep it bounded.

What this buys the product

One data model gives Wunjo Design three features that usually cost three subsystems: real-time collaboration, offline editing, and version history you can roll back. They are not separate engineering efforts stapled together. They are consequences of choosing a data structure that merges by construction. That is the kind of decision that does not show up in the UI but determines whether the UI is even possible.

If you are building anything collaborative in the browser, start from the data model, not the socket. Pick a structure where merge is total, and the network, the offline support, and the history mostly write themselves.

References

Shapiro, Preguiça, Baquero, Zawirski. "Conflict-free Replicated Data Types." SSS 2011. HAL
Nicolaescu, Jahns, Derntl, Klamma. "Near Real-Time Peer-to-Peer Shared Editing on Extensible Data Types." ACM GROUP 2016. (The YATA algorithm behind Yjs.)
Kleppmann, Wiggins, van Hardenberg, McGranaghan. "Local-first software: you own your data, in spite of the cloud." ACM Onward! 2019. PDF