File size: 5,300 Bytes
5723b4b | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 | #!/usr/bin/env python3
"""
Convert Sana DC-AE VAE from PyTorch/Diffusers to MLX format
This avoids Core ML conversion issues by using MLX, Apple's optimized
framework for Apple Silicon.
Usage:
python convert_vae_to_mlx.py \
--model-version Efficient-Large-Model/Sana_600M_512px_diffusers \
--output sana_vae_mlx.npz
"""
import argparse
import json
import numpy as np
from pathlib import Path
import torch
def get_arguments():
parser = argparse.ArgumentParser()
parser.add_argument(
"--model-version",
default="Efficient-Large-Model/Sana_600M_512px_diffusers",
help="Sana model from Hugging Face",
)
parser.add_argument(
"--output",
required=True,
help="Output .npz file for MLX weights",
)
parser.add_argument(
"--component",
choices=["decoder", "encoder", "both"],
default="decoder",
help="Which component to convert",
)
return parser.parse_args()
def convert_pytorch_to_mlx(pytorch_weights, prefix="decoder."):
"""
Convert PyTorch weights to MLX format
MLX uses channels-last format (BHWC) while PyTorch uses (BCHW)
"""
mlx_weights = {}
for key, value in pytorch_weights.items():
if not key.startswith(prefix):
continue
# Remove prefix
mlx_key = key[len(prefix):]
# Convert tensor to numpy
if isinstance(value, torch.Tensor):
value = value.cpu().numpy()
# Convert Conv2d weights from (out_c, in_c, h, w) to (out_c, h, w, in_c)
if "conv" in mlx_key and value.ndim == 4:
value = np.transpose(value, (0, 2, 3, 1))
# Convert Linear weights - MLX uses (out, in) same as PyTorch
# So no conversion needed for linear layers
mlx_weights[mlx_key] = value.astype(np.float32)
return mlx_weights
def main():
args = get_arguments()
print("=" * 80)
print("Converting Sana VAE to MLX Format")
print("=" * 80)
print()
# Download model
print(f"Downloading model: {args.model_version}")
from huggingface_hub import snapshot_download
local_path = snapshot_download(
repo_id=args.model_version,
allow_patterns=["vae/*"],
)
print(f"✓ Downloaded to: {local_path}")
print()
# Load config
config_path = Path(local_path) / "vae" / "config.json"
with open(config_path) as f:
config = json.load(f)
print("VAE Configuration:")
print(f" Latent channels: {config.get('latent_channels', 32)}")
print(f" Scaling factor: {config.get('scaling_factor', 1.0)}")
print(f" Block channels: {config.get('block_out_channels')}")
print()
# Load PyTorch weights
weights_path = Path(local_path) / "vae" / "diffusion_pytorch_model.safetensors"
if not weights_path.exists():
weights_path = Path(local_path) / "vae" / "diffusion_pytorch_model.bin"
print(f"Loading weights from: {weights_path.name}")
if weights_path.suffix == ".safetensors":
from safetensors.torch import load_file
pytorch_weights = load_file(str(weights_path))
else:
pytorch_weights = torch.load(weights_path, map_location="cpu")
print(f"✓ Loaded {len(pytorch_weights)} weight tensors")
print()
# Convert weights
output_weights = {}
if args.component in ["decoder", "both"]:
print("Converting decoder weights...")
decoder_weights = convert_pytorch_to_mlx(pytorch_weights, "decoder.")
print(f" ✓ Converted {len(decoder_weights)} decoder weights")
output_weights.update({f"decoder.{k}": v for k, v in decoder_weights.items()})
if args.component in ["encoder", "both"]:
print("Converting encoder weights...")
encoder_weights = convert_pytorch_to_mlx(pytorch_weights, "encoder.")
print(f" ✓ Converted {len(encoder_weights)} encoder weights")
output_weights.update({f"encoder.{k}": v for k, v in encoder_weights.items()})
print()
# Add config to weights
output_weights["config"] = json.dumps(config)
# Save MLX weights
print(f"Saving MLX weights to: {args.output}")
np.savez(args.output, **output_weights)
print("✓ Conversion complete!")
print()
print("=" * 80)
print("Usage Example:")
print("=" * 80)
print()
print("import mlx.core as mx")
print("from sana_vae_mlx import DCAEDecoder")
print()
print(f'weights = np.load("{args.output}")')
print("decoder = DCAEDecoder(...)")
print("decoder.load_weights([(k, mx.array(v)) for k, v in weights.items()])")
print()
print("# Or use the built-in loader:")
print(f'decoder = DCAEDecoder.from_pretrained("{args.model_version}")')
print()
print("# Decode latents")
print("latents = mx.random.normal((1, 32, 16, 16)) # [B, C, H, W]")
print("image = decoder.decode(latents) # [B, 512, 512, 3]")
print()
# Print size info
total_size = sum(v.nbytes for v in output_weights.values() if isinstance(v, np.ndarray))
print(f"Total size: {total_size / 1024 / 1024:.1f} MB")
if __name__ == "__main__":
main()
|