(
text_encoder_type,
text_encoder_precision=None,
text_encoder_path=None,
logger=None,
device=None,
)
| 70 | |
| 71 | |
| 72 | def load_text_encoder( |
| 73 | text_encoder_type, |
| 74 | text_encoder_precision=None, |
| 75 | text_encoder_path=None, |
| 76 | logger=None, |
| 77 | device=None, |
| 78 | ): |
| 79 | if text_encoder_path is None: |
| 80 | text_encoder_path = TEXT_ENCODER_PATH[text_encoder_type] |
| 81 | if logger is not None: |
| 82 | logger.info( |
| 83 | f"Loading text encoder model ({text_encoder_type}) from: {text_encoder_path}" |
| 84 | ) |
| 85 | |
| 86 | torch_dtype = 'auto' |
| 87 | if text_encoder_precision is not None: |
| 88 | torch_dtype = PRECISION_TO_TYPE[text_encoder_precision] |
| 89 | |
| 90 | if text_encoder_type == "clipL": |
| 91 | text_encoder = CLIPTextModel.from_pretrained(text_encoder_path, torch_dtype=torch_dtype) |
| 92 | text_encoder.final_layer_norm = text_encoder.text_model.final_layer_norm |
| 93 | elif text_encoder_type == "llm": |
| 94 | text_encoder = AutoModel.from_pretrained( |
| 95 | text_encoder_path, low_cpu_mem_usage=True, torch_dtype=torch_dtype |
| 96 | ) |
| 97 | text_encoder.final_layer_norm = text_encoder.norm |
| 98 | else: |
| 99 | raise ValueError(f"Unsupported text encoder type: {text_encoder_type}") |
| 100 | # from_pretrained will ensure that the model is in eval mode. |
| 101 | |
| 102 | text_encoder.requires_grad_(False) |
| 103 | |
| 104 | if logger is not None: |
| 105 | logger.info(f"Text encoder to dtype: {text_encoder.dtype}") |
| 106 | |
| 107 | if device is not None: |
| 108 | text_encoder = text_encoder.to(device) |
| 109 | |
| 110 | return text_encoder, text_encoder_path |
| 111 | |
| 112 | |
| 113 | def train_schedule_steps(self): |
nothing calls this directly
no test coverage detected