Spaces:
Running on Zero
Running on Zero
| """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 | |
| 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) | |
| 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() | |