test

eval.py

@@ -0,0 +1,515 @@
+import os
+import json
+import random
+from typing import Dict, List, Tuple, Optional, Any
+
+import numpy as np
+from PIL import Image
+from tqdm import tqdm
+
+import torch
+import torch.nn.functional as F
+from torch.utils.data import Dataset, DataLoader
+from torchvision.transforms import Compose, Resize, ToTensor, CenterCrop
+from torchvision.utils import save_image
+import lpips
+
+from diffusers import (
+    AutoencoderKL,
+    AutoencoderKLWan,
+    AutoencoderKLLTXVideo,
+    AutoencoderKLQwenImage
+)
+
+from scipy.stats import skew, kurtosis
+
+
+# ========================== Конфиг ==========================
+DEVICE = "cuda"
+DTYPE = torch.float16
+IMAGE_FOLDER = "/home/recoilme/dataset/alchemist"
+MIN_SIZE = 1280
+CROP_SIZE = 512
+BATCH_SIZE = 5
+MAX_IMAGES = 0
+NUM_WORKERS = 4
+SAMPLES_DIR = "test"
+
+VAE_LIST = [
+    ("SD15 VAE", AutoencoderKL, "stable-diffusion-v1-5/stable-diffusion-v1-5", "vae"),
+    ("SDXL VAE fp16 fix", AutoencoderKL, "madebyollin/sdxl-vae-fp16-fix", None),
+    ("AiArtLab/sdxl_vae", AutoencoderKL, "AiArtLab/sdxl_vae", "vae"),
+    ("LTX-Video VAE", AutoencoderKLLTXVideo, "Lightricks/LTX-Video", "vae"),
+    ("Wan2.2-TI2V-5B", AutoencoderKLWan, "Wan-AI/Wan2.2-TI2V-5B-Diffusers", "vae"),
+    ("AiArtLab/wan16x_vae", AutoencoderKLWan, "AiArtLab/wan16x_vae", "vae"),
+    ("Wan2.2-T2V-A14B", AutoencoderKLWan, "Wan-AI/Wan2.2-T2V-A14B-Diffusers", "vae"),
+    ("QwenImage", AutoencoderKLQwenImage, "Qwen/Qwen-Image", "vae"),
+    ("AuraDiffusion/16ch-vae", AutoencoderKL, "AuraDiffusion/16ch-vae", None),
+    ("FLUX.1-schnell VAE", AutoencoderKL, "black-forest-labs/FLUX.1-schnell", "vae"),
+    ("AiArtLab/simplevae", AutoencoderKL, "AiArtLab/simplevae", "vae"),
+]
+
+
+# ========================== Утилиты ==========================
+def to_neg1_1(x: torch.Tensor) -> torch.Tensor:
+    return x * 2 - 1
+
+
+def to_0_1(x: torch.Tensor) -> torch.Tensor:
+    return (x + 1) * 0.5
+
+
+def safe_psnr(mse: float) -> float:
+    if mse <= 1e-12:
+        return float("inf")
+    return 10.0 * float(np.log10(1.0 / mse))
+
+
+def is_video_like_vae(vae) -> bool:
+    # Wan и LTX-Video ждут [B, C, T, H, W]
+    return isinstance(vae, (AutoencoderKLWan, AutoencoderKLLTXVideo,AutoencoderKLQwenImage))
+
+
+def add_time_dim_if_needed(x: torch.Tensor, vae) -> torch.Tensor:
+    if is_video_like_vae(vae) and x.ndim == 4:
+        return x.unsqueeze(2)  # -> [B, C, 1, H, W]
+    return x
+
+
+def strip_time_dim_if_possible(x: torch.Tensor, vae) -> torch.Tensor:
+    if is_video_like_vae(vae) and x.ndim == 5 and x.shape[2] == 1:
+        return x.squeeze(2)  # -> [B, C, H, W]
+    return x
+
+
+@torch.no_grad()
+def sobel_edge_l1(real_0_1: torch.Tensor, fake_0_1: torch.Tensor) -> float:
+    real = to_neg1_1(real_0_1)
+    fake = to_neg1_1(fake_0_1)
+    kx = torch.tensor([[-1, 0, 1], [-2, 0, 2], [-1, 0, 1]], dtype=torch.float32, device=real.device).view(1, 1, 3, 3)
+    ky = torch.tensor([[-1, -2, -1], [0, 0, 0], [1, 2, 1]], dtype=torch.float32, device=real.device).view(1, 1, 3, 3)
+    C = real.shape[1]
+    kx = kx.to(real.dtype).repeat(C, 1, 1, 1)
+    ky = ky.to(real.dtype).repeat(C, 1, 1, 1)
+
+    def grad_mag(x):
+        gx = F.conv2d(x, kx, padding=1, groups=C)
+        gy = F.conv2d(x, ky, padding=1, groups=C)
+        return torch.sqrt(gx * gx + gy * gy + 1e-12)
+
+    return F.l1_loss(grad_mag(fake), grad_mag(real)).item()
+
+
+def flatten_channels(x: torch.Tensor) -> torch.Tensor:
+    # -> [C, N*H*W] или [C, N*T*H*W]
+    if x.ndim == 4:
+        return x.permute(1, 0, 2, 3).reshape(x.shape[1], -1)
+    elif x.ndim == 5:
+        return x.permute(1, 0, 2, 3, 4).reshape(x.shape[1], -1)
+    else:
+        raise ValueError(f"Unexpected tensor ndim={x.ndim}")
+
+
+def _to_numpy_1d(x: Any) -> Optional[np.ndarray]:
+    if x is None:
+        return None
+    if isinstance(x, (int, float)):
+        return None
+    if isinstance(x, torch.Tensor):
+        x = x.detach().cpu().float().numpy()
+    elif isinstance(x, (list, tuple)):
+        x = np.array(x, dtype=np.float32)
+    elif isinstance(x, np.ndarray):
+        x = x.astype(np.float32, copy=False)
+    else:
+        return None
+    x = x.reshape(-1)
+    return x
+
+
+def _to_float(x: Any) -> Optional[float]:
+    if x is None:
+        return None
+    if isinstance(x, (int, float)):
+        return float(x)
+    if isinstance(x, np.ndarray) and x.size == 1:
+        return float(x.item())
+    if isinstance(x, torch.Tensor) and x.numel() == 1:
+        return float(x.item())
+    return None
+
+
+def get_norm_tensors_and_summary(vae, latent_like: torch.Tensor):
+    """
+    Нормализация латентов: глобальная и поканальная.
+    Применение: сначала глобальная (scalar), затем поканальная (vector).
+    Если в конфиге есть несколько ключей — аккумулируем.
+    """
+    cfg = getattr(vae, "config", vae)
+
+    scale_keys = [
+        "latents_std"
+    ]
+    shift_keys = [
+        "latents_mean"
+    ]
+
+    C = latent_like.shape[1]
+    nd = latent_like.ndim  # 4 или 5
+    dev = latent_like.device
+    dt = latent_like.dtype
+
+    scale_global = getattr(vae.config, "scaling_factor", 1.0)
+    shift_global = getattr(vae.config, "shift_factor", 0.0)
+    if scale_global is None:
+        scale_global = 1.0
+    if shift_global is None:
+        shift_global = 0.0
+    
+    scale_channel = np.ones(C, dtype=np.float32)
+    shift_channel = np.zeros(C, dtype=np.float32)
+
+    for k in scale_keys:
+        v = getattr(cfg, k, None)
+        if v is None:
+            continue
+        vec = _to_numpy_1d(v)
+        if vec is not None and vec.size == C:
+            scale_channel *= vec
+        else:
+            s = _to_float(v)
+            if s is not None:
+                scale_global *= s
+
+    for k in shift_keys:
+        v = getattr(cfg, k, None)
+        if v is None:
+            continue
+        vec = _to_numpy_1d(v)
+        if vec is not None and vec.size == C:
+            shift_channel += vec
+        else:
+            s = _to_float(v)
+            if s is not None:
+                shift_global += s
+
+    g_shape = [1] * nd
+    c_shape = [1] * nd
+    c_shape[1] = C
+
+    t_scale_g = torch.tensor(scale_global, dtype=dt, device=dev).view(*g_shape)
+    t_shift_g = torch.tensor(shift_global, dtype=dt, device=dev).view(*g_shape)
+    t_scale_c = torch.from_numpy(scale_channel).to(device=dev, dtype=dt).view(*c_shape)
+    t_shift_c = torch.from_numpy(shift_channel).to(device=dev, dtype=dt).view(*c_shape)
+
+    summary = {
+        "scale_global": float(scale_global),
+        "shift_global": float(shift_global),
+        "scale_channel_min": float(scale_channel.min()),
+        "scale_channel_mean": float(scale_channel.mean()),
+        "scale_channel_max": float(scale_channel.max()),
+        "shift_channel_min": float(shift_channel.min()),
+        "shift_channel_mean": float(shift_channel.mean()),
+        "shift_channel_max": float(shift_channel.max()),
+    }
+    return t_shift_g, t_scale_g, t_shift_c, t_scale_c, summary
+
+
+@torch.no_grad()
+def kl_divergence_per_image(mu: torch.Tensor, logvar: torch.Tensor) -> torch.Tensor:
+    kl_map = -0.5 * (1 + logvar - mu.pow(2) - logvar.exp())  # [B, ...]
+    return kl_map.float().view(kl_map.shape[0], -1).mean(dim=1)  # [B]
+
+
+def sanitize_filename(name: str) -> str:
+    name = name.replace("/", "_").replace("\\", "_").replace(" ", "_")
+    return "".join(ch if (ch.isalnum() or ch in "._-") else "_" for ch in name)
+
+
+# ========================== Датасет ==========================
+class ImageFolderDataset(Dataset):
+    def __init__(self, root_dir: str, extensions=(".png", ".jpg", ".jpeg", ".webp"), min_size=1024, crop_size=512, limit=None):
+        paths = []
+        for root, _, files in os.walk(root_dir):
+            for fname in files:
+                if fname.lower().endswith(extensions):
+                    paths.append(os.path.join(root, fname))
+        if limit:
+            paths = paths[:limit]
+
+        valid = []
+        for p in tqdm(paths, desc="Проверяем файлы"):
+            try:
+                with Image.open(p) as im:
+                    im.verify()
+                valid.append(p)
+            except Exception:
+                pass
+        if not valid:
+            raise RuntimeError(f"Нет валидных изображений в {root_dir}")
+        random.shuffle(valid)
+        self.paths = valid
+        print(f"Найдено {len(self.paths)} изображений")
+
+        self.transform = Compose([
+            Resize(min_size),
+            CenterCrop(crop_size),
+            ToTensor(),  # 0..1, float32
+        ])
+
+    def __len__(self):
+        return len(self.paths)
+
+    def __getitem__(self, idx):
+        with Image.open(self.paths[idx]) as img:
+            img = img.convert("RGB")
+            return self.transform(img)
+
+
+# ========================== Основное ==========================
+def main():
+    torch.set_grad_enabled(False)
+    os.makedirs(SAMPLES_DIR, exist_ok=True)
+
+    dataset = ImageFolderDataset(IMAGE_FOLDER, min_size=MIN_SIZE, crop_size=CROP_SIZE, limit=MAX_IMAGES)
+    loader = DataLoader(dataset, batch_size=BATCH_SIZE, shuffle=False, num_workers=NUM_WORKERS, pin_memory=True)
+
+    lpips_net = lpips.LPIPS(net="vgg").to(DEVICE).eval()
+
+    # Загрузка VAE
+    vaes: List[Tuple[str, object]] = []
+    print("\nЗагрузка VAE...")
+    for human_name, vae_class, model_path, subfolder in VAE_LIST:
+        try:
+            vae = vae_class.from_pretrained(model_path, subfolder=subfolder, torch_dtype=DTYPE)
+            vae = vae.to(DEVICE).eval()
+            vaes.append((human_name, vae))
+            print(f"  ✅ {human_name}")
+        except Exception as e:
+            print(f"  ❌ {human_name}: {e}")
+
+    if not vaes:
+        print("Нет успешно загруженных VAE. Выходим.")
+        return
+
+    # Агрегаторы
+    per_model_metrics: Dict[str, Dict[str, float]] = {
+        name: {"mse": 0.0, "psnr": 0.0, "lpips": 0.0, "edge": 0.0, "kl": 0.0, "count": 0.0}
+        for name, _ in vaes
+    }
+
+    buffers_zmodel: Dict[str, List[torch.Tensor]] = {name: [] for name, _ in vaes}
+    norm_summaries: Dict[str, Dict[str, float]] = {}
+
+    # Флаг для сохранения первой картинки
+    saved_first_for: Dict[str, bool] = {name: False for name, _ in vaes}
+
+    for batch_0_1 in tqdm(loader, desc="Батчи"):
+        batch_0_1 = batch_0_1.to(DEVICE, torch.float32)
+        batch_neg1_1 = to_neg1_1(batch_0_1).to(DTYPE)
+
+        for model_name, vae in vaes:
+            x_in = add_time_dim_if_needed(batch_neg1_1, vae)
+
+            posterior = vae.encode(x_in).latent_dist
+            mu, logvar = posterior.mean, posterior.logvar
+
+            # Реконструкция (детерминированно)
+            z_raw_mode = posterior.mode()
+            x_dec = vae.decode(z_raw_mode).sample  # [-1, 1]
+            x_dec = strip_time_dim_if_possible(x_dec, vae)
+            x_rec_0_1 = to_0_1(x_dec.float()).clamp(0, 1)
+
+            # Латенты для UNet: global -> channelwise
+            z_raw_sample = posterior.sample()
+            t_shift_g, t_scale_g, t_shift_c, t_scale_c, summary = get_norm_tensors_and_summary(vae, z_raw_sample)
+ 
+            if model_name not in norm_summaries:
+                norm_summaries[model_name] = summary
+
+            z_tmp = (z_raw_sample - t_shift_g) * t_scale_g
+            z_model = (z_tmp - t_shift_c) * t_scale_c
+            z_model = strip_time_dim_if_possible(z_model, vae)
+
+            buffers_zmodel[model_name].append(z_model.detach().to("cpu", torch.float32))
+
+            # Сохранить первую картинку (оригинал и реконструкцию) для каждого VAE
+            if not saved_first_for[model_name]:
+                safe = sanitize_filename(model_name)
+                orig_path = os.path.join(SAMPLES_DIR, f"{safe}_original.png")
+                dec_path  = os.path.join(SAMPLES_DIR, f"{safe}_decoded.png")
+                save_image(batch_0_1[0:1].cpu(), orig_path)
+                save_image(x_rec_0_1[0:1].cpu(),  dec_path)
+                saved_first_for[model_name] = True
+
+            # Метрики по картинкам
+            B = batch_0_1.shape[0]
+            for i in range(B):
+                gt = batch_0_1[i:i+1]
+                rec = x_rec_0_1[i:i+1]
+
+                mse = F.mse_loss(gt, rec).item()
+                psnr = safe_psnr(mse)
+                lp = float(lpips_net(gt, rec, normalize=True).mean().item())
+                edge = sobel_edge_l1(gt, rec)
+
+                per_model_metrics[model_name]["mse"] += mse
+                per_model_metrics[model_name]["psnr"] += psnr
+                per_model_metrics[model_name]["lpips"] += lp
+                per_model_metrics[model_name]["edge"] += edge
+
+            # KL per-image
+            kl_pi = kl_divergence_per_image(mu, logvar)  # [B]
+            per_model_metrics[model_name]["kl"] += float(kl_pi.sum().item())
+            per_model_metrics[model_name]["count"] += B
+
+    # Усреднение метрик
+    for name in per_model_metrics:
+        c = max(1.0, per_model_metrics[name]["count"])
+        for k in ["mse", "psnr", "lpips", "edge", "kl"]:
+            per_model_metrics[name][k] /= c
+
+    # Подсчёт статистик латентов и нормальности
+    per_model_latent_stats = {}
+    for name, _ in vaes:
+        if not buffers_zmodel[name]:
+            continue
+        Z = torch.cat(buffers_zmodel[name], dim=0)   # [N, C, H, W]
+
+        # Глобальные
+        z_min = float(Z.min().item())
+        z_mean = float(Z.mean().item())
+        z_max = float(Z.max().item())
+        z_std = float(Z.std(unbiased=True).item())
+
+        # Пер-канально: skew/kurtosis
+        Z_ch = flatten_channels(Z).numpy()  # [C, *]
+        C = Z_ch.shape[0]
+        sk = np.zeros(C, dtype=np.float64)
+        ku = np.zeros(C, dtype=np.float64)
+        for c in range(C):
+            v = Z_ch[c]
+            sk[c] = float(skew(v, bias=False))
+            ku[c] = float(kurtosis(v, fisher=True, bias=False))
+
+        skew_min, skew_mean, skew_max = float(sk.min()), float(sk.mean()), float(sk.max())
+        kurt_min, kurt_mean, kurt_max = float(ku.min()), float(ku.mean()), float(ku.max())
+        mean_abs_skew = float(np.mean(np.abs(sk)))
+        mean_abs_kurt = float(np.mean(np.abs(ku)))
+
+        per_model_latent_stats[name] = {
+            "Z_min": z_min, "Z_mean": z_mean, "Z_max": z_max, "Z_std": z_std,
+            "skew_min": skew_min, "skew_mean": skew_mean, "skew_max": skew_max,
+            "kurt_min": kurt_min, "kurt_mean": kurt_mean, "kurt_max": kurt_max,
+            "mean_abs_skew": mean_abs_skew, "mean_abs_kurt": mean_abs_kurt,
+        }
+
+    # Печать параметров нормализации (shift/scale)
+    print("\n=== Параметры нормализации латентов (как применялись) ===")
+    for name, _ in vaes:
+        if name not in norm_summaries:
+            continue
+        s = norm_summaries[name]
+        print(
+            f"{name:26s} | "
+            f"shift_g={s['shift_global']:.6g} scale_g={s['scale_global']:.6g} | "
+            f"shift_c[min/mean/max]=[{s['shift_channel_min']:.6g}, {s['shift_channel_mean']:.6g}, {s['shift_channel_max']:.6g}] | "
+            f"scale_c[min/mean/max]=[{s['scale_channel_min']:.6g}, {s['scale_channel_mean']:.6g}, {s['scale_channel_max']:.6g}]"
+        )
+
+    # Абсолютные метрики
+    print("\n=== Абсолютные метрики реконструкции и латентов ===")
+    for name, _ in vaes:
+        if name not in per_model_latent_stats:
+            continue
+        m = per_model_metrics[name]
+        s = per_model_latent_stats[name]
+        print(
+            f"{name:26s} | "
+            f"MSE={m['mse']:.3e} PSNR={m['psnr']:.2f} LPIPS={m['lpips']:.3f} Edge={m['edge']:.3f} KL={m['kl']:.3f} | "
+            f"Z[min/mean/max/std]=[{s['Z_min']:.3f}, {s['Z_mean']:.3f}, {s['Z_max']:.3f}, {s['Z_std']:.3f}] | "
+            f"Skew[min/mean/max]=[{s['skew_min']:.3f}, {s['skew_mean']:.3f}, {s['skew_max']:.3f}] | "
+            f"Kurt[min/mean/max]=[{s['kurt_min']:.3f}, {s['kurt_mean']:.3f}, {s['kurt_max']:.3f}]"
+        )
+
+    # Сравнение с первой моделью
+    baseline = vaes[0][0]
+    print("\n=== Сравнение с первой моделью (проценты) ===")
+    print(f"| {'Модель':26s} | {'MSE':>9s} | {'PSNR':>9s} | {'LPIPS':>9s} | {'Edge':>9s} | {'Skew|0':>9s} | {'Kurt|0':>9s} |")
+    print(f"|{'-'*28}|{'-'*11}|{'-'*11}|{'-'*11}|{'-'*11}|{'-'*11}|{'-'*11}|")
+
+    b_m = per_model_metrics[baseline]
+    b_s = per_model_latent_stats[baseline]
+
+    for name, _ in vaes:
+        m = per_model_metrics[name]
+        s = per_model_latent_stats[name]
+
+        mse_pct  = (b_m["mse"] / max(1e-12, m["mse"])) * 100.0               # меньше лучше
+        psnr_pct = (m["psnr"] / max(1e-12, b_m["psnr"])) * 100.0             # больше лучше
+        lpips_pct= (b_m["lpips"] / max(1e-12, m["lpips"])) * 100.0           # меньше лучше
+        edge_pct = (b_m["edge"] / max(1e-12, m["edge"])) * 100.0             # меньше лучше
+
+        skew0_pct = (b_s["mean_abs_skew"] / max(1e-12, s["mean_abs_skew"])) * 100.0
+        kurt0_pct = (b_s["mean_abs_kurt"] / max(1e-12, s["mean_abs_kurt"])) * 100.0
+
+        if name == baseline:
+            print(f"| {name:26s} | {'100%':>9s} | {'100%':>9s} | {'100%':>9s} | {'100%':>9s} | {'100%':>9s} | {'100%':>9s} |")
+        else:
+            print(f"| {name:26s} | {mse_pct:8.1f}% | {psnr_pct:8.1f}% | {lpips_pct:8.1f}% | {edge_pct:8.1f}% | {skew0_pct:8.1f}% | {kurt0_pct:8.1f}% |")
+
+    # ========================== Коррекции для последнего VAE + сохранение в JSON ==========================
+    last_name = vaes[-1][0]
+    if buffers_zmodel[last_name]:
+        Z = torch.cat(buffers_zmodel[last_name], dim=0)  # [N, C, H, W]
+
+        # Глобальная коррекция (по всем каналам/пикселям)
+        z_mean = float(Z.mean().item())
+        z_std  = float(Z.std(unbiased=True).item())
+        correction_global = {
+            "shift": -z_mean,
+            "scale": (1.0 / z_std) if z_std > 1e-12 else 1.0
+        }
+
+        # Поканальная коррекция
+        Z_ch = flatten_channels(Z)  # [C, M]
+        ch_means_t = Z_ch.mean(dim=1)                       # [C]
+        ch_stds_t  = Z_ch.std(dim=1, unbiased=True) + 1e-12 # [C]
+        ch_means = [float(x) for x in ch_means_t.tolist()]
+        ch_stds  = [float(x) for x in ch_stds_t.tolist()]
+
+        correction_per_channel = [
+            {"shift": float(-m), "scale": float(1.0 / s)}
+            for m, s in zip(ch_means, ch_stds)
+        ]
+
+        print(f"\n=== Доп. коррекция для {last_name} (поверх VAE-нормализации) ===")
+        print(f"global_correction = {correction_global}")
+        print(f"channelwise_means = {ch_means}")
+        print(f"channelwise_stds  = {ch_stds}")
+        print(f"channelwise_correction = {correction_per_channel}")
+
+        # Сохранение в JSON
+        json_path = os.path.join(SAMPLES_DIR, f"{sanitize_filename(last_name)}_correction.json")
+        to_save = {
+            "model_name": last_name,
+            "vae_normalization_summary": norm_summaries.get(last_name, {}),
+            "global_correction": correction_global,
+            "per_channel_means": ch_means,
+            "per_channel_stds": ch_stds,
+            "per_channel_correction": correction_per_channel,
+            "apply_order": {
+                "forward": "z_model -> (z - global_shift)*global_scale -> (per-channel: (z - mean_c)/std_c)",
+                "inverse":  "z_corr -> (per-channel: z*std_c + mean_c) -> (z/global_scale + global_shift)"
+            },
+            "note": "Эти коэффициенты рассчитаны по z_model (после встроенных VAE shift/scale), чтобы привести распределение к N(0,1)."
+        }
+        with open(json_path, "w", encoding="utf-8") as f:
+            json.dump(to_save, f, ensure_ascii=False, indent=2)
+        print("Corrections JSON saved to:", os.path.abspath(json_path))
+
+    print("\n✅ Готово. Сэмплы сохранены в:", os.path.abspath(SAMPLES_DIR))
+
+
+if __name__ == "__main__":
+    main()

