Merge branch 'main' into flux3-image
This commit is contained in:
Jaret Burkett committed
commit
9390ca463e
4 files changed
+160
-2
No files matched your search
@@ -36,7 +36,7 @@ from diffusers import EMAModel
|
|||||||
import math
|
import math
|
||||||
from toolkit.train_tools import precondition_model_outputs_flow_match
|
from toolkit.train_tools import precondition_model_outputs_flow_match
|
||||||
from toolkit.models.diffusion_feature_extraction import DiffusionFeatureExtractor, load_dfe
|
from toolkit.models.diffusion_feature_extraction import DiffusionFeatureExtractor, load_dfe
|
||||||
from toolkit.util.losses import wavelet_loss, stepped_loss
|
from toolkit.util.losses import wavelet_loss, stepped_loss, band_loss
|
||||||
import torch.nn.functional as F
|
import torch.nn.functional as F
|
||||||
from toolkit.unloader import unload_text_encoder
|
from toolkit.unloader import unload_text_encoder
|
||||||
from PIL import Image
|
from PIL import Image
|
||||||
@@ -952,6 +952,9 @@ class SDTrainer(BaseSDTrainProcess):
|
|||||||
loss = torch.nn.functional.l1_loss(pred.float(), target.float(), reduction="none")
|
loss = torch.nn.functional.l1_loss(pred.float(), target.float(), reduction="none")
|
||||||
elif self.train_config.loss_type == "wavelet":
|
elif self.train_config.loss_type == "wavelet":
|
||||||
loss = wavelet_loss(pred, batch.latents, noise)
|
loss = wavelet_loss(pred, batch.latents, noise)
|
||||||
|
elif self.train_config.loss_type == "band":
|
||||||
|
# per-sample frequency-weighted MSE: bands the noise has buried at this t drop out
|
||||||
|
loss = band_loss(pred, target, batch.latents, noise, timesteps, self.sd.is_flow_matching, self.sd.noise_scheduler)
|
||||||
elif self.train_config.loss_type == "stepped":
|
elif self.train_config.loss_type == "stepped":
|
||||||
loss = stepped_loss(pred, batch.latents, noise, noisy_latents, timesteps, self.sd.noise_scheduler)
|
loss = stepped_loss(pred, batch.latents, noise, noisy_latents, timesteps, self.sd.noise_scheduler)
|
||||||
# the way this loss works, it is low, increase it to match predictable LR effects
|
# the way this loss works, it is low, increase it to match predictable LR effects
|
||||||
|
|||||||
@@ -522,7 +522,7 @@ class TrainConfig:
|
|||||||
self.correct_pred_norm = kwargs.get('correct_pred_norm', False)
|
self.correct_pred_norm = kwargs.get('correct_pred_norm', False)
|
||||||
self.correct_pred_norm_multiplier = kwargs.get('correct_pred_norm_multiplier', 1.0)
|
self.correct_pred_norm_multiplier = kwargs.get('correct_pred_norm_multiplier', 1.0)
|
||||||
|
|
||||||
self.loss_type = kwargs.get('loss_type', 'mse') # mse, mae, wavelet, pixelspace, mean_flow, pseudo_huber
|
self.loss_type = kwargs.get('loss_type', 'mse') # mse, mae, wavelet, band, pixelspace, mean_flow, pseudo_huber
|
||||||
|
|
||||||
# do the loss on a timestep to 0 prediction
|
# do the loss on a timestep to 0 prediction
|
||||||
self.t0_loss_target = kwargs.get('t0_loss_target', False)
|
self.t0_loss_target = kwargs.get('t0_loss_target', False)
|
||||||
|
|||||||
@@ -0,0 +1,94 @@
|
|||||||
|
<svg xmlns="http://www.w3.org/2000/svg" width="960" height="620" viewBox="0 0 960 620" font-family="Helvetica, Arial, sans-serif">
|
||||||
|
<defs>
|
||||||
|
<pattern id="fine" width="4" height="4" patternUnits="userSpaceOnUse"><rect width="2" height="2" fill="#555"/><rect x="2" y="2" width="2" height="2" fill="#555"/></pattern>
|
||||||
|
<pattern id="mid" width="14" height="14" patternUnits="userSpaceOnUse"><rect width="7" height="7" fill="#777"/><rect x="7" y="7" width="7" height="7" fill="#777"/></pattern>
|
||||||
|
<pattern id="coarse" width="56" height="56" patternUnits="userSpaceOnUse"><rect width="28" height="28" fill="#999"/><rect x="28" y="28" width="28" height="28" fill="#999"/></pattern>
|
||||||
|
<linearGradient id="low" x1="0" y1="0" x2="1" y2="1"><stop offset="0" stop-color="#8a8a8a"/><stop offset="1" stop-color="#d5d5d5"/></linearGradient>
|
||||||
|
<marker id="arr" markerWidth="10" markerHeight="10" refX="9" refY="5" orient="auto"><path d="M0,0 L10,5 L0,10 z" fill="#444"/></marker>
|
||||||
|
<style>
|
||||||
|
.t { font-size: 20px; font-weight: bold; fill: #222; }
|
||||||
|
.h { font-size: 14px; font-weight: bold; fill: #222; }
|
||||||
|
.b { font-size: 13px; fill: #333; }
|
||||||
|
.s { font-size: 11px; fill: #444; }
|
||||||
|
.m { font-size: 12px; font-family: Menlo, Consolas, monospace; fill: #222; }
|
||||||
|
.box { fill: #f7f7f7; stroke: #bbb; stroke-width: 1.5; }
|
||||||
|
.ar { stroke: #444; stroke-width: 2; fill: none; marker-end: url(#arr); }
|
||||||
|
</style>
|
||||||
|
</defs>
|
||||||
|
<rect width="960" height="620" fill="#ffffff"/>
|
||||||
|
<text x="24" y="34" class="t">band loss: frequency-weighted residual (toolkit/util/losses.py: band_loss)</text>
|
||||||
|
<text x="24" y="56" class="b">The plain MSE penalizes detail the noise has already destroyed. This loss keeps each frequency band of the residual</text>
|
||||||
|
<text x="24" y="74" class="b">only as far as that band is still recoverable at the sample's noise level t.</text>
|
||||||
|
|
||||||
|
<!-- step 1: residual -->
|
||||||
|
<rect x="24" y="90" width="150" height="120" class="box"/>
|
||||||
|
<text x="99" y="112" text-anchor="middle" class="h">1. residual</text>
|
||||||
|
<text x="99" y="134" text-anchor="middle" class="m">pred - target</text>
|
||||||
|
<text x="99" y="156" text-anchor="middle" class="s">x0, velocity or epsilon:</text>
|
||||||
|
<text x="99" y="170" text-anchor="middle" class="s">one picture scaled by t,</text>
|
||||||
|
<text x="99" y="184" text-anchor="middle" class="s">no conversion needed</text>
|
||||||
|
<path d="M176,150 L212,150" class="ar"/>
|
||||||
|
|
||||||
|
<!-- step 2: bands -->
|
||||||
|
<rect x="216" y="90" width="310" height="200" class="box"/>
|
||||||
|
<text x="371" y="112" text-anchor="middle" class="h">2. split into Laplacian bands</text>
|
||||||
|
<rect x="236" y="128" width="56" height="56" fill="url(#fine)" stroke="#d9534f" stroke-width="2"/><rect x="308" y="128" width="56" height="56" fill="url(#mid)" stroke="#f0a030" stroke-width="2"/><rect x="380" y="128" width="56" height="56" fill="url(#coarse)" stroke="#3c9d5d" stroke-width="2"/><rect x="452" y="128" width="56" height="56" fill="url(#low)" stroke="#3b7ddd" stroke-width="2"/>
|
||||||
|
<text x="264" y="202" text-anchor="middle" class="s" fill="#d9534f">fine</text>
|
||||||
|
<text x="336" y="202" text-anchor="middle" class="s">mid</text>
|
||||||
|
<text x="408" y="202" text-anchor="middle" class="s">coarse</text>
|
||||||
|
<text x="480" y="202" text-anchor="middle" class="s">low (8x8 pool)</text>
|
||||||
|
<text x="371" y="228" text-anchor="middle" class="m">band_k = low_k - low_{k+1}</text>
|
||||||
|
<text x="371" y="244" text-anchor="middle" class="s">low_k = residual pooled over 2^k x 2^k cells,</text>
|
||||||
|
<text x="371" y="258" text-anchor="middle" class="s">resampled to full size; bands sum to the residual</text>
|
||||||
|
<text x="371" y="272" text-anchor="middle" class="s">pred and target split separately (linear)</text>
|
||||||
|
<text x="371" y="286" text-anchor="middle" class="s">video: per frame over H, W</text>
|
||||||
|
<path d="M528,150 L564,150" class="ar"/>
|
||||||
|
|
||||||
|
<!-- step 3: weights -->
|
||||||
|
<rect x="568" y="90" width="368" height="200" class="box"/>
|
||||||
|
<text x="752" y="112" text-anchor="middle" class="h">3. weight each band by its recoverability at t</text>
|
||||||
|
<text x="588" y="138" class="m">flow: snr(t) = ((1-t)/t)^2, t = steps/1000</text>
|
||||||
|
<text x="588" y="154" class="m">ddpm: snr(t) = alpha_bar / (1 - alpha_bar)</text>
|
||||||
|
<text x="588" y="178" class="m">snr_k = snr(t) * P_clean[k] / P_noise[k]</text>
|
||||||
|
<text x="588" y="194" class="m">w_k = snr_k / (1 + snr_k)</text>
|
||||||
|
<text x="588" y="218" class="s">P_clean[k], P_noise[k]: mean squared band k of the</text>
|
||||||
|
<text x="588" y="232" class="s">batch's clean latents and of its noise, measured per</text>
|
||||||
|
<text x="588" y="246" class="s">batch: no 1/f prior, any VAE, pyramid noise just works.</text>
|
||||||
|
<text x="588" y="268" class="s">t -> 0: every w_k = 1, the loss IS the plain MSE.</text>
|
||||||
|
<text x="588" y="282" class="s">t -> 1: only the low band keeps a weight.</text>
|
||||||
|
|
||||||
|
<!-- step 4: combine -->
|
||||||
|
<path d="M371,292 L371,330" class="ar"/>
|
||||||
|
<rect x="216" y="334" width="310" height="150" class="box"/>
|
||||||
|
<text x="371" y="356" text-anchor="middle" class="h">4. filtered residual, squared</text>
|
||||||
|
<text x="371" y="382" text-anchor="middle" class="m">filtered = sum_k w_k * band_k</text>
|
||||||
|
<text x="371" y="400" text-anchor="middle" class="m">loss = filtered ** 2</text>
|
||||||
|
<text x="371" y="424" text-anchor="middle" class="s">elementwise, pred's shape: masks, timestep weights,</text>
|
||||||
|
<text x="371" y="438" text-anchor="middle" class="s">LossScaler and prior loss apply downstream unchanged</text>
|
||||||
|
<text x="371" y="462" text-anchor="middle" class="s">config: train.loss_type = "band"</text>
|
||||||
|
|
||||||
|
<!-- weight curves plot -->
|
||||||
|
<text x="560" y="318" class="h">w_k(t) for typical latent band powers</text>
|
||||||
|
<rect x="560" y="330" width="340" height="150" fill="#fcfcfc" stroke="#999"/>
|
||||||
|
<polyline points="561.7,330.0 563.4,330.2 565.1,330.4 566.8,330.8 568.5,331.2 570.2,331.8 571.9,332.4 573.6,333.2 575.3,334.1 577.0,335.0 578.7,336.1 580.4,337.3 582.1,338.5 583.8,339.9 585.5,341.4 587.2,343.0 588.9,344.6 590.6,346.3 592.3,348.2 594.0,350.1 595.7,352.0 597.4,354.0 599.1,356.1 600.8,358.3 602.5,360.5 604.2,362.7 605.9,365.0 607.6,367.3 609.3,369.7 611.0,372.0 612.7,374.4 614.4,376.8 616.1,379.2 617.8,381.6 619.5,384.0 621.2,386.4 622.9,388.8 624.6,391.1 626.3,393.5 628.0,395.8 629.7,398.1 631.4,400.4 633.1,402.6 634.8,404.8 636.5,407.0 638.2,409.1 639.9,411.2 641.6,413.2 643.3,415.2 645.0,417.2 646.7,419.1 648.4,421.0 650.1,422.9 651.8,424.6 653.5,426.4 655.2,428.1 656.9,429.8 658.6,431.4 660.3,433.0 662.0,434.5 663.7,436.0 665.4,437.4 667.1,438.8 668.8,440.2 670.5,441.5 672.2,442.8 673.9,444.0 675.6,445.3 677.3,446.4 679.0,447.6 680.7,448.7 682.4,449.7 684.1,450.8 685.8,451.8 687.5,452.7 689.2,453.7 690.9,454.6 692.6,455.4 694.3,456.3 696.0,457.1 697.7,457.9 699.4,458.7 701.1,459.4 702.8,460.1 704.5,460.8 706.2,461.5 707.9,462.2 709.6,462.8 711.3,463.4 713.0,464.0 714.7,464.6 716.4,465.1 718.1,465.6 719.8,466.1 721.5,466.6 723.2,467.1 724.9,467.6 726.6,468.0 728.3,468.5 730.0,468.9 731.7,469.3 733.4,469.7 735.1,470.1 736.8,470.4 738.5,470.8 740.2,471.1 741.9,471.5 743.6,471.8 745.3,472.1 747.0,472.4 748.7,472.7 750.4,472.9 752.1,473.2 753.8,473.5 755.5,473.7 757.2,474.0 758.9,474.2 760.6,474.4 762.3,474.6 764.0,474.8 765.7,475.1 767.4,475.3 769.1,475.4 770.8,475.6 772.5,475.8 774.2,476.0 775.9,476.1 777.6,476.3 779.3,476.5 781.0,476.6 782.7,476.7 784.4,476.9 786.1,477.0 787.8,477.1 789.5,477.3 791.2,477.4 792.9,477.5 794.6,477.6 796.3,477.7 798.0,477.8 799.7,477.9 801.4,478.0 803.1,478.1 804.8,478.2 806.5,478.3 808.2,478.4 809.9,478.5 811.6,478.5 813.3,478.6 815.0,478.7 816.7,478.7 818.4,478.8 820.1,478.9 821.8,478.9 823.5,479.0 825.2,479.1 826.9,479.1 828.6,479.2 830.3,479.2 832.0,479.3 833.7,479.3 835.4,479.3 837.1,479.4 838.8,479.4 840.5,479.5 842.2,479.5 843.9,479.5 845.6,479.6 847.3,479.6 849.0,479.6 850.7,479.7 852.4,479.7 854.1,479.7 855.8,479.7 857.5,479.8 859.2,479.8 860.9,479.8 862.6,479.8 864.3,479.8 866.0,479.9 867.7,479.9 869.4,479.9 871.1,479.9 872.8,479.9 874.5,479.9 876.2,479.9 877.9,479.9 879.6,480.0 881.3,480.0 883.0,480.0 884.7,480.0 886.4,480.0 888.1,480.0 889.8,480.0 891.5,480.0 893.2,480.0 894.9,480.0 896.6,480.0 898.3,480.0" fill="none" stroke="#d9534f" stroke-width="2.5"/>
|
||||||
|
<polyline points="561.7,330.0 563.4,330.0 565.1,330.1 566.8,330.1 568.5,330.2 570.2,330.2 571.9,330.3 573.6,330.4 575.3,330.6 577.0,330.7 578.7,330.8 580.4,331.0 582.1,331.2 583.8,331.4 585.5,331.6 587.2,331.9 588.9,332.1 590.6,332.4 592.3,332.7 594.0,333.0 595.7,333.4 597.4,333.7 599.1,334.1 600.8,334.5 602.5,334.9 604.2,335.4 605.9,335.9 607.6,336.3 609.3,336.9 611.0,337.4 612.7,338.0 614.4,338.6 616.1,339.2 617.8,339.8 619.5,340.5 621.2,341.2 622.9,341.9 624.6,342.6 626.3,343.4 628.0,344.2 629.7,345.0 631.4,345.8 633.1,346.7 634.8,347.6 636.5,348.5 638.2,349.4 639.9,350.4 641.6,351.4 643.3,352.4 645.0,353.4 646.7,354.5 648.4,355.6 650.1,356.7 651.8,357.8 653.5,359.0 655.2,360.2 656.9,361.4 658.6,362.6 660.3,363.9 662.0,365.2 663.7,366.4 665.4,367.8 667.1,369.1 668.8,370.4 670.5,371.8 672.2,373.2 673.9,374.6 675.6,376.0 677.3,377.4 679.0,378.9 680.7,380.3 682.4,381.8 684.1,383.3 685.8,384.8 687.5,386.2 689.2,387.8 690.9,389.3 692.6,390.8 694.3,392.3 696.0,393.8 697.7,395.4 699.4,396.9 701.1,398.4 702.8,400.0 704.5,401.5 706.2,403.0 707.9,404.5 709.6,406.1 711.3,407.6 713.0,409.1 714.7,410.6 716.4,412.1 718.1,413.6 719.8,415.1 721.5,416.6 723.2,418.0 724.9,419.5 726.6,420.9 728.3,422.3 730.0,423.8 731.7,425.1 733.4,426.5 735.1,427.9 736.8,429.3 738.5,430.6 740.2,431.9 741.9,433.2 743.6,434.5 745.3,435.8 747.0,437.0 748.7,438.2 750.4,439.5 752.1,440.6 753.8,441.8 755.5,443.0 757.2,444.1 758.9,445.2 760.6,446.3 762.3,447.4 764.0,448.4 765.7,449.4 767.4,450.5 769.1,451.4 770.8,452.4 772.5,453.4 774.2,454.3 775.9,455.2 777.6,456.1 779.3,456.9 781.0,457.8 782.7,458.6 784.4,459.4 786.1,460.2 787.8,460.9 789.5,461.7 791.2,462.4 792.9,463.1 794.6,463.8 796.3,464.5 798.0,465.1 799.7,465.7 801.4,466.4 803.1,466.9 804.8,467.5 806.5,468.1 808.2,468.6 809.9,469.1 811.6,469.7 813.3,470.1 815.0,470.6 816.7,471.1 818.4,471.5 820.1,472.0 821.8,472.4 823.5,472.8 825.2,473.2 826.9,473.5 828.6,473.9 830.3,474.2 832.0,474.6 833.7,474.9 835.4,475.2 837.1,475.5 838.8,475.8 840.5,476.1 842.2,476.3 843.9,476.6 845.6,476.8 847.3,477.0 849.0,477.2 850.7,477.5 852.4,477.7 854.1,477.8 855.8,478.0 857.5,478.2 859.2,478.3 860.9,478.5 862.6,478.6 864.3,478.8 866.0,478.9 867.7,479.0 869.4,479.1 871.1,479.2 872.8,479.3 874.5,479.4 876.2,479.5 877.9,479.6 879.6,479.6 881.3,479.7 883.0,479.8 884.7,479.8 886.4,479.8 888.1,479.9 889.8,479.9 891.5,479.9 893.2,480.0 894.9,480.0 896.6,480.0 898.3,480.0" fill="none" stroke="#f0a030" stroke-width="2.5"/>
|
||||||
|
<polyline points="561.7,330.0 563.4,330.0 565.1,330.0 566.8,330.0 568.5,330.0 570.2,330.0 571.9,330.0 573.6,330.0 575.3,330.1 577.0,330.1 578.7,330.1 580.4,330.1 582.1,330.1 583.8,330.1 585.5,330.2 587.2,330.2 588.9,330.2 590.6,330.2 592.3,330.3 594.0,330.3 595.7,330.3 597.4,330.4 599.1,330.4 600.8,330.5 602.5,330.5 604.2,330.6 605.9,330.6 607.6,330.7 609.3,330.7 611.0,330.8 612.7,330.8 614.4,330.9 616.1,331.0 617.8,331.0 619.5,331.1 621.2,331.2 622.9,331.3 624.6,331.4 626.3,331.5 628.0,331.5 629.7,331.6 631.4,331.7 633.1,331.9 634.8,332.0 636.5,332.1 638.2,332.2 639.9,332.3 641.6,332.5 643.3,332.6 645.0,332.7 646.7,332.9 648.4,333.0 650.1,333.2 651.8,333.3 653.5,333.5 655.2,333.7 656.9,333.9 658.6,334.1 660.3,334.3 662.0,334.5 663.7,334.7 665.4,334.9 667.1,335.1 668.8,335.3 670.5,335.6 672.2,335.8 673.9,336.1 675.6,336.4 677.3,336.6 679.0,336.9 680.7,337.2 682.4,337.5 684.1,337.8 685.8,338.2 687.5,338.5 689.2,338.8 690.9,339.2 692.6,339.6 694.3,339.9 696.0,340.3 697.7,340.8 699.4,341.2 701.1,341.6 702.8,342.1 704.5,342.5 706.2,343.0 707.9,343.5 709.6,344.0 711.3,344.5 713.0,345.1 714.7,345.6 716.4,346.2 718.1,346.8 719.8,347.4 721.5,348.0 723.2,348.7 724.9,349.3 726.6,350.0 728.3,350.7 730.0,351.4 731.7,352.2 733.4,352.9 735.1,353.7 736.8,354.5 738.5,355.4 740.2,356.2 741.9,357.1 743.6,358.0 745.3,358.9 747.0,359.9 748.7,360.9 750.4,361.9 752.1,362.9 753.8,364.0 755.5,365.1 757.2,366.2 758.9,367.3 760.6,368.5 762.3,369.7 764.0,370.9 765.7,372.2 767.4,373.4 769.1,374.8 770.8,376.1 772.5,377.5 774.2,378.9 775.9,380.3 777.6,381.8 779.3,383.2 781.0,384.8 782.7,386.3 784.4,387.9 786.1,389.5 787.8,391.1 789.5,392.7 791.2,394.4 792.9,396.1 794.6,397.8 796.3,399.6 798.0,401.4 799.7,403.2 801.4,405.0 803.1,406.8 804.8,408.6 806.5,410.5 808.2,412.4 809.9,414.3 811.6,416.2 813.3,418.1 815.0,420.0 816.7,421.9 818.4,423.8 820.1,425.8 821.8,427.7 823.5,429.6 825.2,431.5 826.9,433.4 828.6,435.3 830.3,437.2 832.0,439.1 833.7,440.9 835.4,442.8 837.1,444.6 838.8,446.4 840.5,448.1 842.2,449.8 843.9,451.5 845.6,453.2 847.3,454.8 849.0,456.4 850.7,457.9 852.4,459.4 854.1,460.9 855.8,462.3 857.5,463.6 859.2,464.9 860.9,466.2 862.6,467.4 864.3,468.6 866.0,469.7 867.7,470.7 869.4,471.7 871.1,472.6 872.8,473.5 874.5,474.3 876.2,475.1 877.9,475.8 879.6,476.4 881.3,477.0 883.0,477.5 884.7,478.0 886.4,478.5 888.1,478.8 889.8,479.1 891.5,479.4 893.2,479.6 894.9,479.8 896.6,479.9 898.3,480.0" fill="none" stroke="#3c9d5d" stroke-width="2.5"/>
|
||||||
|
<polyline points="561.7,330.0 563.4,330.0 565.1,330.0 566.8,330.0 568.5,330.0 570.2,330.0 571.9,330.0 573.6,330.0 575.3,330.0 577.0,330.0 578.7,330.0 580.4,330.0 582.1,330.0 583.8,330.0 585.5,330.0 587.2,330.0 588.9,330.0 590.6,330.0 592.3,330.0 594.0,330.0 595.7,330.0 597.4,330.0 599.1,330.0 600.8,330.0 602.5,330.1 604.2,330.1 605.9,330.1 607.6,330.1 609.3,330.1 611.0,330.1 612.7,330.1 614.4,330.1 616.1,330.1 617.8,330.1 619.5,330.1 621.2,330.1 622.9,330.1 624.6,330.1 626.3,330.1 628.0,330.2 629.7,330.2 631.4,330.2 633.1,330.2 634.8,330.2 636.5,330.2 638.2,330.2 639.9,330.2 641.6,330.2 643.3,330.3 645.0,330.3 646.7,330.3 648.4,330.3 650.1,330.3 651.8,330.3 653.5,330.4 655.2,330.4 656.9,330.4 658.6,330.4 660.3,330.4 662.0,330.5 663.7,330.5 665.4,330.5 667.1,330.5 668.8,330.6 670.5,330.6 672.2,330.6 673.9,330.6 675.6,330.7 677.3,330.7 679.0,330.7 680.7,330.8 682.4,330.8 684.1,330.8 685.8,330.9 687.5,330.9 689.2,330.9 690.9,331.0 692.6,331.0 694.3,331.1 696.0,331.1 697.7,331.1 699.4,331.2 701.1,331.2 702.8,331.3 704.5,331.4 706.2,331.4 707.9,331.5 709.6,331.5 711.3,331.6 713.0,331.7 714.7,331.7 716.4,331.8 718.1,331.9 719.8,331.9 721.5,332.0 723.2,332.1 724.9,332.2 726.6,332.3 728.3,332.4 730.0,332.5 731.7,332.6 733.4,332.7 735.1,332.8 736.8,332.9 738.5,333.0 740.2,333.1 741.9,333.2 743.6,333.4 745.3,333.5 747.0,333.6 748.7,333.8 750.4,333.9 752.1,334.1 753.8,334.3 755.5,334.4 757.2,334.6 758.9,334.8 760.6,335.0 762.3,335.2 764.0,335.4 765.7,335.6 767.4,335.9 769.1,336.1 770.8,336.4 772.5,336.6 774.2,336.9 775.9,337.2 777.6,337.5 779.3,337.8 781.0,338.2 782.7,338.5 784.4,338.9 786.1,339.2 787.8,339.6 789.5,340.1 791.2,340.5 792.9,341.0 794.6,341.4 796.3,341.9 798.0,342.5 799.7,343.0 801.4,343.6 803.1,344.2 804.8,344.9 806.5,345.6 808.2,346.3 809.9,347.0 811.6,347.8 813.3,348.7 815.0,349.6 816.7,350.5 818.4,351.5 820.1,352.5 821.8,353.6 823.5,354.8 825.2,356.0 826.9,357.3 828.6,358.6 830.3,360.1 832.0,361.6 833.7,363.2 835.4,364.9 837.1,366.7 838.8,368.5 840.5,370.5 842.2,372.6 843.9,374.9 845.6,377.2 847.3,379.7 849.0,382.3 850.7,385.0 852.4,387.9 854.1,390.9 855.8,394.1 857.5,397.4 859.2,400.9 860.9,404.5 862.6,408.3 864.3,412.2 866.0,416.2 867.7,420.3 869.4,424.5 871.1,428.8 872.8,433.2 874.5,437.6 876.2,441.9 877.9,446.3 879.6,450.5 881.3,454.7 883.0,458.6 884.7,462.4 886.4,465.8 888.1,469.0 889.8,471.9 891.5,474.3 893.2,476.3 894.9,477.9 896.6,479.1 898.3,479.8" fill="none" stroke="#3b7ddd" stroke-width="2.5"/>
|
||||||
|
<rect x="780" y="338" width="12" height="4" fill="#d9534f"/><text x="796" y="343" class="s">fine band</text>
|
||||||
|
<rect x="780" y="354" width="12" height="4" fill="#f0a030"/><text x="796" y="359" class="s">mid band</text>
|
||||||
|
<rect x="780" y="370" width="12" height="4" fill="#3c9d5d"/><text x="796" y="375" class="s">coarse band</text>
|
||||||
|
<rect x="780" y="386" width="12" height="4" fill="#3b7ddd"/><text x="796" y="391" class="s">low band</text>
|
||||||
|
<text x="554" y="334" text-anchor="end" class="s">1</text>
|
||||||
|
<text x="554" y="484" text-anchor="end" class="s">0</text>
|
||||||
|
<text x="560" y="496" class="s">t = 0 (clean)</text>
|
||||||
|
<text x="900" y="496" text-anchor="end" class="s">t = 1 (noise)</text>
|
||||||
|
<text x="730.0" y="512" text-anchor="middle" class="s">fine detail drops out first, global layout last</text>
|
||||||
|
|
||||||
|
<!-- why -->
|
||||||
|
<rect x="24" y="500" width="912" height="96" class="box"/>
|
||||||
|
<text x="40" y="522" class="h">why</text>
|
||||||
|
<text x="40" y="542" class="b">Images carry most power at low frequencies and white noise is flat, so at noise level t each band has its own</text>
|
||||||
|
<text x="40" y="558" class="b">signal-to-noise ratio and the fine ones go under first. Above that cutoff the best possible estimate is a blur, and an</text>
|
||||||
|
<text x="40" y="574" class="b">MSE gradient on those bands is a coin flip. Dropping them removes gradient noise at high t; at low t nothing changes.</text>
|
||||||
|
<text x="40" y="590" class="b">The depth-wise code heads give coarse-to-fine supervision across layers; this gives it across timesteps. They compose.</text>
|
||||||
|
</svg>
|
||||||
|
After Width: | Height: | Size: 17 KiB |
@@ -1,4 +1,5 @@
|
|||||||
import torch
|
import torch
|
||||||
|
import torch.nn.functional as F
|
||||||
|
|
||||||
|
|
||||||
_dwt = None
|
_dwt = None
|
||||||
@@ -91,3 +92,63 @@ def stepped_loss(model_pred, latents, noise, noisy_latents, timesteps, scheduler
|
|||||||
latents.float().to(device=predicted_images.device),
|
latents.float().to(device=predicted_images.device),
|
||||||
reduction="none",
|
reduction="none",
|
||||||
)
|
)
|
||||||
|
|
||||||
|
|
||||||
|
def _band_lows(x: torch.Tensor, levels: int) -> list:
|
||||||
|
"""Gaussian-pyramid low-passes of (N, C, H, W) at the input size: lows[0] = x, lows[k] = x
|
||||||
|
averaged over 2^k x 2^k cells and bilinearly resampled back."""
|
||||||
|
lows = [x]
|
||||||
|
cur = x
|
||||||
|
for _ in range(levels):
|
||||||
|
cur = F.avg_pool2d(cur, 2, ceil_mode=True)
|
||||||
|
lows.append(F.interpolate(cur, size=x.shape[-2:], mode="bilinear", align_corners=False))
|
||||||
|
return lows
|
||||||
|
|
||||||
|
|
||||||
|
def _bands(x: torch.Tensor, levels: int) -> list:
|
||||||
|
"""Laplacian bands at the input size, fine to coarse; they sum back to x."""
|
||||||
|
lows = _band_lows(x, levels)
|
||||||
|
return [lows[k] - lows[k + 1] for k in range(levels)] + [lows[levels]]
|
||||||
|
|
||||||
|
|
||||||
|
def band_loss(pred, target, latents, noise, timesteps, is_flow_matching, scheduler, levels: int = 3):
|
||||||
|
"""MSE on a per-sample frequency-filtered residual: each Laplacian band of the residual is
|
||||||
|
weighted by how recoverable that band is at the sample's noise level, snr_k / (1 + snr_k) with
|
||||||
|
snr_k = snr(t) * band power of the clean latents / band power of the noise (both measured on
|
||||||
|
the batch). At t -> 0 every weight is 1 and this is the plain MSE; at high t the bands the
|
||||||
|
noise has buried drop out instead of adding gradient noise. The filter is linear, so it
|
||||||
|
applies to pred / target as they are (x0, velocity or epsilon residuals differ only by a
|
||||||
|
scalar in t). Returns an elementwise loss, pred's shape."""
|
||||||
|
pred, target = pred.float(), target.float()
|
||||||
|
shape = pred.shape
|
||||||
|
if pred.dim() not in (4, 5):
|
||||||
|
return F.mse_loss(pred, target, reduction="none")
|
||||||
|
n = shape[0]
|
||||||
|
if pred.dim() == 5: # video: bands over (H, W) of every frame
|
||||||
|
flat = lambda x: x.float().transpose(1, 2).flatten(0, 1)
|
||||||
|
else:
|
||||||
|
flat = lambda x: x.float()
|
||||||
|
with torch.no_grad():
|
||||||
|
clean_b = _bands(flat(latents.to(pred.device)), levels)
|
||||||
|
noise_b = _bands(flat(noise.to(pred.device)), levels)
|
||||||
|
# band power over the whole batch: a prior-level quantity, not per image
|
||||||
|
p_clean = torch.stack([b.pow(2).mean() for b in clean_b])
|
||||||
|
p_noise = torch.stack([b.pow(2).mean() for b in noise_b]).clamp_min(1e-12)
|
||||||
|
ts = timesteps.to(pred.device).float().flatten()
|
||||||
|
if is_flow_matching:
|
||||||
|
t = (ts / 1000.0).clamp(1e-4, 1 - 1e-4)
|
||||||
|
snr = ((1 - t) / t) ** 2
|
||||||
|
else:
|
||||||
|
ac = scheduler.alphas_cumprod.to(pred.device, torch.float32)[ts.long().clamp(0, scheduler.alphas_cumprod.shape[0] - 1)]
|
||||||
|
snr = ac / (1 - ac).clamp_min(1e-8)
|
||||||
|
s = snr.view(n, 1) * (p_clean / p_noise).view(1, -1) # (N, levels + 1)
|
||||||
|
w = s / (1 + s)
|
||||||
|
if pred.dim() == 5:
|
||||||
|
w = w.repeat_interleave(shape[2], dim=0) # per frame row
|
||||||
|
w = w.t().reshape(levels + 1, -1, 1, 1, 1)
|
||||||
|
pb, tb = _bands(flat(pred), levels), _bands(flat(target), levels)
|
||||||
|
filtered = sum(w[k] * (pb[k] - tb[k]) for k in range(levels + 1))
|
||||||
|
loss = filtered.pow(2)
|
||||||
|
if pred.dim() == 5:
|
||||||
|
loss = loss.view(n, shape[2], shape[1], *shape[-2:]).transpose(1, 2)
|
||||||
|
return loss
|
||||||
Reference in new issue
Block a user