(
self,
ckpt_name: str,
dtype: str,
attention_mode: str = "auto",
rms_norm_function: str = "default",
quantization: str = "disabled",
load_device_name: str = "main_device",
)
| 1606 | CATEGORY = "loaders/bitdance" |
| 1607 | |
| 1608 | def load_text_encoder( |
| 1609 | self, |
| 1610 | ckpt_name: str, |
| 1611 | dtype: str, |
| 1612 | attention_mode: str = "auto", |
| 1613 | rms_norm_function: str = "default", |
| 1614 | quantization: str = "disabled", |
| 1615 | load_device_name: str = "main_device", |
| 1616 | ): |
| 1617 | target_dtype = _resolve_dtype(dtype) |
| 1618 | selected_path = _get_full_path_from_folder_types(TEXT_ENCODER_FOLDER_TYPES, ckpt_name) |
| 1619 | model_root = _try_resolve_model_root_from_path(selected_path) |
| 1620 | use_fp8_runtime = str(quantization or "disabled") in { |
| 1621 | "fp8_e4m3fn_scaled", |
| 1622 | "fp8_e4m3fn", |
| 1623 | } |
| 1624 | if model_root is not None: |
| 1625 | if use_fp8_runtime: |
| 1626 | LOGGER.warning( |
| 1627 | "BitDanceQwenFP8 runtime is currently implemented for single-file merged text encoder checkpoints. " |
| 1628 | "Bundle/sharded text encoder will use standard HF runtime." |
| 1629 | ) |
| 1630 | tokenizer, llm_model, llm_config = local_model.build_text_model_and_tokenizer( |
| 1631 | model_root, |
| 1632 | target_dtype, |
| 1633 | attention_mode=attention_mode, |
| 1634 | rms_norm_function=rms_norm_function, |
| 1635 | ) |
| 1636 | if quantization != "disabled": |
| 1637 | LOGGER.info( |
| 1638 | "BitDance text encoder quantization='%s' selected. " |
| 1639 | "Bundle loader still dequantizes FP8/scale weights to %s during load.", |
| 1640 | quantization, |
| 1641 | str(target_dtype).replace("torch.", ""), |
| 1642 | ) |
| 1643 | _load_llm_shards(llm_model, model_root, target_dtype) |
| 1644 | # Place text encoder after loading (bundle mode still loads shards via CPU dict path). |
| 1645 | try: |
| 1646 | llm_model = llm_model.to(_named_device(load_device_name)).eval() |
| 1647 | except Exception as e: |
| 1648 | LOGGER.warning("Failed to place BitDance text encoder on %s: %s", load_device_name, e) |
| 1649 | |
| 1650 | hidden_size = int(getattr(llm_config, "hidden_size", local_model.load_hidden_size(model_root))) # type: ignore |
| 1651 | runtime = BitDanceTextRuntime( |
| 1652 | root=model_root, |
| 1653 | tokenizer=tokenizer, |
| 1654 | llm_model=llm_model.eval(), |
| 1655 | hidden_size=hidden_size, |
| 1656 | ) |
| 1657 | else: |
| 1658 | if use_fp8_runtime: |
| 1659 | runtime = _build_text_runtime_from_single_file_fp8( |
| 1660 | selected_path, |
| 1661 | target_dtype=target_dtype, |
| 1662 | attention_mode=attention_mode, |
| 1663 | rms_norm_function=rms_norm_function, |
| 1664 | load_device_name=load_device_name, |
| 1665 | ) |
no test coverage detected