Spaces:
Running on Zero
Running on Zero
File size: 2,954 Bytes
36cdb93 | 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 | """Prompt conditioning for the demo: the 96 prompt bank, plus free text encoding.
Two sources of conditioning, deliberately kept distinct:
* **The bank** (`data/prompts.pt`) -- 96 umt5-xxl embeddings that every clip in the
project was generated and trained under. Zero cost, and exactly the conditioning
the published numbers refer to.
* **Free text** -- encoded here, on this machine, with umt5-xxl (11.4 GB, loaded
lazily on first use). Note that umt5 embeddings are mildly hardware-dependent, so
text encoded here is not numerically identical to what the same string would give
on the training box. It looks fine; it just isn't the *same* conditioning, so
free-text results are not strictly comparable to the bank's published metrics.
Both paths end in the same shape: [1, 512, 4096], zero padded, which is what
`WanModel` expects.
"""
from pathlib import Path
import torch
TEXT_LEN = 512
THEMES = [
("People", 0, 40),
("Animals", 40, 56),
("Nature", 56, 72),
("City", 72, 86),
("Objects", 86, 96),
]
def theme_of(idx):
for name, lo, hi in THEMES:
if lo <= idx < hi:
return name
return "Other"
class PromptBank:
def __init__(self, path):
d = torch.load(path, map_location="cpu", weights_only=False)
self.texts = d["prompts"]
self.pos = d["pos"] # [96, 512, 4096] fp16
self.neg = d["neg"]
self.neg_prompt = d["neg_prompt"]
def __len__(self):
return len(self.texts)
def embedding(self, idx):
return self.pos[idx : idx + 1].float()
def catalogue(self):
return [
{"idx": i, "text": t, "theme": theme_of(i)} for i, t in enumerate(self.texts)
]
class TextEncoder:
"""Lazy umt5-xxl. 11.4 GB -- only loaded if someone actually types a prompt."""
def __init__(self, checkpoint, tokenizer_path, wan_repo, device="cuda"):
self.checkpoint = str(checkpoint)
self.tokenizer_path = str(tokenizer_path)
self.wan_repo = str(wan_repo)
self.device = device
self._model = None
@property
def loaded(self):
return self._model is not None
def load(self):
if self._model is not None:
return
from wan.modules.t5 import T5EncoderModel
self._model = T5EncoderModel(
text_len=TEXT_LEN, dtype=torch.bfloat16, device=self.device,
checkpoint_path=self.checkpoint, tokenizer_path=self.tokenizer_path)
@torch.no_grad()
def encode(self, text):
"""-> [1, 512, 4096] float32, zero-padded exactly as the bank is."""
self.load()
ctx = self._model([text], self.device)[0] # [L, 4096], L <= 512
out = torch.zeros(TEXT_LEN, ctx.shape[1], dtype=torch.float32, device=ctx.device)
out[: ctx.shape[0]] = ctx.float()
return out.unsqueeze(0).cpu()
def unload(self):
self._model = None
torch.cuda.empty_cache()
|