Quantize model weights in-place Args: weights: Model state dictionary w_bit: Quantization bit width target_keys: List of module names to quantize Returns: Modified state dictionary with quantized weights and scales
(
weights,
w_bit=8,
target_keys=["attn", "ffn"],
adapter_keys=None,
key_idx=2,
ignore_key=None,
ignore_quant_keys=None,
linear_type="int8",
non_linear_dtype=torch.float,
comfyui_mode=False,
comfyui_keys=[],
)
| 312 | raise ValueError(f"Invalid direction: {direction}") |
| 313 | elif model_type == "h3": |
| 314 | # MiniMax-H3 checkpoints under the Diffusers ``transformer`` or |
| 315 | # ``transformer_ref`` directory already use LightX2V's runtime keys. |
| 316 | return [] |
| 317 | else: |
| 318 | raise ValueError(f"Unsupported model type: {model_type}") |
| 319 | |
| 320 | |
| 321 | def quantize_model( |
| 322 | weights, |
| 323 | w_bit=8, |
| 324 | target_keys=["attn", "ffn"], |
| 325 | adapter_keys=None, |
| 326 | key_idx=2, |
| 327 | ignore_key=None, |
| 328 | ignore_quant_keys=None, |
| 329 | linear_type="int8", |
| 330 | non_linear_dtype=torch.float, |
| 331 | preserve_non_quant_dtype=False, |
| 332 | comfyui_mode=False, |
| 333 | comfyui_keys=[], |
| 334 | ): |
| 335 | """ |
| 336 | Quantize model weights in-place |
| 337 | |
| 338 | Args: |
| 339 | weights: Model state dictionary |
| 340 | w_bit: Quantization bit width |
| 341 | target_keys: List of module names to quantize |
| 342 | |
| 343 | Returns: |
| 344 | Modified state dictionary with quantized weights and scales |
| 345 | """ |
| 346 | total_quantized = 0 |
| 347 | original_size = 0 |
| 348 | quantized_size = 0 |
| 349 | non_quantized_size = 0 |
| 350 | keys = list(weights.keys()) |
| 351 | |
| 352 | with tqdm(keys, desc="Quantizing weights") as pbar: |
| 353 | for key in pbar: |
| 354 | pbar.set_postfix(current_key=key, refresh=False) |
| 355 | |
| 356 | if ignore_key is not None and any(ig_key in key for ig_key in ignore_key): |
| 357 | del weights[key] |
| 358 | continue |
| 359 | |
| 360 | tensor = weights[key] |
| 361 | |
| 362 | # Skip non-tensors and non-2D tensors |
| 363 | if not isinstance(tensor, torch.Tensor) or tensor.dim() != 2: |
| 364 | if not preserve_non_quant_dtype and tensor.dtype != non_linear_dtype: |
| 365 | weights[key] = tensor.to(non_linear_dtype) |
| 366 | non_quantized_size += weights[key].numel() * weights[key].element_size() |
| 367 | else: |
| 368 | non_quantized_size += tensor.numel() * tensor.element_size() |
| 369 | continue |
| 370 | |
| 371 | # Check if key matches target modules |
no test coverage detected