test/AiArtLab_simplevae_correction.json

@@ -0,0 +1,124 @@
+{
+  "model_name": "AiArtLab/simplevae",
+  "vae_normalization_summary": {
+    "scale_global": 1.0,
+    "shift_global": 0.0,
+    "scale_channel_min": 1.0,
+    "scale_channel_mean": 1.0,
+    "scale_channel_max": 1.0,
+    "shift_channel_min": 0.0,
+    "shift_channel_mean": 0.0,
+    "shift_channel_max": 0.0
+  },
+  "global_correction": {
+    "shift": 0.060732152312994,
+    "scale": 1.035888502366052
+  },
+  "per_channel_means": [
+    -0.17934872210025787,
+    0.08837347477674484,
+    -0.18308678269386292,
+    0.2742660641670227,
+    0.002165613230317831,
+    0.03449296951293945,
+    -0.16568207740783691,
+    -0.26266154646873474,
+    -0.1797001212835312,
+    -0.12722325325012207,
+    -0.03235034644603729,
+    -0.04355641081929207,
+    -0.0853145644068718,
+    0.08101209998130798,
+    -0.12256482243537903,
+    -0.07053600996732712
+  ],
+  "per_channel_stds": [
+    0.9051793217658997,
+    0.8686285614967346,
+    0.8835359811782837,
+    1.3679808378219604,
+    0.889316201210022,
+    1.010545253753662,
+    0.9772343635559082,
+    0.9337192177772522,
+    0.988460123538971,
+    0.9164251685142517,
+    0.8944668769836426,
+    0.9558082222938538,
+    0.9226727485656738,
+    0.8971477746963501,
+    0.9111431837081909,
+    0.8696402907371521
+  ],
+  "per_channel_correction": [
+    {
+      "shift": 0.17934872210025787,
+      "scale": 1.1047534736532825
+    },
+    {
+      "shift": -0.08837347477674484,
+      "scale": 1.1512400631599071
+    },
+    {
+      "shift": 0.18308678269386292,
+      "scale": 1.131815818826529
+    },
+    {
+      "shift": -0.2742660641670227,
+      "scale": 0.7310043915470018
+    },
+    {
+      "shift": -0.002165613230317831,
+      "scale": 1.124459442703708
+    },
+    {
+      "shift": -0.03449296951293945,
+      "scale": 0.9895647882027134
+    },
+    {
+      "shift": 0.16568207740783691,
+      "scale": 1.023295984354514
+    },
+    {
+      "shift": 0.26266154646873474,
+      "scale": 1.0709857749105038
+    },
+    {
+      "shift": 0.1797001212835312,
+      "scale": 1.0116745999016257
+    },
+    {
+      "shift": 0.12722325325012207,
+      "scale": 1.091196569405927
+    },
+    {
+      "shift": 0.03235034644603729,
+      "scale": 1.1179843834712364
+    },
+    {
+      "shift": 0.04355641081929207,
+      "scale": 1.046234983833985
+    },
+    {
+      "shift": 0.0853145644068718,
+      "scale": 1.0838078848156445
+    },
+    {
+      "shift": -0.08101209998130798,
+      "scale": 1.1146435717777503
+    },
+    {
+      "shift": 0.12256482243537903,
+      "scale": 1.0975223410333572
+    },
+    {
+      "shift": 0.07053600996732712,
+      "scale": 1.1499007240710388
+    }
+  ],
+  "apply_order": {
+    "forward": "z_model -> (z - global_shift)*global_scale -> (per-channel: (z - mean_c)/std_c)",
+    "inverse": "z_corr -> (per-channel: z*std_c + mean_c) -> (z/global_scale + global_shift)"
+  },
+  "note": "Эти коэффициенты рассчитаны по z_model (после встроенных VAE shift/scale), чтобы привести распределение к N(0,1)."
+}

