File size: 5,038 Bytes
5c049df | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 | """d1_lune_sampler.py — SD15-Lune rectified-flow sampler + 225-token CLIP encoder.
PORTED VERBATIM (read-only law) from E:\\coworkers\\qwen_runner_1\\colab\\
sd15_lune_before_after.py (the sampler of record — repos/geolip-aleph-diffusion.md):
encode_clip_225 — 3 x 75-token chunked CLIP -> [B, 227, 768]. ALL judging in
this line conditions through THIS, never a bare 77-token encode.
flow_sample — Euler ODE sigma 1->0 on the SHIFT-warped grid
sigma = (shift*u)/(1+(shift-1)*u), v = noise - x0, batched CFG.
NOTE: pod/v35_expP1_sd15.py used an UNSHIFTED grid — that
discrepancy is ledgered in R0a; THIS grid matches the trainer.
Defaults from the source bed: STEPS 30, SHIFT 2.5, VAE_SCALE 0.18215, 64x64
latents (512px). Guidance differs across Phil's beds (6.0 before/after bed,
4.0 expP1) — beds pass it explicitly and ledger it.
CPU cert: python pod2/d1_lune_sampler.py (sigma-grid math + dummy-UNet
integration shape check only; real sampling is pod work).
"""
from __future__ import annotations
import torch
VAE_SCALE = 0.18215
SHIFT = 2.5
STEPS = 30
def encode_clip_225(prompts, tokenizer, text_encoder, device):
"""Encode text as 3 x 75-token chunks -> one 227-length hidden sequence."""
chunk_len = tokenizer.model_max_length # 77
body_len = chunk_len - 2 # 75
n_chunks = 3
ids = tokenizer(prompts, padding="max_length",
max_length=body_len * n_chunks + 2,
truncation=True, return_tensors="pt").input_ids # [B, 227]
bos, eos = ids[:, :1], ids[:, -1:]
chunks = []
for k in range(n_chunks):
s = 1 + k * body_len
chunks.append(torch.cat([bos, ids[:, s:s + body_len], eos], dim=1))
ids = torch.stack(chunks, dim=1).reshape(-1, chunk_len).to(device)
hs = text_encoder(ids)[0] # [B*3, 77, 768]
hs = hs.reshape(len(prompts), n_chunks * chunk_len, -1)
out = [hs[:, :1]]
for k in range(n_chunks):
s = k * chunk_len + 1
out.append(hs[:, s:s + body_len])
out.append(hs[:, -1:])
return torch.cat(out, dim=1) # [B, 227, 768]
def shifted_sigmas(n_steps: int, shift: float = SHIFT,
device="cpu") -> torch.Tensor:
u = torch.linspace(1.0, 0.0, n_steps + 1, device=device)
return (shift * u) / (1 + (shift - 1) * u)
@torch.no_grad()
def flow_sample(unet, ehs_cond, *, n_steps=STEPS, guidance=6.0, shift=SHIFT,
seed=1234, device="cuda", latent_hw=64,
encoder_attention_mask=None):
"""Integrate the rectified-flow ODE sigma 1->0; unet predicts v = noise - x0.
encoder_attention_mask (optional, [2B, L]): the Tier-A masked-append toggle
path — passed through to the UNet when provided."""
B = ehs_cond.shape[0]
g = torch.Generator(device=device).manual_seed(seed)
x = torch.randn(B, 4, latent_hw, latent_hw, generator=g, device=device,
dtype=torch.float32).to(ehs_cond.dtype) # noise drawn fp32
# for seed comparability across dtypes, then cast to the UNet's dtype
ehs_in = torch.cat([ehs_cond, torch.zeros_like(ehs_cond)], dim=0)
sigmas = shifted_sigmas(n_steps, shift, device)
kw = {}
if encoder_attention_mask is not None:
kw["encoder_attention_mask"] = encoder_attention_mask
for i in range(n_steps):
s, s_next = float(sigmas[i]), float(sigmas[i + 1])
t = torch.full((2 * B,), s * 1000.0, device=device)
v = unet(torch.cat([x, x], dim=0), t, ehs_in, return_dict=False,
**kw)[0]
v_cond, v_uncond = v.chunk(2)
v = v_uncond + guidance * (v_cond - v_uncond)
x = x + (s_next - s) * v # Euler (sigma down)
return x
@torch.no_grad()
def decode(vae, latents):
img = vae.decode(latents / VAE_SCALE).sample
img = (img / 2 + 0.5).clamp(0, 1)
return img.permute(0, 2, 3, 1).float().cpu().numpy()
def _smoke():
s = shifted_sigmas(30)
assert abs(float(s[0]) - 1.0) < 1e-6 and abs(float(s[-1])) < 1e-6, \
"shifted grid must span 1 -> 0 exactly"
assert (s[:-1] > s[1:]).all(), "sigmas must be strictly decreasing"
mid_shifted = float(shifted_sigmas(2)[1]) # u=0.5 -> 2.5*.5/1.75
assert abs(mid_shifted - (1.25 / 1.75)) < 1e-6, "SHIFT=2.5 warp wrong"
class DummyUNet:
def __call__(self, x, t, ehs, return_dict=False, **kw):
return (torch.zeros_like(x),)
out = flow_sample(DummyUNet(), torch.zeros(2, 227, 768), n_steps=4,
guidance=1.0, seed=0, device="cpu")
assert out.shape == (2, 4, 64, 64)
print("d1_lune_sampler smoke PASSED (shifted grid endpoints + monotone, "
"warp value, dummy integration shape)")
if __name__ == "__main__":
_smoke()
|