(config: Dict[str, Any])
| 164 | |
| 165 | |
| 166 | def load_wan_text_encoder(config: Dict[str, Any]): |
| 167 | from lightx2v.models.input_encoders.hf.wan.t5.model import T5EncoderModel |
| 168 | |
| 169 | # offload config |
| 170 | t5_offload = config.get("t5_cpu_offload", config.get("cpu_offload")) |
| 171 | if t5_offload: |
| 172 | t5_device = torch.device("cpu") |
| 173 | else: |
| 174 | t5_device = torch.device(AI_DEVICE) |
| 175 | tokenizer_path = os.path.join(config["model_path"], "google/umt5-xxl") |
| 176 | # quant_config |
| 177 | t5_quantized = config.get("t5_quantized", False) |
| 178 | if t5_quantized: |
| 179 | t5_quant_scheme = config.get("t5_quant_scheme", None) |
| 180 | assert t5_quant_scheme is not None |
| 181 | tmp_t5_quant_scheme = t5_quant_scheme.split("-")[0] |
| 182 | t5_model_name = f"models_t5_umt5-xxl-enc-{tmp_t5_quant_scheme}.pth" |
| 183 | t5_quantized_ckpt = find_torch_model_path(config, "t5_quantized_ckpt", t5_model_name) |
| 184 | t5_original_ckpt = None |
| 185 | else: |
| 186 | t5_quant_scheme = None |
| 187 | t5_quantized_ckpt = None |
| 188 | t5_model_name = "models_t5_umt5-xxl-enc-bf16.pth" |
| 189 | t5_original_ckpt = find_torch_model_path(config, "t5_original_ckpt", t5_model_name) |
| 190 | |
| 191 | text_encoder = T5EncoderModel( |
| 192 | text_len=config["text_len"], |
| 193 | dtype=torch.bfloat16, |
| 194 | device=t5_device, |
| 195 | checkpoint_path=t5_original_ckpt, |
| 196 | tokenizer_path=tokenizer_path, |
| 197 | shard_fn=None, |
| 198 | cpu_offload=t5_offload, |
| 199 | t5_quantized=t5_quantized, |
| 200 | t5_quantized_ckpt=t5_quantized_ckpt, |
| 201 | quant_scheme=t5_quant_scheme, |
| 202 | load_from_rank0=config.get("load_from_rank0", False), |
| 203 | lazy_load=config.get("t5_lazy_load", False), |
| 204 | ) |
| 205 | # Return single encoder to match original returning list |
| 206 | text_encoders = [text_encoder] |
| 207 | return text_encoders |
| 208 | |
| 209 | |
| 210 | def load_wan_image_encoder(config: Dict[str, Any]): |
no test coverage detected