test/FLUX.1-schnell_VAE_correction.json

@@ -0,0 +1,124 @@
+{
+  "model_name": "FLUX.1-schnell VAE",
+  "vae_normalization_summary": {
+    "scale_global": 0.3611,
+    "shift_global": 0.1159,
+    "scale_channel_min": 1.0,
+    "scale_channel_mean": 1.0,
+    "scale_channel_max": 1.0,
+    "shift_channel_min": 0.0,
+    "shift_channel_mean": 0.0,
+    "shift_channel_max": 0.0
+  },
+  "global_correction": {
+    "shift": 0.033536698669195175,
+    "scale": 1.0893412605691914
+  },
+  "per_channel_means": [
+    -0.2011537104845047,
+    -0.23432493209838867,
+    0.1418181210756302,
+    -0.051936931908130646,
+    0.05777733772993088,
+    0.2514500916004181,
+    -0.3927314877510071,
+    0.23873785138130188,
+    -0.1971667855978012,
+    0.5086520910263062,
+    0.18148651719093323,
+    0.18726322054862976,
+    -0.029747387394309044,
+    -0.30456721782684326,
+    -0.4305008351802826,
+    -0.2616432309150696
+  ],
+  "per_channel_stds": [
+    0.7963889241218567,
+    1.28511381149292,
+    0.8507877588272095,
+    0.5032204985618591,
+    0.527312159538269,
+    0.6365501284599304,
+    1.104862928390503,
+    0.8536051511764526,
+    0.6959906220436096,
+    1.011374592781067,
+    0.6488218307495117,
+    0.8073245882987976,
+    0.9629114866256714,
+    1.241254448890686,
+    0.7992448806762695,
+    0.9001436829566956
+  ],
+  "per_channel_correction": [
+    {
+      "shift": 0.2011537104845047,
+      "scale": 1.2556678900358345
+    },
+    {
+      "shift": 0.23432493209838867,
+      "scale": 0.7781411973452356
+    },
+    {
+      "shift": -0.1418181210756302,
+      "scale": 1.1753812741481802
+    },
+    {
+      "shift": 0.051936931908130646,
+      "scale": 1.987200447632547
+    },
+    {
+      "shift": -0.05777733772993088,
+      "scale": 1.8964099004954318
+    },
+    {
+      "shift": -0.2514500916004181,
+      "scale": 1.5709681850499353
+    },
+    {
+      "shift": 0.3927314877510071,
+      "scale": 0.9050896489546799
+    },
+    {
+      "shift": -0.23873785138130188,
+      "scale": 1.1715018338652052
+    },
+    {
+      "shift": 0.1971667855978012,
+      "scale": 1.4368009687598085
+    },
+    {
+      "shift": -0.5086520910263062,
+      "scale": 0.9887533334708467
+    },
+    {
+      "shift": -0.18148651719093323,
+      "scale": 1.5412551683792932
+    },
+    {
+      "shift": -0.18726322054862976,
+      "scale": 1.2386591644721363
+    },
+    {
+      "shift": 0.029747387394309044,
+      "scale": 1.0385170536331412
+    },
+    {
+      "shift": 0.30456721782684326,
+      "scale": 0.8056365887700978
+    },
+    {
+      "shift": 0.4305008351802826,
+      "scale": 1.2511809886775433
+    },
+    {
+      "shift": 0.2616432309150696,
+      "scale": 1.1109337530596306
+    }
+  ],
+  "apply_order": {
+    "forward": "z_model -> (z - global_shift)*global_scale -> (per-channel: (z - mean_c)/std_c)",
+    "inverse": "z_corr -> (per-channel: z*std_c + mean_c) -> (z/global_scale + global_shift)"
+  },
+  "note": "Эти коэффициенты рассчитаны по z_model (после встроенных VAE shift/scale), чтобы привести распределение к N(0,1)."
+}