(self, config)
| 183 | ] |
| 184 | |
| 185 | def __init__(self, config): |
| 186 | self.config = config |
| 187 | self.model_config = self.config['model'] |
| 188 | self.offloader = ModelOffloader('dummy', [], 0, 0, True, torch.device('cuda'), False, debug=False) |
| 189 | dtype = self.model_config['dtype'] |
| 190 | self.cache_text_embeddings = self.model_config.get('cache_text_embeddings', True) |
| 191 | self.multiscale_loss_weight = self.model_config.get('multiscale_loss_weight', None) |
| 192 | |
| 193 | # This isn't a nn.Module. |
| 194 | self.vae = WanVAE( |
| 195 | vae_pth=self.model_config['vae_path'], |
| 196 | device='cpu', |
| 197 | dtype=dtype, |
| 198 | ) |
| 199 | # These need to be on the device the VAE will be moved to during caching. |
| 200 | self.vae.mean = self.vae.mean.to('cuda') |
| 201 | self.vae.std = self.vae.std.to('cuda') |
| 202 | self.vae.scale = [self.vae.mean, 1.0 / self.vae.std] |
| 203 | |
| 204 | self.is_generic_llm = False |
| 205 | self.t5_tokenizer = T5TokenizerFast( |
| 206 | vocab_file='configs/t5_old/spiece.model', |
| 207 | tokenizer_file='configs/t5_old/tokenizer.json', |
| 208 | ) |
| 209 | |
| 210 | if 't5_path' in self.model_config: |
| 211 | self.tokenizer = self.t5_tokenizer |
| 212 | t5_state_dict = load_state_dict(self.model_config['t5_path']) |
| 213 | if self.model_config.get('text_encoder_nf4', False): |
| 214 | quantization_config = transformers.BitsAndBytesConfig( |
| 215 | load_in_4bit=True, |
| 216 | bnb_4bit_quant_type='nf4', |
| 217 | bnb_4bit_compute_dtype=dtype, |
| 218 | ) |
| 219 | else: |
| 220 | quantization_config = None |
| 221 | self.text_encoder = T5EncoderModel.from_pretrained( |
| 222 | None, |
| 223 | config='configs/t5_old/config.json', |
| 224 | state_dict=t5_state_dict, |
| 225 | torch_dtype='auto', |
| 226 | local_files_only=True, |
| 227 | quantization_config=quantization_config, |
| 228 | ) |
| 229 | if quantization_config is None and self.model_config.get('text_encoder_fp8', False): |
| 230 | for name, p in self.text_encoder.named_parameters(): |
| 231 | if p.ndim == 2 and not ('shared' in name or 'relative_attention_bias' in name): |
| 232 | p.data = p.data.to(torch.float8_e4m3fn) |
| 233 | elif 'llm_path' in self.model_config: |
| 234 | llm_path = self.model_config['llm_path'] |
| 235 | if os.path.isdir(llm_path): |
| 236 | # generic Transformers LLM |
| 237 | self.tokenizer = AutoTokenizer.from_pretrained(llm_path, local_files_only=True) |
| 238 | text_encoder = AutoModelForCausalLM.from_pretrained(llm_path, dtype=dtype, local_files_only=True) |
| 239 | else: |
| 240 | # assume Qwen3-0.6b (Anima) |
| 241 | self.tokenizer = AutoTokenizer.from_pretrained('configs/qwen3_06b', local_files_only=True) |
| 242 | llm_config = transformers.Qwen3Config.from_pretrained('configs/qwen3_06b', local_files_only=True) |
nothing calls this directly
no test coverage detected