Add cifar10_cd_ct_lpips: pkl + cm.py + load.py

cifar10_cd_ct_lpips/cifar10_cd_ct_lpips.pkl

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:67cc54eecab61615a8f019399e18b3c8b28d04275ed3007969ed8a65caec6858
+size 225605629

cifar10_cd_ct_lpips/cm.py

@@ -0,0 +1,179 @@
+"""Consistency-Distillation / Consistency-Training generator wrapper for CIFAR10.
+
+Wraps OpenAI's CIFAR10 consistency model (`consistency_models_cifar10`, package
+`jcm`) so it can stand in as a generator with a decode-style API: build once,
+then `decode(z)` maps a noise-space input to a 32x32x3 image in [0, 1].
+
+Latents are flat vectors of shape `(batch, 3072)` (== 32*32*3), drawn from
+N(0, sigma_max**2 * I) and reshaped internally.
+
+Only the generator path is exercised - LPIPS / classifier heads / training-time
+losses are not imported. `distiller_fn` math is inlined (KVE-SDE Karras et al.
+2022 conditioning) so we don't depend on `jcm.sde_lib`.
+"""
+
+from pathlib import Path
+
+import jax
+import jax.numpy as jnp
+import ml_collections
+
+# Registers "ncsnpp" in jcm.models.utils._MODELS via decorator side-effect.
+from jcm.models import ncsnpp  # noqa: F401
+from jcm.models import utils as mutils
+
+Array = jax.Array
+
+IMAGE_SIZE = 32
+NUM_CHANNELS = 3
+NOISE_DIM = IMAGE_SIZE * IMAGE_SIZE * NUM_CHANNELS  # 3072
+SIGMA_MAX = 80.0
+SIGMA_MIN = 0.002
+DATA_STD = 0.5
+
+
+def _default_config() -> ml_collections.ConfigDict:
+    """NCSN++ config used by the CT-LPIPS CIFAR10 checkpoint."""
+    c = ml_collections.ConfigDict()
+    c.data = ml_collections.ConfigDict(
+        dict(
+            dataset="CIFAR10",
+            image_size=IMAGE_SIZE,
+            num_channels=NUM_CHANNELS,
+            random_flip=False,
+            uniform_dequantization=False,
+        )
+    )
+    c.model = ml_collections.ConfigDict(
+        dict(
+            name="ncsnpp",
+            ema_rate=0.9999,
+            normalization="GroupNorm",
+            nonlinearity="swish",
+            nf=128,
+            ch_mult=(2, 2, 2),
+            num_res_blocks=4,
+            attn_resolutions=(16,),
+            resamp_with_conv=True,
+            conditional=True,
+            fir=True,
+            fir_kernel=[1, 3, 3, 1],
+            skip_rescale=True,
+            resblock_type="biggan",
+            progressive="none",
+            progressive_input="residual",
+            progressive_combine="sum",
+            attention_type="ddpm",
+            embedding_type="fourier",
+            init_scale=0.0,
+            fourier_scale=16,
+            conv_size=3,
+            rho=7.0,
+            data_std=DATA_STD,
+            num_scales=18,
+            dropout=0.0,
+            sigma_min=0.02,
+            sigma_max=100,
+            beta_min=0.1,
+            beta_max=20.0,
+            t_min=SIGMA_MIN,
+            t_max=SIGMA_MAX,
+            double_heads=False,
+        )
+    )
+    c.training = ml_collections.ConfigDict(dict(sde="kvesde"))
+    c.sampling = ml_collections.ConfigDict(dict(method="onestep", std=SIGMA_MAX))
+    return c
+
+
+def _build_model_and_shapes(config: ml_collections.ConfigDict, seed: int):
+    import functools
+
+    model_def = functools.partial(mutils.get_model("ncsnpp"), config=config)
+    model = model_def()
+    rng = jax.random.PRNGKey(seed)
+    p_rng, d_rng = jax.random.split(rng)
+    fake_x = jnp.zeros((1, IMAGE_SIZE, IMAGE_SIZE, NUM_CHANNELS))
+    fake_t = jnp.zeros((1,), dtype=jnp.float32)
+    variables = model.init({"params": p_rng, "dropout": d_rng}, fake_x, fake_t)
+    return model, variables["params"]
+
+
+class CD:
+    """Consistency-Distillation / Consistency-Training generator.
+
+    Instantiate with `pretrained` naming a checkpoint kind and an explicit
+    `checkpoint_path` pointing at the slim pickle bundle
+    (`{"params", "meta"}`). Latents are flat vectors of shape `(batch, 3072)`
+    for consistency with the target's expected `n_dim`; the wrapper reshapes
+    internally. `decode(z)` returns 32x32x3 images in [0, 1].
+    """
+
+    z_dim: int = NOISE_DIM
+    image_size: int = IMAGE_SIZE
+    num_channels: int = NUM_CHANNELS
+
+    def __init__(
+        self,
+        pretrained: str = "ct-lpips",
+        checkpoint_path: str | Path | None = None,
+        seed: int = 0,
+    ) -> None:
+        if pretrained != "ct-lpips":
+            raise ValueError(
+                f"only 'ct-lpips' is wired up; got {pretrained!r}. "
+                "Add another config in `_default_config` to support others."
+            )
+        if checkpoint_path is None:
+            raise ValueError("checkpoint_path is required (path to the slim .pkl bundle)")
+
+        self.pretrained = pretrained
+        self.sigma_max = SIGMA_MAX
+        self.sigma_min = SIGMA_MIN
+        self.data_std = DATA_STD
+
+        self.config = _default_config()
+        self.model, init_params = _build_model_and_shapes(self.config, seed)
+
+        import pickle
+
+        from flax import serialization
+
+        with open(checkpoint_path, "rb") as f:
+            bundle = pickle.load(f)
+        self.params = serialization.from_state_dict(init_params, bundle["params"])
+
+        def _decode(z_flat: Array) -> Array:
+            batch = z_flat.shape[0]
+            # Reshape flat noise vector to image; scale by sigma_max (prior std).
+            x = z_flat.reshape(batch, IMAGE_SIZE, IMAGE_SIZE, NUM_CHANNELS) * SIGMA_MAX
+            img = _distill_one_step(self.model, self.params, x, SIGMA_MAX)
+            # Consistency models output pixels in ~[-1, 1]; map to [0, 1].
+            return jnp.clip((img + 1.0) * 0.5, 0.0, 1.0)
+
+        self._decode = jax.jit(_decode)
+
+    def decode(self, z: Array) -> Array:
+        return self._decode(z)
+
+    def sample_prior(self, key: Array, n: int) -> Array:
+        z = jax.random.normal(key, (n, self.z_dim))
+        return self.decode(z)
+
+
+def _distill_one_step(model, params, x: Array, sigma: float) -> Array:
+    """Inline KVE-SDE consistency-model one-step distillation (Karras et al. 2022).
+
+    x: (B, H, W, C) noise-scaled inputs (i.e. x ~ N(0, sigma**2 * I) reshaped).
+    sigma: scalar noise level (== SIGMA_MAX for one-step from prior).
+    """
+    t = jnp.asarray(sigma, dtype=x.dtype)
+    pred_t = SIGMA_MIN
+    d2 = DATA_STD * DATA_STD
+    in_scale = 1.0 / jnp.sqrt(t * t + d2)
+    cond_t = 0.25 * jnp.log(t)
+    cond_batch = jnp.full((x.shape[0],), cond_t, dtype=x.dtype)
+    raw = model.apply({"params": params}, x * in_scale, cond_batch, train=False)
+    out_scale = (t - pred_t) * DATA_STD / jnp.sqrt(t * t + d2)
+    skip_scale = d2 / ((t - pred_t) ** 2 + d2)
+    return skip_scale * x + out_scale * raw

cifar10_cd_ct_lpips/load.py

@@ -0,0 +1,3 @@
+from cm import CD
+
+model = CD(pretrained="ct-lpips", checkpoint_path="cifar10_cd_ct_lpips.pkl")