(self, checkpoint_path)
| 13 | |
| 14 | class WanTextEncoder(torch.nn.Module): |
| 15 | def __init__(self, checkpoint_path) -> None: |
| 16 | super().__init__() |
| 17 | |
| 18 | self.text_encoder = umt5_xxl( |
| 19 | encoder_only=True, |
| 20 | return_tokenizer=False, |
| 21 | dtype=torch.float32, |
| 22 | device=torch.device('cpu') |
| 23 | ).eval().requires_grad_(False) |
| 24 | |
| 25 | self.text_encoder.load_state_dict( |
| 26 | torch.load(f"{checkpoint_path}/Wan2.1-T2V-1.3B/models_t5_umt5-xxl-enc-bf16.pth", |
| 27 | map_location='cpu', weights_only=False) |
| 28 | ) |
| 29 | |
| 30 | self.tokenizer = HuggingfaceTokenizer( |
| 31 | name=f"{checkpoint_path}/Wan2.1-T2V-1.3B/google/umt5-xxl/", seq_len=512, clean='whitespace') |
| 32 | |
| 33 | @property |
| 34 | def device(self): |
nothing calls this directly
no test coverage detected