42 lines
1.7 KiB
Python
42 lines
1.7 KiB
Python
"""
|
|
Learnable self-denoiser refinement pass, ported from Apple ml-starflow
|
|
utils/inference.py::apply_denoising (the `learnable_self_denoiser` branch)
|
|
at commit 20984834c1a68ee679faa8ed9cda2842f02b042d.
|
|
|
|
Upstream's process_denoising() runs this unconditionally whenever
|
|
finetuned_vae == 'none' and vae_adapter is None and return_sequence == 0 --
|
|
which is always true for this dependency-light suite (we support neither a
|
|
finetuned VAE nor a DiT VAE adapter, and never return the full block
|
|
sequence), and always calls apply_denoising with sigma_next defaulting to 0,
|
|
which is the only code path implemented here.
|
|
|
|
With sigma_next=0, the update simplifies exactly to:
|
|
x_denoised = x - denoiser_output * NOISE_STD_CONST
|
|
since (sigma_curr - 0) / sigma_curr == 1 regardless of sigma_curr's value --
|
|
this is why NOISE_STD_CONST (0.3) is a fixed constant, independent of the
|
|
model's configured noise_std (0.5).
|
|
"""
|
|
import torch
|
|
|
|
NOISE_STD_CONST = 0.3
|
|
|
|
|
|
@torch.no_grad()
|
|
def apply_learnable_denoiser(transformer, x: torch.Tensor, y_text_embed: torch.Tensor) -> torch.Tensor:
|
|
window, overlap = 21, 10
|
|
if x.shape[1] <= window:
|
|
denoiser_output = transformer(x, y_text_embed, denoiser=True)
|
|
return x - denoiser_output * NOISE_STD_CONST
|
|
result = torch.empty_like(x)
|
|
start = 0
|
|
while start < x.shape[1]:
|
|
end = min(start + window, x.shape[1])
|
|
chunk = x[:, start:end]
|
|
denoised = chunk - transformer(chunk, y_text_embed, denoiser=True) * NOISE_STD_CONST
|
|
write_start = start if start == 0 else start + overlap
|
|
result[:, write_start:end] = denoised[:, write_start - start:]
|
|
if end == x.shape[1]:
|
|
break
|
|
start = end - overlap
|
|
return result
|