Spaces:
Running on Zero
Running on Zero
| """Reference-image encoding and result decoding in the MOSS latent space. | |
| Port of https://github.com/Chenkin-x/ComfyUI-Moss-Image-Edit | |
| (moss_runtime/processing.py, GPLv3), plus the reference resize logic from | |
| that repo's `nodes.py::_reference_tensor` as `prepare_reference`. The | |
| ComfyUI component loader (`load_encoding_components`) is dropped — the Space | |
| reuses the VAE already loaded by the bundled diffusers pipeline. | |
| """ | |
| from __future__ import annotations | |
| import math | |
| def image_to_tensor(image): | |
| """PIL RGB -> CPU FP32 [1, 3, H, W], normalized to [-1, 1].""" | |
| import numpy as np | |
| import torch | |
| pixels = np.array(image.convert("RGB"), dtype=np.float32, copy=True) | |
| return torch.from_numpy(pixels).permute(2, 0, 1).unsqueeze(0).div_(127.5).sub_(1) | |
| def prepare_reference(image, max_edge: int): | |
| """PIL image -> the same Lanczos resize and /16 rounding as CLI inference. | |
| From ComfyUI-Moss-Image-Edit `nodes.py::_reference_tensor`, adapted to | |
| take a PIL image directly instead of a ComfyUI IMAGE batch. | |
| """ | |
| from PIL import Image | |
| picture = image.convert("RGB") | |
| width, height = picture.size | |
| limit = max_edge // 16 * 16 | |
| scale = min(1.0, limit / max(width, height)) | |
| def aligned(value): | |
| return min(limit, max(16, int(math.floor(value * scale / 16 + 0.5)) * 16)) | |
| picture = picture.resize((aligned(width), aligned(height)), Image.Resampling.LANCZOS) | |
| return image_to_tensor(picture) | |
| def _vae_placement(vae): | |
| import torch | |
| parameter = next(vae.parameters(), None) | |
| if parameter is not None: | |
| return parameter.device, parameter.dtype | |
| return vae.bn.running_mean.device, getattr(vae, "dtype", torch.float32) | |
| def _bn_stats(vae, channels: int, device): | |
| import torch | |
| mean = vae.bn.running_mean.to(device=device, dtype=torch.float32).reshape(1, -1, 1, 1) | |
| variance = vae.bn.running_var.to(device=device, dtype=torch.float32).reshape(1, -1, 1, 1) | |
| eps = float(vae.config.batch_norm_eps) | |
| if mean.numel() != channels or variance.numel() != channels: | |
| raise ValueError("Packed latent channels must match FLUX.2 batch-normalization statistics.") | |
| if not math.isfinite(eps) or eps <= 0 or not torch.isfinite(mean).all() or not torch.isfinite(variance).all(): | |
| raise ValueError("Invalid FLUX.2 batch-normalization statistics.") | |
| if (variance < 0).any(): | |
| raise ValueError("FLUX.2 running variance must be nonnegative.") | |
| return mean, (variance + eps).sqrt() | |
| def encode_image(vae, image, *, sample_posterior: bool = False, generator=None): | |
| """Frozen FLUX.2 encode -> 2x2 channel pack -> BN normalization. | |
| Input is [B, 3, H, W] in [-1, 1]; result is FP32 [B, 128, H/16, W/16] | |
| for the official 32-channel VAE. Posterior mode matches Flux2Pipeline. | |
| """ | |
| import torch | |
| if not isinstance(image, torch.Tensor) or image.ndim != 4 or image.shape[1] != 3: | |
| raise ValueError("image must be a [B, 3, H, W] tensor.") | |
| if not image.is_floating_point() or min(image.shape) < 1 or any(s % 16 for s in image.shape[-2:]): | |
| raise ValueError("Image dimensions must be nonempty multiples of 16 and pixels floating point.") | |
| if not torch.isfinite(image).all() or image.min() < -1.001 or image.max() > 1.001: | |
| raise ValueError("Image pixels must be finite and normalized to [-1, 1].") | |
| device, dtype = _vae_placement(vae) | |
| with torch.no_grad(): | |
| vae.eval() | |
| posterior = vae.encode(image.to(device=device, dtype=dtype)).latent_dist | |
| z = posterior.sample(generator=generator) if sample_posterior else posterior.mode() | |
| if z.ndim != 4 or z.shape[-2] % 2 or z.shape[-1] % 2: | |
| raise ValueError("VAE latent spatial dimensions must be even for 2x2 packing.") | |
| batch, channels, height, width = z.shape | |
| # Official Flux2Pipeline._patchify_latents ordering: c, dy, dx. | |
| z = z.reshape(batch, channels, height // 2, 2, width // 2, 2) | |
| z = z.permute(0, 1, 3, 5, 2, 4).reshape(batch, 4 * channels, height // 2, width // 2).float() | |
| mean, std = _bn_stats(vae, z.shape[1], z.device) | |
| normalized = (z - mean) / std | |
| if not torch.isfinite(normalized).all(): | |
| raise ValueError("VAE encoding produced nonfinite latents.") | |
| return normalized | |
| def decode_latents(vae, latents): | |
| """Inverse normalized packed latent transform, then frozen VAE decode. | |
| Returns floating [B, 3, H, W], clamped to [-1, 1] as the official MOSS | |
| pipeline. Normalization arithmetic is FP32 before conversion to VAE dtype. | |
| """ | |
| import torch | |
| if not isinstance(latents, torch.Tensor) or latents.ndim != 4 or min(latents.shape) < 1: | |
| raise ValueError("latents must be a nonempty [B, C, H, W] tensor.") | |
| if not latents.is_floating_point() or latents.shape[1] % 4 or not torch.isfinite(latents).all(): | |
| raise ValueError("Latents must be finite floating point, with channels divisible by four.") | |
| device, dtype = _vae_placement(vae) | |
| with torch.no_grad(): | |
| vae.eval() | |
| z = latents.to(device=device, dtype=torch.float32) | |
| mean, std = _bn_stats(vae, z.shape[1], device) | |
| z = z * std + mean | |
| batch, channels, height, width = z.shape | |
| z = z.reshape(batch, channels // 4, 2, 2, height, width) | |
| z = z.permute(0, 1, 4, 2, 5, 3).reshape(batch, channels // 4, 2 * height, 2 * width) | |
| result = vae.decode(z.to(dtype=dtype), return_dict=False)[0] | |
| if not torch.isfinite(result).all(): | |
| raise ValueError("VAE decoding produced nonfinite images.") | |
| return result.clamp(-1, 1) | |