pearsonkyle/SDXL-Model-Merger

1"""Model export functionality for SDXL Model Merger."""2 3import torch4from safetensors.torch import save_file5 6from . import config7from .config import SCRIPT_DIR8from .gpu_decorator import GPU9 10 11def _quantize_model(model, qtype: str):12    """Apply torchao quantization to a model using quantize_."""13    from torchao.quantization import quantize_14 15    if qtype == "int8":16        from torchao.quantization import Int8WeightOnlyConfig17 18        print("  ⚙️ Quantizing with int8_weight_only...")19        config = Int8WeightOnlyConfig()20        quantize_(model, config)21 22    elif qtype == "int4":23        from torchao.quantization import Int4WeightOnlyConfig24 25        print("  ⚙️ Quantizing with int4_weight_only (group_size=32)...")26        config = Int4WeightOnlyConfig(group_size=32)27        quantize_(model, config)28 29    elif qtype == "float8":30        from torchao.quantization import Float8DynamicActivationFloat8WeightConfig31 32        print("  ⚙️ Quantizing with float8_dynamic_activation_float8_weight...")33        config = Float8DynamicActivationFloat8WeightConfig()34        quantize_(model, config)35 36    else:37        raise ValueError(f"Unsupported qtype: {qtype}. Must be one of: int8, int4, float8")38 39 40@GPU(duration=180)41def _extract_and_save(pipe, include_lora, quantize, qtype, save_format):42    """GPU-decorated helper that extracts weights and saves the model."""43    if include_lora:44        try:45            pipe.unload_lora_weights()46        except Exception as e:47            print(f"  ℹ️ Could not unload LoRAs: {e}")48 49    # Quantize components in-place before extracting state dicts50    if quantize and qtype != "none":51        _quantize_model(pipe.unet, qtype)52        # torchao quantized tensors cannot be saved with safetensors, use torch.save instead53        # Don't dequantize - keep the quantized format for smaller file size54 55    merged_state_dict = {}56 57    # Extract UNet weights58    for k, v in pipe.unet.state_dict().items():59        # For quantized tensors, save directly; otherwise convert to half60        if hasattr(v, 'dequantize'):61            # Keep quantized tensor as-is for smaller file size62            merged_state_dict[f"unet.{k}"] = v63        else:64            merged_state_dict[f"unet.{k}"] = v.contiguous().half()65 66    # Extract text encoder weights67    if pipe.text_encoder is not None:68        for k, v in pipe.text_encoder.state_dict().items():69            if hasattr(v, 'dequantize'):70                merged_state_dict[f"text_encoder.{k}"] = v71            else:72                merged_state_dict[f"text_encoder.{k}"] = v.contiguous().half()73    if pipe.text_encoder_2 is not None:74        for k, v in pipe.text_encoder_2.state_dict().items():75            if hasattr(v, 'dequantize'):76                merged_state_dict[f"text_encoder_2.{k}"] = v77            else:78                merged_state_dict[f"text_encoder_2.{k}"] = v.contiguous().half()79 80    # Extract VAE weights81    if pipe.vae is not None:82        for k, v in pipe.vae.state_dict().items():83            if hasattr(v, 'dequantize'):84                merged_state_dict[f"first_stage_model.{k}"] = v85            else:86                merged_state_dict[f"first_stage_model.{k}"] = v.contiguous().half()87 88    # Save model89    ext = ".bin" if save_format == "bin" else ".safetensors"90    prefix = f"{qtype}_" if quantize and qtype != "none" else ""91    out_path = SCRIPT_DIR / f"merged_{prefix}checkpoint{ext}"92 93    if quantize and qtype != "none":94        # torchao quantized tensors are not compatible with safetensors95        # Use torch.save instead which preserves the quantization format96        ext = ".pt"97        out_path = SCRIPT_DIR / f"merged_{qtype}_checkpoint.pt"98        torch.save(merged_state_dict, str(out_path))99    elif ext == ".bin":100        torch.save(merged_state_dict, str(out_path))101    else:102        save_file(merged_state_dict, str(out_path))103 104    return out_path105 106 107def export_merged_model(108    include_lora: bool,109    quantize: bool,110    qtype: str,111    save_format: str = "safetensors",112):113    """114    Export the merged pipeline model with optional LoRA baking and quantization.115 116    Args:117        include_lora: Whether to include fused LoRAs in export118        quantize: Whether to apply quantization119        qtype: Quantization type - 'none', 'int8', 'int4', or 'float8'120        save_format: Output format - 'safetensors' or 'bin'121 122    Returns:123        Tuple of (output_path or None, status message)124    """125    # Fetch the pipeline at call time — avoids the stale import-by-value problem.126    pipe = config.get_pipe()127 128    if not pipe:129        return None, "⚠️ Please load a pipeline first."130 131    try:132        # Validate quantization type133        valid_qtypes = ("none", "int8", "int4", "float8")134        if qtype not in valid_qtypes:135            return None, f"❌ Invalid quantization type: {qtype}. Must be one of: {valid_qtypes}"136 137        out_path = _extract_and_save(pipe, include_lora, quantize, qtype, save_format)138 139        size_gb = out_path.stat().st_size / 1024**3140 141        if quantize and qtype != "none":142            msg = f"✅ Quantized checkpoint saved: `{out_path}` ({size_gb:.2f} GB)"143        else:144            msg = f"✅ Merged checkpoint saved: `{out_path}` ({size_gb:.2f} GB)"145 146        return str(out_path), msg147 148    except ImportError as e:149        return None, f"❌ Missing dependency: {str(e)}"150    except Exception as e:151        import traceback152        print(traceback.format_exc())153        return None, f"❌ Export failed: {str(e)}"154 155 156def get_export_status() -> str:157    """Get current export capability status."""158    try:159        from torchao.quantization import quantize_, Int4WeightOnlyConfig, Int8WeightOnlyConfig, Float8DynamicActivationFloat8WeightConfig160        return "✅ torchao available for quantization"161    except ImportError:162        return "ℹ️ Install torchao for quantization support: pip install torchao"163