| """Make stickers with a trained Bhaav-E checkpoint. (OUR code.) |
| |
| Generation = walk back along the straight line from noise (t = 1) to clean (t = 0): |
| repeat 50 times: x <- x - v * (t_now - t_next) |
| Two standard extras: |
| - guidance (CFG): ask the model twice, with the caption and with the empty caption, then push further in the |
| direction the caption adds: v = v_empty + w * (v_caption - v_empty). Bigger w = follows the words harder. |
| - time shift: spend more of the 50 steps at high noise, where the overall layout is decided: |
| t' = s * t / (1 + (s - 1) * t). |
| |
| Usage (pod): |
| python sample.py --ckpt /workspace/runs/r1/last.pt --cache /workspace/cache \ |
| --prompt "an angry emoji shaking, steam puffing from its head, black background" --out /workspace/samples |
| Usage (released weights from Hugging Face: bhaav_e.safetensors + config.json + text/info.json + latents/vae_config.json): |
| python sample.py --ckpt <download>/bhaav_e.safetensors --cache <download> --prompt "..." --out samples |
| """ |
| import argparse, json |
| from pathlib import Path |
|
|
| import numpy as np |
| import torch |
| from PIL import Image |
|
|
| from videodit import SIZES, VideoDiT |
|
|
|
|
| def timesteps(n, shift): |
| t = torch.linspace(1.0, 0.0, n + 1) |
| return shift * t / (1 + (shift - 1) * t) |
|
|
|
|
| @torch.no_grad() |
| def sample(model, txt, mask, empty_txt, empty_mask, steps=50, cfg=3.0, shift=5.0, first=None, seed=0, |
| guide="cfg", bad_model=None, zero_steps=2): |
| """txt/mask: [B, L, 768] caption vectors. first: optional clean latent frame [B, 48, 16, 16] (image -> sticker). |
| guide picks how the two answers are combined (cfg = the guidance weight w in every mode): |
| cfg v = v_empty + w * (v_caption - v_empty) the standard recipe |
| cfgzero CFG-Zero* (Fan et al. 2025): first rescale v_empty by s = <v_caption, v_empty> / |v_empty|^2, the |
| amount of v_empty that actually points along v_caption, then guide from there; and take no step at all |
| for the first `zero_steps` steps, where an undertrained model's velocity is least reliable. |
| auto autoguidance (Karras et al. 2024): the "worse" answer comes from a weaker copy of the model (bad_model, |
| e.g. an early snapshot) given the SAME caption, not from the empty caption. Pushing away from the weak |
| copy's mistakes, instead of away from "no caption", sharpens quality without exaggerating the caption.""" |
| c = model.cfg |
| B, dev = txt.shape[0], txt.device |
| x = torch.randn(B, c.in_ch, c.frames, c.size, c.size, generator=torch.Generator().manual_seed(seed)).to(dev) |
| given = torch.full((B,), first is not None, device=dev) |
| ts = timesteps(steps, shift).to(dev) |
| for i in range(steps): |
| if first is not None: |
| x[:, :, 0] = first |
| t = ts[i].expand(B) |
| v_c = model(x, t, txt, mask, given) |
| if guide == "auto": |
| v_u = bad_model(x, t, txt, mask, given) |
| else: |
| v_u = model(x, t, empty_txt.expand(B, -1, -1), empty_mask.expand(B, -1), given) |
| if guide == "cfgzero": |
| if i < zero_steps: |
| continue |
| vc, vu = v_c.float().flatten(1), v_u.float().flatten(1) |
| s = ((vc * vu).sum(1) / (vu * vu).sum(1).clamp(min=1e-8)).view(B, 1, 1, 1, 1).to(v_u.dtype) |
| v_u = v_u * s |
| x = x - (v_u + cfg * (v_c - v_u)) * (ts[i] - ts[i + 1]) |
| if first is not None: |
| x[:, :, 0] = first |
| return x |
|
|
|
|
| def main(): |
| ap = argparse.ArgumentParser() |
| ap.add_argument("--ckpt", required=True) |
| ap.add_argument("--cache", required=True) |
| ap.add_argument("--prompt", action="append", required=True) |
| ap.add_argument("--out", required=True) |
| ap.add_argument("--steps", type=int, default=50) |
| ap.add_argument("--cfg", type=float, default=3.0) |
| ap.add_argument("--shift", type=float, default=5.0) |
| ap.add_argument("--seed", type=int, default=0) |
| a = ap.parse_args() |
| dev = "cuda" |
| from diffusers import AutoencoderKLWan |
| from transformers import AutoTokenizer, T5EncoderModel |
|
|
| if a.ckpt.endswith(".safetensors"): |
| from safetensors.torch import load_file |
| size = json.loads((Path(a.ckpt).parent / "config.json").read_text())["size"] |
| ema = load_file(a.ckpt, device=dev) |
| else: |
| s = torch.load(a.ckpt, map_location=dev) |
| size, ema = s["size"], s["ema"] |
| model = VideoDiT(SIZES[size]).to(dev).eval() |
| model.load_state_dict({k: v.float() for k, v in ema.items()}) |
|
|
| info = json.loads((Path(a.cache) / "text" / "info.json").read_text()) |
| tok = AutoTokenizer.from_pretrained(info["model"]) |
| t5 = T5EncoderModel.from_pretrained(info["model"], torch_dtype=torch.bfloat16).to(dev).eval() |
| def encode(texts): |
| t = tok(texts, padding="max_length", truncation=True, max_length=info["max_len"], return_tensors="pt").to(dev) |
| return (t5(input_ids=t.input_ids, attention_mask=t.attention_mask).last_hidden_state |
| * t.attention_mask[..., None]).float(), t.attention_mask |
| txt, mask = encode(a.prompt) |
| etxt, emask = encode([""]) |
|
|
| vae_cfg = json.loads((Path(a.cache) / "latents" / "vae_config.json").read_text()) |
| mean = torch.tensor(vae_cfg["latents_mean"], device=dev).view(1, -1, 1, 1, 1) |
| std = torch.tensor(vae_cfg["latents_std"], device=dev).view(1, -1, 1, 1, 1) |
| with torch.autocast("cuda", dtype=torch.bfloat16): |
| x = sample(model, txt, mask, etxt, emask, a.steps, a.cfg, a.shift, seed=a.seed) |
| vae = AutoencoderKLWan.from_pretrained("Wan-AI/Wan2.2-TI2V-5B-Diffusers", subfolder="vae", |
| torch_dtype=torch.float32).to(dev).eval() |
| with torch.no_grad(): |
| video = torch.cat([vae.decode(z[None]).sample for z in (x.float() * std + mean)]) |
| frames = ((video.clamp(-1, 1) + 1) * 127.5).round().byte().permute(0, 2, 3, 4, 1).cpu().numpy() |
| out = Path(a.out); out.mkdir(parents=True, exist_ok=True) |
| for k, (p, clip) in enumerate(zip(a.prompt, frames)): |
| imgs = [Image.fromarray(f) for f in clip[:-1]] |
| path = out / f"{k:02d}_seed{a.seed}.gif" |
| imgs[0].save(path, save_all=True, append_images=imgs[1:], duration=125, loop=0) |
| print(path, "|", p) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
