Changes
diff --git a/diffusion_cli/cli.py b/diffusion_cli/cli.py
index 69fba26..4709c55 100644
--- a/diffusion_cli/cli.py
+++ b/diffusion_cli/cli.py
@@ -3,6 +3,7 @@
from __future__ import annotations
import argparse
+import gc
from pathlib import Path
import sys
@@ -19,8 +20,13 @@ from diffusion_cli.config import (
buildGenerationConfig,
)
from diffusion_cli.errors import DiffusionCliError
+from diffusion_cli.image_io import saveImages
from diffusion_cli.model_inspect import formatSummary, inspectSafetensors
from diffusion_cli.paths import resolveModelFiles
+from diffusion_cli.sampling import sampleLatents
+from diffusion_cli.text_encoder import ZImageTextEncoder
+from diffusion_cli.vae import ZImageVae
+from diffusion_cli.zimage_model import ZImageModel
def buildParser() -> argparse.ArgumentParser:
@@ -80,16 +86,61 @@ def inspectModels(model_root: Path | None) -> None:
print(output)
+def releaseMemory() -> None:
+ """Release Python and CUDA caches between large model stages."""
+
+ gc.collect()
+ try:
+ import torch
+
+ if torch.cuda.is_available():
+ torch.cuda.empty_cache()
+ except ImportError:
+ pass
+
+
def generate(args) -> list[Path]:
"""Validate generation arguments and run the current milestone."""
- buildGenerationConfig(args)
- resolveModelFiles(args.model_root)
- raise DiffusionCliError(
- "Full image generation is not implemented yet. Run "
- "`diffusion-cli --inspect-models` to verify local checkpoint "
- "compatibility before the Z-Image model port is added."
+ config = buildGenerationConfig(args)
+ model_files = resolveModelFiles(args.model_root)
+
+ text_encoder = ZImageTextEncoder(
+ model_files.text_encoder,
+ config.tokenizer_path,
+ config.device,
+ config.dtype,
)
+ conditioning = text_encoder.encodePrompts(
+ config.prompt,
+ config.negative_prompt,
+ )
+ del text_encoder
+ releaseMemory()
+
+ model = ZImageModel(
+ model_files.diffusion_model,
+ config.device,
+ config.dtype,
+ )
+ latent = sampleLatents(
+ model,
+ conditioning,
+ batch_size=config.batch_size,
+ height=config.height,
+ width=config.width,
+ seed=config.seed,
+ steps=config.steps,
+ cfg=config.cfg,
+ device=config.device,
+ dtype=config.dtype,
+ )
+ del model
+ releaseMemory()
+
+ vae = ZImageVae(model_files.vae, config.device, config.dtype)
+ images = vae.decode(latent)
+ return saveImages(images, config.output)
def main(argv: list[str] | None = None) -> int:
diff --git a/diffusion_cli/sampling.py b/diffusion_cli/sampling.py
index 6fcc2c1..671f799 100644
--- a/diffusion_cli/sampling.py
+++ b/diffusion_cli/sampling.py
@@ -46,6 +46,16 @@ class ModelSamplingDiscreteFlow:
timesteps = indices / self.timesteps * self.multiplier
return self.sigma(timesteps)
+ def percentToSigma(self, percent: float) -> torch.Tensor:
+ """Convert a denoise percentage to shifted flow sigma."""
+
+ if percent <= 0.0:
+ return torch.tensor(1.0)
+ if percent >= 1.0:
+ return torch.tensor(0.0)
+ timestep = torch.tensor(1.0 - percent) * self.multiplier
+ return self.sigma(timestep)
+
def latentShape(
batch_size: int,
@@ -116,13 +126,31 @@ def betaScheduler(
def sigmaToHalfLogSnr(sigma: torch.Tensor) -> torch.Tensor:
"""Convert sigma to half-log-SNR for flow samplers."""
- return sigma.log().neg()
+ return torch.logit(sigma).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()
+ return half_log_snr.neg().sigmoid()
+
+
+def offsetFirstSigmaForSnr(
+ sigmas: torch.Tensor,
+ sampling: ModelSamplingDiscreteFlow | None = None,
+ percent_offset: float = 1e-4,
+) -> torch.Tensor:
+ """Offset an initial sigma of 1 to avoid infinite flow log-SNR."""
+
+ if len(sigmas) <= 1 or sigmas[0] < 1:
+ return sigmas
+ sampling = sampling or ModelSamplingDiscreteFlow()
+ sigmas = sigmas.clone()
+ sigmas[0] = sampling.percentToSigma(percent_offset).to(
+ device=sigmas.device,
+ dtype=sigmas.dtype,
+ )
+ return sigmas
def eiHPhi1(h: torch.Tensor) -> torch.Tensor:
@@ -198,10 +226,11 @@ def sampleSeeds3(
if len(sigmas) < 2:
raise ValueError("sigmas must contain at least two values")
- x = latent
- sigmas = sigmas.to(device=x.device, dtype=x.dtype)
+ x = latent.float()
+ sigmas = sigmas.to(device=x.device, dtype=torch.float32)
+ sigmas = offsetFirstSigmaForSnr(sigmas)
noise_sampler = noise_sampler or defaultNoiseSampler(x, seed=seed)
- s_in = x.new_ones([x.shape[0]])
+ s_in = torch.ones(x.shape[0], device=x.device, dtype=torch.float32)
inject_noise = eta > 0 and s_noise > 0
for i in range(len(sigmas) - 1):
@@ -323,14 +352,18 @@ def sampleLatents(
width,
seed,
device,
- dtype,
+ torch.float32,
)
- sigmas = betaScheduler(steps, device=device, dtype=dtype)
+ sigmas = betaScheduler(steps, device=device, dtype=torch.float32)
+ # Keep the first standalone milestone deterministic; the stochastic
+ # SEEDS-3 path is more sensitive to numeric parity with ComfyUI internals.
return sampleSeeds3(
model,
latent,
sigmas,
conditioning,
cfg,
+ eta=0.0,
+ s_noise=0.0,
seed=seed,
)
diff --git a/diffusion_cli/text_encoder.py b/diffusion_cli/text_encoder.py
index 829988b..27f3fa1 100644
--- a/diffusion_cli/text_encoder.py
+++ b/diffusion_cli/text_encoder.py
@@ -160,6 +160,7 @@ class ZImageTextEncoder:
with torch.device("meta"):
model = Qwen3Model(buildQwen3_4BConfig())
+ model.to_empty(device="cpu")
state_dict = normalizeQwenStateDict(load_file(model_path, device="cpu"))
missing_keys, unexpected_keys = model.load_state_dict(
diff --git a/diffusion_cli/vae.py b/diffusion_cli/vae.py
index 43fd38e..1e4e262 100644
--- a/diffusion_cli/vae.py
+++ b/diffusion_cli/vae.py
@@ -1,10 +1,387 @@
-"""VAE loading boundary for Z-Image Turbo decoding."""
+"""VAE loading and decode support for Z-Image Turbo."""
from __future__ import annotations
+from dataclasses import dataclass
+from pathlib import Path
+
+import torch
+from safetensors.torch import load_file
+from torch import nn
+from torch.nn import functional as F
+
+from diffusion_cli.config import LATENT_CHANNELS, LATENT_DOWNSCALE
+from diffusion_cli.errors import DiffusionCliError
+
+
+FLUX_SCALE_FACTOR = 0.3611
+FLUX_SHIFT_FACTOR = 0.1159
+VAE_CHANNELS = 128
+VAE_CHANNEL_MULTIPLIERS = (1, 2, 4, 4)
+VAE_NUM_RES_BLOCKS = 2
+VAE_OUT_CHANNELS = 3
+VAE_RESOLUTION = 256
+
+
+@dataclass(frozen=True)
+class VaeConfig:
+ """Architecture settings for the supported image VAE decoder."""
+
+ channels: int = VAE_CHANNELS
+ channel_multipliers: tuple[int, ...] = VAE_CHANNEL_MULTIPLIERS
+ num_res_blocks: int = VAE_NUM_RES_BLOCKS
+ latent_channels: int = LATENT_CHANNELS
+ out_channels: int = VAE_OUT_CHANNELS
+ resolution: int = VAE_RESOLUTION
+
+
+def fluxLatentToVae(latent: torch.Tensor) -> torch.Tensor:
+ """Convert Flux-format model latents into VAE latents."""
+
+ return latent / FLUX_SCALE_FACTOR + FLUX_SHIFT_FACTOR
+
+
+def postprocessVaeOutput(image: torch.Tensor) -> torch.Tensor:
+ """Convert decoded VAE output from [-1, 1] to clamped [0, 1]."""
+
+ return image.add(1.0).div(2.0).clamp(0.0, 1.0)
+
+
+def detectVaeConfig(state_dict: dict[str, torch.Tensor]) -> VaeConfig:
+ """Detect the exact Z-Image VAE decoder architecture."""
+
+ try:
+ conv_in = state_dict["decoder.conv_in.weight"]
+ conv_out = state_dict["decoder.conv_out.weight"]
+ except KeyError as exc:
+ raise DiffusionCliError(
+ "Unsupported VAE state dict: no known decoder keys found"
+ ) from exc
+
+ if conv_in.ndim != 4 or conv_in.shape[1] != LATENT_CHANNELS:
+ raise DiffusionCliError(
+ "Unsupported VAE latent channels: expected "
+ f"{LATENT_CHANNELS}, got {tuple(conv_in.shape)}"
+ )
+ if conv_out.ndim != 4 or conv_out.shape[0] != VAE_OUT_CHANNELS:
+ raise DiffusionCliError(
+ "Unsupported VAE output channels: expected "
+ f"{VAE_OUT_CHANNELS}, got {tuple(conv_out.shape)}"
+ )
+
+ base_channels = conv_out.shape[1]
+ block_in = conv_in.shape[0]
+ if base_channels != VAE_CHANNELS or block_in != VAE_CHANNELS * 4:
+ raise DiffusionCliError(
+ "Unsupported VAE channel layout: expected Z-Image image VAE"
+ )
+
+ return VaeConfig()
+
+
+def _normalize(channels: int) -> nn.GroupNorm:
+ return nn.GroupNorm(
+ num_groups=32,
+ num_channels=channels,
+ eps=1e-6,
+ affine=True,
+ )
+
+
+class Upsample(nn.Module):
+ """Nearest-neighbor 2x upsample followed by a 3x3 convolution."""
+
+ def __init__(self, channels: int, *, dtype=None, device=None) -> None:
+ super().__init__()
+ self.conv = nn.Conv2d(
+ channels,
+ channels,
+ kernel_size=3,
+ padding=1,
+ dtype=dtype,
+ device=device,
+ )
+
+ def forward(self, x: torch.Tensor) -> torch.Tensor:
+ """Apply 2x nearest-neighbor upsampling."""
+
+ x = F.interpolate(x, scale_factor=2.0, mode="nearest")
+ return self.conv(x)
+
+
+class ResnetBlock(nn.Module):
+ """Residual block used by the Z-Image VAE decoder."""
+
+ def __init__(
+ self,
+ in_channels: int,
+ out_channels: int | None = None,
+ *,
+ dtype=None,
+ device=None,
+ ) -> None:
+ super().__init__()
+ out_channels = out_channels if out_channels is not None else in_channels
+ self.in_channels = in_channels
+ self.out_channels = out_channels
+ self.norm1 = _normalize(in_channels).to(device=device, dtype=dtype)
+ self.conv1 = nn.Conv2d(
+ in_channels,
+ out_channels,
+ kernel_size=3,
+ padding=1,
+ dtype=dtype,
+ device=device,
+ )
+ self.norm2 = _normalize(out_channels).to(device=device, dtype=dtype)
+ self.dropout = nn.Dropout(0.0)
+ self.conv2 = nn.Conv2d(
+ out_channels,
+ out_channels,
+ kernel_size=3,
+ padding=1,
+ dtype=dtype,
+ device=device,
+ )
+ if in_channels != out_channels:
+ self.nin_shortcut = nn.Conv2d(
+ in_channels,
+ out_channels,
+ kernel_size=1,
+ dtype=dtype,
+ device=device,
+ )
+ else:
+ self.nin_shortcut = None
+
+ def forward(self, x: torch.Tensor) -> torch.Tensor:
+ """Apply the residual update."""
+
+ h = self.norm1(x)
+ h = F.silu(h)
+ h = self.conv1(h)
+ h = self.norm2(h)
+ h = F.silu(h)
+ h = self.dropout(h)
+ h = self.conv2(h)
+ if self.nin_shortcut is not None:
+ x = self.nin_shortcut(x)
+ return x + h
+
+
+class AttnBlock(nn.Module):
+ """Single spatial self-attention block used at the VAE bottleneck."""
+
+ def __init__(self, channels: int, *, dtype=None, device=None) -> None:
+ super().__init__()
+ self.norm = _normalize(channels).to(device=device, dtype=dtype)
+ self.q = nn.Conv2d(
+ channels,
+ channels,
+ kernel_size=1,
+ dtype=dtype,
+ device=device,
+ )
+ self.k = nn.Conv2d(
+ channels,
+ channels,
+ kernel_size=1,
+ dtype=dtype,
+ device=device,
+ )
+ self.v = nn.Conv2d(
+ channels,
+ channels,
+ kernel_size=1,
+ dtype=dtype,
+ device=device,
+ )
+ self.proj_out = nn.Conv2d(
+ channels,
+ channels,
+ kernel_size=1,
+ dtype=dtype,
+ device=device,
+ )
+
+ def forward(self, x: torch.Tensor) -> torch.Tensor:
+ """Apply full spatial attention over latent positions."""
+
+ batch_size, channels, height, width = x.shape
+ h = self.norm(x)
+ query = self.q(h).reshape(batch_size, channels, -1).transpose(1, 2)
+ key = self.k(h).reshape(batch_size, channels, -1).transpose(1, 2)
+ value = self.v(h).reshape(batch_size, channels, -1).transpose(1, 2)
+ h = F.scaled_dot_product_attention(
+ query.unsqueeze(1),
+ key.unsqueeze(1),
+ value.unsqueeze(1),
+ )
+ h = h.squeeze(1).transpose(1, 2).reshape(
+ batch_size,
+ channels,
+ height,
+ width,
+ )
+ return x + self.proj_out(h)
+
+
+class VaeDecoder(nn.Module):
+ """Image-only decoder matching ComfyUI's Z-Image VAE branch."""
+
+ def __init__(
+ self,
+ config: VaeConfig,
+ *,
+ dtype=None,
+ device=None,
+ ) -> None:
+ super().__init__()
+ self.config = config
+ self.num_resolutions = len(config.channel_multipliers)
+ self.num_res_blocks = config.num_res_blocks
+ block_in = (
+ config.channels
+ * config.channel_multipliers[self.num_resolutions - 1]
+ )
+ self.conv_in = nn.Conv2d(
+ config.latent_channels,
+ block_in,
+ kernel_size=3,
+ padding=1,
+ dtype=dtype,
+ device=device,
+ )
+ self.mid = nn.Module()
+ self.mid.block_1 = ResnetBlock(
+ block_in,
+ block_in,
+ dtype=dtype,
+ device=device,
+ )
+ self.mid.attn_1 = AttnBlock(block_in, dtype=dtype, device=device)
+ self.mid.block_2 = ResnetBlock(
+ block_in,
+ block_in,
+ dtype=dtype,
+ device=device,
+ )
+
+ self.up = nn.ModuleList()
+ for level in reversed(range(self.num_resolutions)):
+ block = nn.ModuleList()
+ block_out = config.channels * config.channel_multipliers[level]
+ for _ in range(self.num_res_blocks + 1):
+ block.append(
+ ResnetBlock(
+ block_in,
+ block_out,
+ dtype=dtype,
+ device=device,
+ )
+ )
+ block_in = block_out
+ up = nn.Module()
+ up.block = block
+ up.attn = nn.ModuleList()
+ if level != 0:
+ up.upsample = Upsample(block_in, dtype=dtype, device=device)
+ self.up.insert(0, up)
+
+ self.norm_out = _normalize(block_in).to(device=device, dtype=dtype)
+ self.conv_out = nn.Conv2d(
+ block_in,
+ config.out_channels,
+ kernel_size=3,
+ padding=1,
+ dtype=dtype,
+ device=device,
+ )
+
+ def forward(self, latent: torch.Tensor) -> torch.Tensor:
+ """Decode a VAE latent tensor into an image tensor in [-1, 1]."""
+
+ h = self.conv_in(latent)
+ h = self.mid.block_1(h)
+ h = self.mid.attn_1(h)
+ h = self.mid.block_2(h)
+
+ for level in reversed(range(self.num_resolutions)):
+ for block in self.up[level].block:
+ h = block(h)
+ if level != 0:
+ h = self.up[level].upsample(h)
+
+ h = self.norm_out(h)
+ h = F.silu(h)
+ return self.conv_out(h)
+
class ZImageVae:
- """Placeholder boundary for the future VAE decoder implementation."""
+ """Load the local Z-Image VAE and decode final latents."""
+
+ def __init__(
+ self,
+ model_path: Path | None,
+ device,
+ dtype,
+ *,
+ model: nn.Module | None = None,
+ decoder_config: VaeConfig | None = None,
+ ) -> None:
+ self.device = device
+ self.dtype = dtype
+ if model is not None:
+ self.model = model.to(device=device, dtype=dtype).eval()
+ return
+ if model_path is None:
+ raise ValueError("model_path is required when model is absent")
+
+ state_dict = load_file(model_path, device="cpu")
+ config = decoder_config or detectVaeConfig(state_dict)
+ decoder = VaeDecoder(config, dtype=dtype, device=device)
+ decoder_state = {
+ key.removeprefix("decoder."): value
+ for key, value in state_dict.items()
+ if key.startswith("decoder.")
+ }
+ missing, unexpected = decoder.load_state_dict(
+ decoder_state,
+ strict=False,
+ )
+ if missing or unexpected:
+ raise DiffusionCliError(
+ "Unsupported VAE state dict: decoder keys did not match"
+ )
+ self.model = decoder.eval()
+
+ @torch.inference_mode()
+ def decode(self, latent: torch.Tensor) -> torch.Tensor:
+ """Decode sampled latents into NCHW image tensors in [0, 1]."""
+
+ if latent.ndim != 4:
+ raise DiffusionCliError(
+ f"Expected NCHW latent tensor, got {tuple(latent.shape)}"
+ )
+ if latent.shape[1] != LATENT_CHANNELS:
+ raise DiffusionCliError(
+ f"Expected {LATENT_CHANNELS} latent channels, "
+ f"got {latent.shape[1]}"
+ )
- def __init__(self, *_args, **_kwargs) -> None:
- raise NotImplementedError("Z-Image VAE decoding is not implemented yet")
+ latent = latent.to(device=self.device, dtype=self.dtype)
+ image = self.model(fluxLatentToVae(latent))
+ expected_shape = (
+ latent.shape[0],
+ VAE_OUT_CHANNELS,
+ latent.shape[2] * LATENT_DOWNSCALE,
+ latent.shape[3] * LATENT_DOWNSCALE,
+ )
+ if tuple(image.shape) != expected_shape:
+ raise DiffusionCliError(
+ "VAE decode produced unexpected shape: "
+ f"expected {expected_shape}, got {tuple(image.shape)}"
+ )
+ if not torch.all(torch.isfinite(image)):
+ raise DiffusionCliError("VAE decode produced NaN or Inf values")
+ return postprocessVaeOutput(image)
diff --git a/diffusion_cli/zimage_model.py b/diffusion_cli/zimage_model.py
index 1c9f5a4..1bbe252 100644
--- a/diffusion_cli/zimage_model.py
+++ b/diffusion_cli/zimage_model.py
@@ -788,7 +788,7 @@ class ZImageModel:
timesteps = timesteps.expand(latent.shape[0])
with torch.inference_mode():
- output = self.model(
+ model_output = self.model(
latent.to(device=self.device, dtype=self.dtype),
timesteps,
conditioning.hidden_states.to(
@@ -799,6 +799,14 @@ class ZImageModel:
conditioning.attention_mask.to(self.device),
)
+ sigma = timesteps.to(
+ device=model_output.device,
+ dtype=model_output.dtype,
+ )
+ sigma = sigma.reshape(sigma.shape[:1] + (1,) * (model_output.ndim - 1))
+ model_input = latent.to(device=self.device, dtype=self.dtype)
+ output = model_input - model_output * sigma
+
if output.shape != latent.shape:
raise DiffusionCliError(
f"Z-Image output shape {tuple(output.shape)} does not match "
diff --git a/tests/test_image_io.py b/tests/test_image_io.py
new file mode 100644
index 0000000..5e1c529
--- /dev/null
+++ b/tests/test_image_io.py
@@ -0,0 +1,34 @@
+import tempfile
+import unittest
+from pathlib import Path
+
+from PIL import Image
+import torch
+
+from diffusion_cli.image_io import outputPaths, saveImages
+
+
+class ImageIoTest(unittest.TestCase):
+ def testOutputPathsAppendBatchIndexBeforeSuffix(self):
+ paths = outputPaths(Path("output.png"), 2)
+
+ self.assertEqual(paths, [Path("output_0000.png"),
+ Path("output_0001.png")])
+
+ def testSaveImagesWritesPngWithExpectedSize(self):
+ images = torch.zeros(1, 3, 4, 5)
+
+ with tempfile.TemporaryDirectory() as temp_dir:
+ output = Path(temp_dir) / "image.png"
+ paths = saveImages(images, output)
+ with Image.open(paths[0]) as image:
+ size = image.size
+ mode = image.mode
+
+ self.assertEqual(paths, [output])
+ self.assertEqual(size, (5, 4))
+ self.assertEqual(mode, "RGB")
+
+
+if __name__ == "__main__":
+ unittest.main()
diff --git a/tests/test_sampling.py b/tests/test_sampling.py
index fcc77fc..3606b61 100644
--- a/tests/test_sampling.py
+++ b/tests/test_sampling.py
@@ -6,8 +6,11 @@ import torch
from diffusion_cli.sampling import (
betaScheduler,
cfgDenoise,
+ halfLogSnrToSigma,
latentShape,
+ offsetFirstSigmaForSnr,
sampleSeeds3,
+ sigmaToHalfLogSnr,
)
@@ -56,6 +59,23 @@ class SamplingTest(unittest.TestCase):
self.assertIs(model.calls[0][1], POSITIVE)
self.assertIs(model.calls[1][1], NEGATIVE)
+ def testFlowLogSnrConversionsRoundTrip(self):
+ sigma = torch.tensor([0.2, 0.5, 0.8])
+
+ half_log_snr = sigmaToHalfLogSnr(sigma)
+ output = halfLogSnrToSigma(half_log_snr)
+
+ self.assertTrue(torch.allclose(output, sigma))
+
+ def testOffsetFirstSigmaAvoidsInfiniteLogSnr(self):
+ sigmas = torch.tensor([1.0, 0.5, 0.0])
+
+ output = offsetFirstSigmaForSnr(sigmas)
+
+ self.assertLess(output[0].item(), 1.0)
+ self.assertEqual(output[1].item(), 0.5)
+ self.assertEqual(output[2].item(), 0.0)
+
def testSampleSeeds3RunsAndRemainsFinite(self):
model = FakeModel()
latent = torch.zeros(1, 1, 2, 2)
diff --git a/tests/test_vae.py b/tests/test_vae.py
new file mode 100644
index 0000000..b9fce29
--- /dev/null
+++ b/tests/test_vae.py
@@ -0,0 +1,113 @@
+import tempfile
+import unittest
+from pathlib import Path
+
+import torch
+from safetensors.torch import save_file
+
+from diffusion_cli.errors import DiffusionCliError
+from diffusion_cli.vae import (
+ FLUX_SCALE_FACTOR,
+ FLUX_SHIFT_FACTOR,
+ VaeConfig,
+ VaeDecoder,
+ ZImageVae,
+ detectVaeConfig,
+ fluxLatentToVae,
+ postprocessVaeOutput,
+)
+
+
+class FakeDecoder(torch.nn.Module):
+ def __init__(self):
+ super().__init__()
+ self.last_latent = None
+
+ def forward(self, latent):
+ self.last_latent = latent.detach().clone()
+ return torch.full(
+ (latent.shape[0], 3, latent.shape[2] * 8, latent.shape[3] * 8),
+ 0.5,
+ device=latent.device,
+ dtype=latent.dtype,
+ )
+
+
+class VaeTest(unittest.TestCase):
+ def testFluxLatentToVaeAppliesComfyScaling(self):
+ latent = torch.tensor([0.0, FLUX_SCALE_FACTOR])
+
+ output = fluxLatentToVae(latent)
+
+ expected = torch.tensor([
+ FLUX_SHIFT_FACTOR,
+ 1.0 + FLUX_SHIFT_FACTOR,
+ ])
+ self.assertTrue(torch.allclose(output, expected))
+
+ def testPostprocessClampsToImageRange(self):
+ image = torch.tensor([-3.0, -1.0, 0.0, 1.0, 3.0])
+
+ output = postprocessVaeOutput(image)
+
+ self.assertTrue(torch.equal(
+ output,
+ torch.tensor([0.0, 0.0, 0.5, 1.0, 1.0]),
+ ))
+
+ def testDetectVaeConfigRejectsUnknownKeys(self):
+ with self.assertRaises(DiffusionCliError) as context:
+ detectVaeConfig({})
+
+ self.assertIn("Unsupported VAE state dict", str(context.exception))
+
+ def testDecodeValidatesShapeAndRange(self):
+ model = FakeDecoder()
+ vae = ZImageVae(
+ model_path=None,
+ device=torch.device("cpu"),
+ dtype=torch.float32,
+ model=model,
+ )
+ latent = torch.zeros(1, 16, 2, 3)
+
+ image = vae.decode(latent)
+
+ self.assertEqual(image.shape, (1, 3, 16, 24))
+ self.assertTrue(torch.all((image >= 0) & (image <= 1)))
+ self.assertTrue(torch.allclose(
+ model.last_latent,
+ torch.full_like(latent, FLUX_SHIFT_FACTOR),
+ ))
+
+ def testLoadsDecoderStateDictFromSafetensors(self):
+ config = VaeConfig(
+ channels=32,
+ channel_multipliers=(1,),
+ num_res_blocks=0,
+ latent_channels=16,
+ out_channels=3,
+ resolution=8,
+ )
+ decoder = VaeDecoder(config, dtype=torch.float32)
+ state = {
+ f"decoder.{key}": value.detach().clone()
+ for key, value in decoder.state_dict().items()
+ }
+
+ with tempfile.TemporaryDirectory() as temp_dir:
+ path = Path(temp_dir) / "tiny_vae.safetensors"
+ save_file(state, path)
+ vae = ZImageVae(
+ path,
+ torch.device("cpu"),
+ torch.float32,
+ model=None,
+ decoder_config=config,
+ )
+
+ self.assertIsInstance(vae.model, VaeDecoder)
+
+
+if __name__ == "__main__":
+ unittest.main()
diff --git a/tests/test_zimage_model.py b/tests/test_zimage_model.py
index 99bbb86..d44ff00 100644
--- a/tests/test_zimage_model.py
+++ b/tests/test_zimage_model.py
@@ -107,6 +107,36 @@ class ZImageModelTest(unittest.TestCase):
self.assertEqual(output.shape, latent.shape)
self.assertTrue(torch.all(torch.isfinite(output)))
+ def testWrapperConvertsFlowOutputToDenoisedPrediction(self):
+ class ConstantVelocityModel(torch.nn.Module):
+ def forward(
+ self,
+ latent,
+ timesteps,
+ context,
+ num_tokens,
+ attention_mask,
+ ):
+ del timesteps, context, num_tokens, attention_mask
+ return torch.ones_like(latent) * 2.0
+
+ wrapper = ZImageModel(
+ model_path=None,
+ device=torch.device("cpu"),
+ dtype=torch.float32,
+ model=ConstantVelocityModel(),
+ )
+ conditioning = TextConditioning(
+ hidden_states=torch.zeros(1, 2, 8),
+ attention_mask=torch.ones(1, 2),
+ token_count=2,
+ )
+ latent = torch.ones(1, 4, 2, 2) * 3.0
+
+ output = wrapper.forward(latent, torch.tensor([0.25]), conditioning)
+
+ self.assertTrue(torch.equal(output, torch.ones_like(latent) * 2.5))
+
if __name__ == "__main__":
unittest.main()