(self, model, diffusion_model_dtype)
| 526 | self.text_encoders.append(ModelWrapper(load_fn)) |
| 527 | |
| 528 | def dequantize(self, model, diffusion_model_dtype): |
| 529 | operations = comfy.ops.disable_weight_init |
| 530 | for mod_name, module in model.named_children(): |
| 531 | is_quantized = False |
| 532 | for p_name, p in module.named_parameters(recurse=False): |
| 533 | if p.__class__.__name__ == 'QuantizedTensor': |
| 534 | module.register_parameter(p_name, nn.Parameter(p.dequantize())) |
| 535 | p = getattr(module, p_name) |
| 536 | is_quantized = True |
| 537 | |
| 538 | name = f'{mod_name}.{p_name}' |
| 539 | if any(keyword in name for keyword in self.keep_in_high_precision) or p.ndim == 1: |
| 540 | continue |
| 541 | p.data = p.data.to(diffusion_model_dtype) |
| 542 | |
| 543 | if is_quantized: |
| 544 | bias = module.bias is not None |
| 545 | with accelerate.init_empty_weights(): |
| 546 | new_linear = operations.Linear(module.in_features, module.out_features, bias=bias) |
| 547 | new_linear.comfy_cast_weights = True # model_patcher.set_model_compute_dtype() would normally set this |
| 548 | new_linear.weight = module.weight |
| 549 | if bias: |
| 550 | new_linear.bias = module.bias |
| 551 | model._modules[mod_name] = new_linear |
| 552 | |
| 553 | if len(list(module.children())) > 0: |
| 554 | self.dequantize(module, diffusion_model_dtype) |
| 555 | |
| 556 | def load_diffusion_model(self): |
| 557 | dtype = self.model_config['dtype'] |
no test coverage detected