Skip to content

fix(inpainting): keep LanPaint sigma on the input device - #2099

Open
acesof wants to merge 1 commit into
deepbeepmeep:mainfrom
acesof:fix/lanpaint-device-mismatch
Open

acesof wants to merge 1 commit into
deepbeepmeep:mainfrom
acesof:fix/lanpaint-device-mismatch

Conversation

@acesof

@acesof acesof commented Aug 7, 2026

Copy link
Copy Markdown

Problem

Inpainting crashes on any non-CPU device that isn't the default, e.g. Apple Silicon:

RuntimeError: Expected all tensors to be on the same device,
but found at least two devices, mps:0 and cpu!

Traceback:

File "models/qwen/pipeline_qwenimage.py", line 979, in __call__
  latents = lanpaint_proc(denoise, cfg_predictions, true_cfg_scale, 1., latents,
                          original_image_latents, randn, t/1000, image_mask_latents, ...)
File "shared/inpainting/lanpaint.py", line 87, in __call__
  out = self.LanPaint(...)
File "shared/inpainting/lanpaint.py", line 142, in LanPaint
  noisy_image = self.latent_image * (1.0 - sigma) + self.noise * sigma

Cause

The pipeline passes t/1000, which is a 0-dim tensor, so this branch fires:

if len(sigma.shape) == 0:
    sigma = torch.tensor([sigma.item()])   # no device= -> always CPU

torch.tensor(...) without device= allocates on CPU regardless of where
sigma came from. Line 142 then multiplies device-resident latents by that CPU
tensor.

Note this is a shape-(1,) tensor, not a 0-dim one, so it does not qualify for
PyTorch's CPU-scalar promotion and the op genuinely fails.

Fix

-            sigma = torch.tensor([sigma.item()])
+            sigma = torch.tensor([sigma.item()], device=sigma.device)

I used device= rather than sigma.unsqueeze(0) deliberately:
torch.tensor([x.item()]) infers float32 regardless of the input dtype, so
adding device= changes placement only and leaves numerics identical.
unsqueeze(0) would preserve bf16 and silently change the precision of the
sigma schedule downstream.

shared/inpainting/ was checked for the same pattern — this is the only
device-less tensor construction. utils.py:270,280,293 already pass
device=y0.device.

Verification

Apple M4 Pro / MPS, torch 2.13.0, qwen_image_20B, 512x512,
video_prompt_type: "VA", model_mode: 2, 3 steps.

Before: crashes at the first LanPaint call.
After: completes in ~1m07s, exit 0, no errors in the log.

To confirm the mask is genuinely respected rather than the run merely not
crashing, mean absolute pixel difference from the source image:

region mean abs diff
inside mask 51.81
outside mask 0.96

The masked region is regenerated; the rest is preserved (0.96 is JPEG noise).

Not exercised: the IS_FLUX / IS_FLOW branches and Flux/Wan inpainting. The
changed line runs before that branch, so they are covered by the same fix, but I
only ran the Qwen path.

🤖 Generated with Claude Code

The pipeline passes a 0-dim sigma (`t/1000`), so the `len(sigma.shape) == 0`
branch rebuilds it with `torch.tensor([sigma.item()])`, which has no `device`
argument and therefore always lands on CPU.

A few lines later:

    noisy_image = self.latent_image * (1.0 - sigma) + self.noise * sigma

mixes that CPU tensor with device-resident latents and raises:

    RuntimeError: Expected all tensors to be on the same device,
    but found at least two devices, mps:0 and cpu!

Reproduced on Apple Silicon (MPS) with qwen_image_20B inpainting; the same
mismatch applies to any non-CPU device.

Passing `device=sigma.device` changes placement only. `torch.tensor([x.item()])`
infers float32 regardless of the input dtype, so the numerics are unchanged --
unlike `sigma.unsqueeze(0)`, which would preserve bf16 and silently alter the
precision of the sigma schedule.

Verified end-to-end on MPS (qwen_image_20B, 512x512, video_prompt_type "VA",
model_mode 2): completes cleanly, and the mask is correctly respected -- mean
absolute difference from the source image is 51.8 inside the mask vs 0.96
outside.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
@acesof acesof closed this Aug 7, 2026
@acesof acesof reopened this Aug 7, 2026
@acesof

acesof commented Aug 7, 2026

Copy link
Copy Markdown
Author

Reopening, closed by mistake.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant