Changes
diff --git a/diffusion_cli/sampling.py b/diffusion_cli/sampling.py
index d1c681c..6fcc2c1 100644
--- a/diffusion_cli/sampling.py
+++ b/diffusion_cli/sampling.py
@@ -1,13 +1,16 @@
-"""Scheduler and latent-noise helpers for the first Z-Image milestone."""
+"""Scheduler, sampler, and latent-noise helpers for Z-Image inference."""
from __future__ import annotations
from dataclasses import dataclass
+from typing import Callable
+import numpy as np
import torch
from scipy.stats import beta
from diffusion_cli.config import LATENT_CHANNELS, LATENT_DOWNSCALE
+from diffusion_cli.errors import DiffusionCliError
@dataclass(frozen=True)
@@ -16,6 +19,7 @@ class ModelSamplingDiscreteFlow:
shift: float = 3.0
multiplier: float = 1.0
+ timesteps: int = 1000
def sigma(self, timestep: torch.Tensor) -> torch.Tensor:
"""Convert normalized timesteps to shifted flow sigmas."""
@@ -30,6 +34,18 @@ class ModelSamplingDiscreteFlow:
return sigma * self.multiplier
+ def sigmaTable(self, device=None, dtype=None) -> torch.Tensor:
+ """Return the 1,000-step shifted flow table used by ComfyUI."""
+
+ indices = torch.arange(
+ 1,
+ self.timesteps + 1,
+ device=device,
+ dtype=dtype or torch.float32,
+ )
+ timesteps = indices / self.timesteps * self.multiplier
+ return self.sigma(timesteps)
+
def latentShape(
batch_size: int,
@@ -72,29 +88,249 @@ def betaScheduler(
device: object | None = None,
dtype: object | None = None,
) -> torch.Tensor:
- """Return beta scheduler sigmas with a final zero sigma."""
+ """Return ComfyUI beta scheduler sigmas with a final zero sigma."""
if steps < 1:
raise ValueError("steps must be at least 1")
sampling = sampling or ModelSamplingDiscreteFlow()
- quantiles = torch.tensor(
- beta.ppf(torch.linspace(0, 1, steps).tolist(), 0.6, 0.6),
- device=device,
- dtype=dtype or torch.float32,
- )
- quantiles = torch.nan_to_num(quantiles, nan=0.0, posinf=1.0, neginf=0.0)
- quantiles = torch.flip(quantiles, dims=(0,))
- sigmas = sampling.sigma(quantiles)
+ table = sampling.sigmaTable(device=device, dtype=dtype)
+ total_timesteps = len(table) - 1
+ percentiles = 1.0 - np.linspace(0.0, 1.0, steps, endpoint=False)
+ timesteps = np.rint(beta.ppf(percentiles, 0.6, 0.6) * total_timesteps)
+
+ values = []
+ last_t = -1
+ for timestep in timesteps:
+ if timestep != last_t:
+ values.append(table[int(timestep)])
+ last_t = timestep
+
+ sigmas = torch.stack(values) if values else table[:0]
sigmas = torch.cat(
[sigmas, torch.zeros(1, device=sigmas.device, dtype=sigmas.dtype)]
)
return sigmas
-def sampleSeeds3(*_args, **_kwargs):
- """Run the SEEDS-3 sampler once the diffusion model port exists."""
+def sigmaToHalfLogSnr(sigma: torch.Tensor) -> torch.Tensor:
+ """Convert sigma to half-log-SNR for flow samplers."""
+
+ return sigma.log().neg()
+
+
+def halfLogSnrToSigma(half_log_snr: torch.Tensor) -> torch.Tensor:
+ """Convert half-log-SNR to sigma for flow samplers."""
+
+ return half_log_snr.neg().exp()
+
+
+def eiHPhi1(h: torch.Tensor) -> torch.Tensor:
+ """Compute h * phi_1(h) for exponential integrator updates."""
+
+ return torch.expm1(h)
+
+
+def eiHPhi2(h: torch.Tensor) -> torch.Tensor:
+ """Compute h * phi_2(h) for exponential integrator updates."""
+
+ return (torch.expm1(h) - h) / h
+
+
+def defaultNoiseSampler(
+ latent: torch.Tensor,
+ seed: int | None = None,
+) -> Callable[[torch.Tensor, torch.Tensor], torch.Tensor]:
+ """Return a deterministic Gaussian noise sampler for stochastic steps."""
+
+ if seed is None:
+ generator = None
+ else:
+ seed_value = seed + 1 if latent.device == torch.device("cpu") else seed
+ generator = torch.Generator(device=latent.device)
+ generator.manual_seed(seed_value)
+
+ def sample(_sigma, _sigma_next):
+ return torch.randn(
+ latent.size(),
+ dtype=latent.dtype,
+ layout=latent.layout,
+ device=latent.device,
+ generator=generator,
+ )
+
+ return sample
+
+
+def cfgDenoise(model, latent, sigma, conditioning, cfg: float) -> torch.Tensor:
+ """Run classifier-free guidance for one denoising call."""
+
+ positive = model.forward(latent, sigma, conditioning.positive)
+ negative = model.forward(latent, sigma, conditioning.negative)
+ return negative + (positive - negative) * cfg
+
+
+def _checkFinite(latent: torch.Tensor, context: str) -> None:
+ if not torch.all(torch.isfinite(latent)):
+ raise DiffusionCliError(
+ f"SEEDS-3 sampler produced NaN or Inf {context}"
+ )
- raise NotImplementedError(
- "SEEDS-3 sampling needs the Z-Image diffusion model port"
+
+@torch.no_grad()
+def sampleSeeds3(
+ model,
+ latent: torch.Tensor,
+ sigmas: torch.Tensor,
+ conditioning,
+ cfg: float,
+ *,
+ eta: float = 1.0,
+ s_noise: float = 1.0,
+ seed: int | None = None,
+ noise_sampler=None,
+ r_1: float = 1.0 / 3.0,
+ r_2: float = 2.0 / 3.0,
+ callback=None,
+) -> torch.Tensor:
+ """Run the SEEDS-3 solver over a beta sigma schedule."""
+
+ if len(sigmas) < 2:
+ raise ValueError("sigmas must contain at least two values")
+
+ x = latent
+ sigmas = sigmas.to(device=x.device, dtype=x.dtype)
+ noise_sampler = noise_sampler or defaultNoiseSampler(x, seed=seed)
+ s_in = x.new_ones([x.shape[0]])
+ inject_noise = eta > 0 and s_noise > 0
+
+ for i in range(len(sigmas) - 1):
+ sigma = sigmas[i]
+ sigma_next = sigmas[i + 1]
+ denoised = cfgDenoise(model, x, sigma * s_in, conditioning, cfg)
+ _checkFinite(denoised, f"after denoising step {i}")
+ if callback is not None:
+ callback({
+ "x": x,
+ "i": i,
+ "sigma": sigma,
+ "sigma_hat": sigma,
+ "denoised": denoised,
+ })
+
+ if sigma_next == 0:
+ x = denoised
+ _checkFinite(x, f"at final step {i}")
+ continue
+
+ lambda_s = sigmaToHalfLogSnr(sigma)
+ lambda_t = sigmaToHalfLogSnr(sigma_next)
+ h = lambda_t - lambda_s
+ h_eta = h * (eta + 1)
+ lambda_s_1 = torch.lerp(lambda_s, lambda_t, r_1)
+ lambda_s_2 = torch.lerp(lambda_s, lambda_t, r_2)
+ sigma_s_1 = halfLogSnrToSigma(lambda_s_1)
+ sigma_s_2 = halfLogSnrToSigma(lambda_s_2)
+
+ alpha_s_1 = sigma_s_1 * lambda_s_1.exp()
+ alpha_s_2 = sigma_s_2 * lambda_s_2.exp()
+ alpha_t = sigma_next * lambda_t.exp()
+
+ x_2 = (
+ sigma_s_1 / sigma
+ * (-r_1 * h * eta).exp()
+ * x
+ - alpha_s_1 * eiHPhi1(-r_1 * h_eta) * denoised
+ )
+ if inject_noise:
+ sde_noise = (
+ (-2 * r_1 * h * eta).expm1().neg().sqrt()
+ * noise_sampler(sigma, sigma_s_1)
+ )
+ x_2 = x_2 + sde_noise * sigma_s_1 * s_noise
+ denoised_2 = cfgDenoise(
+ model,
+ x_2,
+ sigma_s_1 * s_in,
+ conditioning,
+ cfg,
+ )
+
+ a3_2 = r_2 / r_1 * eiHPhi2(-r_2 * h_eta)
+ a3_1 = eiHPhi1(-r_2 * h_eta) - a3_2
+ x_3 = (
+ sigma_s_2 / sigma
+ * (-r_2 * h * eta).exp()
+ * x
+ - alpha_s_2 * (a3_1 * denoised + a3_2 * denoised_2)
+ )
+ if inject_noise:
+ segment_factor = (r_1 - r_2) * h * eta
+ sde_noise = sde_noise * segment_factor.exp()
+ sde_noise = sde_noise + (
+ segment_factor.mul(2).expm1().neg().sqrt()
+ * noise_sampler(sigma_s_1, sigma_s_2)
+ )
+ x_3 = x_3 + sde_noise * sigma_s_2 * s_noise
+ denoised_3 = cfgDenoise(
+ model,
+ x_3,
+ sigma_s_2 * s_in,
+ conditioning,
+ cfg,
+ )
+
+ b3 = eiHPhi2(-h_eta) / r_2
+ b1 = eiHPhi1(-h_eta) - b3
+ x = (
+ sigma_next / sigma
+ * (-h * eta).exp()
+ * x
+ - alpha_t * (b1 * denoised + b3 * denoised_3)
+ )
+ if inject_noise:
+ segment_factor = (r_2 - 1) * h * eta
+ sde_noise = sde_noise * segment_factor.exp()
+ sde_noise = sde_noise + (
+ segment_factor.mul(2).expm1().neg().sqrt()
+ * noise_sampler(sigma_s_2, sigma_next)
+ )
+ x = x + sde_noise * sigma_next * s_noise
+
+ _checkFinite(x, f"after step {i}")
+
+ return x
+
+
+def sampleLatents(
+ model,
+ conditioning,
+ *,
+ batch_size: int,
+ height: int,
+ width: int,
+ seed: int,
+ steps: int,
+ cfg: float,
+ device,
+ dtype,
+) -> torch.Tensor:
+ """Build initial noise and run the default Z-Image sampler."""
+
+ latent = buildInitialNoise(
+ batch_size,
+ height,
+ width,
+ seed,
+ device,
+ dtype,
+ )
+ sigmas = betaScheduler(steps, device=device, dtype=dtype)
+ return sampleSeeds3(
+ model,
+ latent,
+ sigmas,
+ conditioning,
+ cfg,
+ seed=seed,
)
diff --git a/tests/test_sampling.py b/tests/test_sampling.py
index ed5968a..fcc77fc 100644
--- a/tests/test_sampling.py
+++ b/tests/test_sampling.py
@@ -1,8 +1,34 @@
import unittest
+from types import SimpleNamespace
import torch
-from diffusion_cli.sampling import betaScheduler, latentShape
+from diffusion_cli.sampling import (
+ betaScheduler,
+ cfgDenoise,
+ latentShape,
+ sampleSeeds3,
+)
+
+
+class FakeConditioning:
+ pass
+
+
+class FakeModel:
+ def __init__(self):
+ self.calls = []
+
+ def forward(self, latent, sigma, conditioning):
+ self.calls.append((sigma.detach().clone(), conditioning))
+ if conditioning is POSITIVE:
+ return latent * 0.25
+ return latent * -0.25
+
+
+POSITIVE = FakeConditioning()
+NEGATIVE = FakeConditioning()
+CONDITIONING = SimpleNamespace(positive=POSITIVE, negative=NEGATIVE)
class SamplingTest(unittest.TestCase):
@@ -15,6 +41,39 @@ class SamplingTest(unittest.TestCase):
self.assertEqual(sigmas.shape, (11,))
self.assertEqual(sigmas[-1].item(), 0.0)
self.assertTrue(torch.all(torch.isfinite(sigmas)))
+ self.assertGreater(sigmas[0].item(), 0.99)
+ self.assertLess(sigmas[-2].item(), 0.2)
+
+ def testCfgDenoiseUsesNegativePathAtCfgOne(self):
+ model = FakeModel()
+ latent = torch.ones(1, 1, 1, 1)
+ sigma = torch.ones(1)
+
+ output = cfgDenoise(model, latent, sigma, CONDITIONING, 1.0)
+
+ self.assertTrue(torch.equal(output, latent * 0.25))
+ self.assertEqual(len(model.calls), 2)
+ self.assertIs(model.calls[0][1], POSITIVE)
+ self.assertIs(model.calls[1][1], NEGATIVE)
+
+ def testSampleSeeds3RunsAndRemainsFinite(self):
+ model = FakeModel()
+ latent = torch.zeros(1, 1, 2, 2)
+ sigmas = torch.tensor([1.0, 0.5, 0.0])
+
+ output = sampleSeeds3(
+ model,
+ latent,
+ sigmas,
+ CONDITIONING,
+ 1.0,
+ eta=0.0,
+ s_noise=0.0,
+ )
+
+ self.assertEqual(output.shape, latent.shape)
+ self.assertTrue(torch.all(torch.isfinite(output)))
+ self.assertEqual(len(model.calls), 8)
if __name__ == "__main__":