Refactor LTX-2 model structure
This commit is contained in:
Prince Canuma
785
mlx_video/models/ltx_2/convert.py
Normal file
785
mlx_video/models/ltx_2/convert.py
Normal file
@@ -0,0 +1,785 @@
|
||||
"""Convert LTX-2/2.3 safetensors to MLX directory layout.
|
||||
|
||||
Converts from the single-file format (e.g. Lightricks/LTX-2/ltx-2-19b-distilled.safetensors
|
||||
or Lightricks/LTX-2.3/ltx-2.3-22b-distilled.safetensors) to the modular directory structure:
|
||||
|
||||
output/
|
||||
├── transformer/ # DiT transformer weights (sharded)
|
||||
│ ├── config.json
|
||||
│ ├── model-00001-of-N.safetensors
|
||||
│ └── model.safetensors.index.json
|
||||
├── vae/
|
||||
│ ├── decoder/ # Video VAE decoder
|
||||
│ │ ├── config.json
|
||||
│ │ └── model.safetensors
|
||||
│ └── encoder/ # Video VAE encoder
|
||||
│ ├── config.json
|
||||
│ └── model.safetensors
|
||||
├── audio_vae/ # Audio VAE decoder
|
||||
│ ├── config.json
|
||||
│ └── model.safetensors
|
||||
├── vocoder/ # Audio vocoder
|
||||
│ ├── config.json
|
||||
│ └── model.safetensors
|
||||
└── text_projections/ # Text projection connectors
|
||||
└── model.safetensors
|
||||
|
||||
Usage:
|
||||
# From HF repo ID
|
||||
python -m mlx_video.models.ltx_2.convert --source Lightricks/LTX-2 --output LTX-2-distilled --variant distilled
|
||||
python -m mlx_video.models.ltx_2.convert --source Lightricks/LTX-2.3 --output LTX-2.3-distilled --variant distilled
|
||||
|
||||
# From local folder containing the monolithic safetensors
|
||||
python -m mlx_video.models.ltx_2.convert --source ./Lightricks-LTX-2/ --output LTX-2-distilled --variant distilled
|
||||
|
||||
# From a direct safetensors file path
|
||||
python -m mlx_video.models.ltx_2.convert --source ./ltx-2-19b-distilled.safetensors --output LTX-2-distilled --variant distilled
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import json
|
||||
import re
|
||||
import shutil
|
||||
from pathlib import Path
|
||||
from typing import Dict
|
||||
|
||||
import mlx.core as mx
|
||||
|
||||
|
||||
# ─── Key prefix routing ──────────────────────────────────────────────────────
|
||||
|
||||
TRANSFORMER_PREFIX = "model.diffusion_model."
|
||||
VAE_DECODER_PREFIX = "vae.decoder."
|
||||
VAE_ENCODER_PREFIX = "vae.encoder."
|
||||
VAE_STATS_PREFIX = "vae.per_channel_statistics."
|
||||
AUDIO_DECODER_PREFIX = "audio_vae.decoder."
|
||||
AUDIO_ENCODER_PREFIX = "audio_vae.encoder."
|
||||
AUDIO_STATS_PREFIX = "audio_vae.per_channel_statistics."
|
||||
VOCODER_PREFIX = "vocoder."
|
||||
TEXT_PROJ_PREFIX = "text_embedding_projection."
|
||||
VIDEO_CONNECTOR_PREFIX = "model.diffusion_model.video_embeddings_connector."
|
||||
AUDIO_CONNECTOR_PREFIX = "model.diffusion_model.audio_embeddings_connector."
|
||||
|
||||
|
||||
# ─── Sanitization functions ──────────────────────────────────────────────────
|
||||
|
||||
|
||||
def sanitize_transformer(weights: Dict[str, mx.array]) -> Dict[str, mx.array]:
|
||||
"""Sanitize transformer keys: strip prefix, rename layers, cast to bfloat16."""
|
||||
sanitized = {}
|
||||
for key, value in weights.items():
|
||||
if not key.startswith(TRANSFORMER_PREFIX):
|
||||
continue
|
||||
# Skip connector weights (they go to text_projections)
|
||||
if "audio_embeddings_connector" in key or "video_embeddings_connector" in key:
|
||||
continue
|
||||
|
||||
new_key = key[len(TRANSFORMER_PREFIX):]
|
||||
new_key = new_key.replace(".to_out.0.", ".to_out.")
|
||||
new_key = new_key.replace(".ff.net.0.proj.", ".ff.proj_in.")
|
||||
new_key = new_key.replace(".ff.net.2.", ".ff.proj_out.")
|
||||
new_key = new_key.replace(".audio_ff.net.0.proj.", ".audio_ff.proj_in.")
|
||||
new_key = new_key.replace(".audio_ff.net.2.", ".audio_ff.proj_out.")
|
||||
new_key = new_key.replace(".linear_1.", ".linear1.")
|
||||
new_key = new_key.replace(".linear_2.", ".linear2.")
|
||||
|
||||
# Cast all weights to bfloat16 (matches MLX model loading behavior)
|
||||
if value.dtype != mx.bfloat16:
|
||||
value = value.astype(mx.bfloat16)
|
||||
|
||||
sanitized[new_key] = value
|
||||
return sanitized
|
||||
|
||||
|
||||
def sanitize_vae_decoder(weights: Dict[str, mx.array]) -> Dict[str, mx.array]:
|
||||
"""Sanitize VAE decoder keys: strip prefix, transpose Conv3d, wrap .conv."""
|
||||
sanitized = {}
|
||||
for key, value in weights.items():
|
||||
new_key = None
|
||||
|
||||
if key.startswith(VAE_STATS_PREFIX):
|
||||
if key == "vae.per_channel_statistics.mean-of-means":
|
||||
new_key = "per_channel_statistics.mean"
|
||||
elif key == "vae.per_channel_statistics.std-of-means":
|
||||
new_key = "per_channel_statistics.std"
|
||||
else:
|
||||
continue
|
||||
elif key.startswith(VAE_DECODER_PREFIX):
|
||||
new_key = key[len(VAE_DECODER_PREFIX):]
|
||||
else:
|
||||
continue
|
||||
|
||||
# Conv3d weight transpose: PyTorch (O, I, D, H, W) -> MLX (O, D, H, W, I)
|
||||
if ".conv.weight" in key and value.ndim == 5:
|
||||
value = mx.transpose(value, (0, 2, 3, 4, 1))
|
||||
|
||||
# Wrap .conv.weight -> .conv.conv.weight (CausalConv3d wrapper)
|
||||
if ".conv.weight" in new_key or ".conv.bias" in new_key:
|
||||
if ".conv.conv.weight" not in new_key and ".conv.conv.bias" not in new_key:
|
||||
new_key = new_key.replace(".conv.weight", ".conv.conv.weight")
|
||||
new_key = new_key.replace(".conv.bias", ".conv.conv.bias")
|
||||
|
||||
sanitized[new_key] = value
|
||||
return sanitized
|
||||
|
||||
|
||||
def sanitize_vae_encoder(weights: Dict[str, mx.array]) -> Dict[str, mx.array]:
|
||||
"""Sanitize VAE encoder keys: strip prefix, transpose Conv3d/Conv2d."""
|
||||
sanitized = {}
|
||||
for key, value in weights.items():
|
||||
new_key = None
|
||||
|
||||
if "position_ids" in key:
|
||||
continue
|
||||
|
||||
if key.startswith(VAE_STATS_PREFIX):
|
||||
if key == "vae.per_channel_statistics.mean-of-means":
|
||||
new_key = "per_channel_statistics.mean"
|
||||
elif key == "vae.per_channel_statistics.std-of-means":
|
||||
new_key = "per_channel_statistics.std"
|
||||
else:
|
||||
continue
|
||||
# Per-channel statistics must stay float32 for precision
|
||||
if value.dtype != mx.float32:
|
||||
value = value.astype(mx.float32)
|
||||
elif key.startswith(VAE_ENCODER_PREFIX):
|
||||
new_key = key[len(VAE_ENCODER_PREFIX):]
|
||||
else:
|
||||
continue
|
||||
|
||||
# Conv3d: PyTorch (O, I, D, H, W) -> MLX (O, D, H, W, I)
|
||||
if "conv" in new_key.lower() and "weight" in new_key and value.ndim == 5:
|
||||
value = mx.transpose(value, (0, 2, 3, 4, 1))
|
||||
|
||||
# Conv2d: PyTorch (O, I, H, W) -> MLX (O, H, W, I)
|
||||
if "conv" in new_key.lower() and "weight" in new_key and value.ndim == 4:
|
||||
value = mx.transpose(value, (0, 2, 3, 1))
|
||||
|
||||
sanitized[new_key] = value
|
||||
return sanitized
|
||||
|
||||
|
||||
def sanitize_audio_decoder(weights: Dict[str, mx.array]) -> Dict[str, mx.array]:
|
||||
"""Sanitize audio VAE decoder keys: strip prefix, transpose Conv2d."""
|
||||
sanitized = {}
|
||||
for key, value in weights.items():
|
||||
new_key = None
|
||||
|
||||
if key.startswith(AUDIO_DECODER_PREFIX):
|
||||
new_key = key[len(AUDIO_DECODER_PREFIX):]
|
||||
elif key.startswith(AUDIO_STATS_PREFIX):
|
||||
if "mean-of-means" in key:
|
||||
new_key = "per_channel_statistics.mean_of_means"
|
||||
elif "std-of-means" in key:
|
||||
new_key = "per_channel_statistics.std_of_means"
|
||||
else:
|
||||
continue
|
||||
else:
|
||||
continue
|
||||
|
||||
# Conv2d: PyTorch (O, I, H, W) -> MLX (O, H, W, I)
|
||||
if "conv" in new_key.lower() and "weight" in new_key and value.ndim == 4:
|
||||
value = mx.transpose(value, (0, 2, 3, 1))
|
||||
|
||||
sanitized[new_key] = value
|
||||
return sanitized
|
||||
|
||||
|
||||
def sanitize_vocoder(weights: Dict[str, mx.array]) -> Dict[str, mx.array]:
|
||||
"""Sanitize vocoder keys: strip prefix, transpose Conv1d/ConvTranspose1d."""
|
||||
sanitized = {}
|
||||
for key, value in weights.items():
|
||||
if not key.startswith(VOCODER_PREFIX):
|
||||
continue
|
||||
|
||||
new_key = key[len(VOCODER_PREFIX):]
|
||||
|
||||
# Handle Conv1d/ConvTranspose1d weight shape conversion
|
||||
if "weight" in new_key and value.ndim == 3:
|
||||
if "ups" in new_key:
|
||||
# ConvTranspose1d: PyTorch (in_ch, out_ch, kernel) -> MLX (out_ch, kernel, in_ch)
|
||||
value = mx.transpose(value, (1, 2, 0))
|
||||
else:
|
||||
# Conv1d: PyTorch (out_ch, in_ch, kernel) -> MLX (out_ch, kernel, in_ch)
|
||||
value = mx.transpose(value, (0, 2, 1))
|
||||
|
||||
sanitized[new_key] = value
|
||||
return sanitized
|
||||
|
||||
|
||||
def sanitize_connector_key(key: str) -> str:
|
||||
"""Sanitize connector sub-key names."""
|
||||
key = key.replace(".ff.net.0.proj.", ".ff.proj_in.")
|
||||
key = key.replace(".ff.net.2.", ".ff.proj_out.")
|
||||
key = key.replace(".to_out.0.", ".to_out.")
|
||||
return key
|
||||
|
||||
|
||||
def extract_text_projections(weights: Dict[str, mx.array]) -> Dict[str, mx.array]:
|
||||
"""Extract text projection weights (aggregate_embed + connectors).
|
||||
|
||||
Handles both LTX-2 (aggregate_embed.weight) and LTX-2.3
|
||||
(video_aggregate_embed.*, audio_aggregate_embed.*) formats.
|
||||
"""
|
||||
extracted = {}
|
||||
|
||||
# aggregate_embed weights (text_embedding_projection.*)
|
||||
for key, value in weights.items():
|
||||
if key.startswith(TEXT_PROJ_PREFIX):
|
||||
new_key = key[len(TEXT_PROJ_PREFIX):]
|
||||
extracted[new_key] = value
|
||||
|
||||
# video_embeddings_connector
|
||||
for key, value in weights.items():
|
||||
if key.startswith(VIDEO_CONNECTOR_PREFIX):
|
||||
suffix = key[len(VIDEO_CONNECTOR_PREFIX):]
|
||||
new_key = "video_embeddings_connector." + sanitize_connector_key(suffix)
|
||||
extracted[new_key] = value
|
||||
|
||||
# audio_embeddings_connector
|
||||
for key, value in weights.items():
|
||||
if key.startswith(AUDIO_CONNECTOR_PREFIX):
|
||||
suffix = key[len(AUDIO_CONNECTOR_PREFIX):]
|
||||
new_key = "audio_embeddings_connector." + sanitize_connector_key(suffix)
|
||||
extracted[new_key] = value
|
||||
|
||||
return extracted
|
||||
|
||||
|
||||
# ─── Saving utilities ─────────────────────────────────────────────────────────
|
||||
|
||||
|
||||
def save_sharded(
|
||||
weights: Dict[str, mx.array],
|
||||
output_dir: Path,
|
||||
max_shard_size_bytes: int = 5 * 1024 * 1024 * 1024, # 5GB per shard
|
||||
):
|
||||
"""Save weights as sharded safetensors with an index file."""
|
||||
output_dir.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
# Sort keys for deterministic output
|
||||
sorted_keys = sorted(weights.keys())
|
||||
|
||||
# Calculate total size
|
||||
total_size = sum(weights[k].nbytes for k in sorted_keys)
|
||||
|
||||
# Determine sharding
|
||||
shards = []
|
||||
current_shard = {}
|
||||
current_size = 0
|
||||
|
||||
for key in sorted_keys:
|
||||
tensor = weights[key]
|
||||
tensor_size = tensor.nbytes
|
||||
|
||||
if current_size + tensor_size > max_shard_size_bytes and current_shard:
|
||||
shards.append(current_shard)
|
||||
current_shard = {}
|
||||
current_size = 0
|
||||
|
||||
current_shard[key] = tensor
|
||||
current_size += tensor_size
|
||||
|
||||
if current_shard:
|
||||
shards.append(current_shard)
|
||||
|
||||
num_shards = len(shards)
|
||||
weight_map = {}
|
||||
|
||||
for i, shard in enumerate(shards):
|
||||
if num_shards == 1:
|
||||
filename = "model.safetensors"
|
||||
else:
|
||||
filename = f"model-{i+1:05d}-of-{num_shards:05d}.safetensors"
|
||||
|
||||
mx.save_safetensors(str(output_dir / filename), shard)
|
||||
|
||||
for key in shard:
|
||||
weight_map[key] = filename
|
||||
|
||||
# Write index
|
||||
index = {
|
||||
"metadata": {"total_size": total_size},
|
||||
"weight_map": weight_map,
|
||||
}
|
||||
with open(output_dir / "model.safetensors.index.json", "w") as f:
|
||||
json.dump(index, f, indent=2, sort_keys=True)
|
||||
|
||||
return num_shards
|
||||
|
||||
|
||||
def save_single(weights: Dict[str, mx.array], output_dir: Path):
|
||||
"""Save weights as a single safetensors file with an index."""
|
||||
output_dir.mkdir(parents=True, exist_ok=True)
|
||||
mx.save_safetensors(str(output_dir / "model.safetensors"), weights)
|
||||
|
||||
# Also write index for consistency
|
||||
total_size = sum(v.nbytes for v in weights.values())
|
||||
weight_map = {k: "model.safetensors" for k in sorted(weights.keys())}
|
||||
index = {
|
||||
"metadata": {"total_size": total_size},
|
||||
"weight_map": weight_map,
|
||||
}
|
||||
with open(output_dir / "model.safetensors.index.json", "w") as f:
|
||||
json.dump(index, f, indent=2, sort_keys=True)
|
||||
|
||||
|
||||
def save_config(config: dict, output_dir: Path):
|
||||
"""Save config.json to a directory."""
|
||||
output_dir.mkdir(parents=True, exist_ok=True)
|
||||
with open(output_dir / "config.json", "w") as f:
|
||||
json.dump(config, f, indent=4)
|
||||
f.write("\n")
|
||||
|
||||
|
||||
# ─── Source resolution ─────────────────────────────────────────────────────────
|
||||
|
||||
# Matches monolithic model files: ltx-2-19b-distilled.safetensors, ltx-2.3-22b-dev.safetensors, etc.
|
||||
MONOLITHIC_PATTERN = re.compile(r"^ltx-[\d.]+-\d+b-(?P<variant>distilled|dev)\.safetensors$")
|
||||
|
||||
# Matches upscaler files like ltx-2-spatial-upscaler-x2-1.0.safetensors,
|
||||
# ltx-2.3-spatial-upscaler-x2-1.0.safetensors, etc.
|
||||
UPSCALER_PATTERN = re.compile(r"^ltx-[\d.]+-(?:spatial|temporal)-upscaler-.+\.safetensors$")
|
||||
|
||||
|
||||
def resolve_source(source: str, variant: str) -> Path:
|
||||
"""Resolve source to a monolithic safetensors file path.
|
||||
|
||||
Args:
|
||||
source: HF repo ID (e.g. "Lightricks/LTX-2"), local directory, or direct file path.
|
||||
variant: Model variant ("distilled" or "dev") to select the right file.
|
||||
|
||||
Returns:
|
||||
Path to the monolithic safetensors file.
|
||||
"""
|
||||
source_path = Path(source)
|
||||
|
||||
# Direct file path
|
||||
if source_path.is_file():
|
||||
return source_path
|
||||
|
||||
# Local directory — find the variant's safetensors file
|
||||
if source_path.is_dir():
|
||||
matches = []
|
||||
for f in sorted(source_path.glob("ltx-*b-*.safetensors")):
|
||||
m = MONOLITHIC_PATTERN.match(f.name)
|
||||
if m and m.group("variant") == variant:
|
||||
matches.append(f)
|
||||
|
||||
if matches:
|
||||
return matches[0]
|
||||
|
||||
# Broader fallback
|
||||
all_mono = sorted(source_path.glob("ltx-*.safetensors"))
|
||||
for f in all_mono:
|
||||
if variant in f.name and MONOLITHIC_PATTERN.match(f.name):
|
||||
return f
|
||||
|
||||
raise FileNotFoundError(
|
||||
f"No monolithic *-{variant}.safetensors found in {source_path}. "
|
||||
f"Files found: {[f.name for f in all_mono]}"
|
||||
)
|
||||
|
||||
# HF repo ID — download via huggingface_hub
|
||||
if "/" in source and not source_path.exists():
|
||||
from huggingface_hub import hf_hub_download, list_repo_files
|
||||
|
||||
# Find the right file in the repo
|
||||
repo_files = list_repo_files(source)
|
||||
target = None
|
||||
for f in repo_files:
|
||||
m = MONOLITHIC_PATTERN.match(f)
|
||||
if m and m.group("variant") == variant:
|
||||
target = f
|
||||
break
|
||||
|
||||
if not target:
|
||||
raise FileNotFoundError(
|
||||
f"No *-{variant}.safetensors found in {source}. "
|
||||
f"Available: {[f for f in repo_files if f.endswith('.safetensors')]}"
|
||||
)
|
||||
|
||||
print(f"Downloading {target} from {source}...")
|
||||
local_path = hf_hub_download(repo_id=source, filename=target)
|
||||
return Path(local_path)
|
||||
|
||||
raise FileNotFoundError(
|
||||
f"Source not found: {source}. Provide an HF repo ID, local directory, or file path."
|
||||
)
|
||||
|
||||
|
||||
# ─── Config inference ─────────────────────────────────────────────────────────
|
||||
|
||||
|
||||
def infer_transformer_config(weights: Dict[str, mx.array]) -> dict:
|
||||
"""Infer transformer config from weight shapes."""
|
||||
# Count transformer layers
|
||||
max_layer = -1
|
||||
for key in weights:
|
||||
if "transformer_blocks." in key:
|
||||
parts = key.split(".")
|
||||
try:
|
||||
idx = parts.index("transformer_blocks") + 1
|
||||
if idx < len(parts) and parts[idx].isdigit():
|
||||
max_layer = max(max_layer, int(parts[idx]))
|
||||
except ValueError:
|
||||
pass
|
||||
num_layers = max_layer + 1 if max_layer >= 0 else 48
|
||||
|
||||
# Detect cross_attention_dim from attn2.to_k (cross-attention input dim)
|
||||
cross_attention_dim = 4096
|
||||
for key, value in weights.items():
|
||||
if "transformer_blocks.0.attn2.to_k.weight" in key:
|
||||
cross_attention_dim = value.shape[-1]
|
||||
break
|
||||
|
||||
# Check for prompt_adaln_single (LTX-2.3 feature)
|
||||
has_prompt_adaln = any("prompt_adaln_single" in k for k in weights)
|
||||
|
||||
config = {
|
||||
"attention_head_dim": 128,
|
||||
"attention_type": "default",
|
||||
"audio_attention_head_dim": 64,
|
||||
"audio_caption_channels": 3840,
|
||||
"audio_cross_attention_dim": 2048,
|
||||
"audio_in_channels": 128,
|
||||
"audio_num_attention_heads": 32,
|
||||
"audio_out_channels": 128,
|
||||
"audio_positional_embedding_max_pos": [20],
|
||||
"av_ca_timestep_scale_multiplier": 1000,
|
||||
"caption_channels": 3840,
|
||||
"cross_attention_dim": cross_attention_dim,
|
||||
"double_precision_rope": True,
|
||||
"in_channels": 128,
|
||||
"model_type": "ltx av model",
|
||||
"norm_eps": 1e-06,
|
||||
"num_attention_heads": 32,
|
||||
"num_layers": num_layers,
|
||||
"out_channels": 128,
|
||||
"positional_embedding_max_pos": [20, 2048, 2048],
|
||||
"positional_embedding_theta": 10000.0,
|
||||
"rope_type": "split",
|
||||
"timestep_scale_multiplier": 1000,
|
||||
"use_middle_indices_grid": True,
|
||||
}
|
||||
|
||||
if has_prompt_adaln:
|
||||
config["has_prompt_adaln"] = True
|
||||
|
||||
return config
|
||||
|
||||
|
||||
def infer_vae_decoder_config(weights: Dict[str, mx.array], variant: str) -> dict:
|
||||
"""Infer VAE decoder config from weights."""
|
||||
# Check for timestep conditioning keys
|
||||
has_timestep = any("last_time_embedder" in k or "last_scale_shift_table" in k for k in weights)
|
||||
|
||||
# Count channel multipliers from up_blocks
|
||||
max_block = -1
|
||||
for key in weights:
|
||||
if "up_blocks." in key:
|
||||
parts = key.split(".")
|
||||
try:
|
||||
idx = parts.index("up_blocks") + 1
|
||||
if idx < len(parts) and parts[idx].isdigit():
|
||||
max_block = max(max_block, int(parts[idx]))
|
||||
except ValueError:
|
||||
pass
|
||||
|
||||
# Default config
|
||||
config = {
|
||||
"ch": 128,
|
||||
"ch_mult": [1, 2, 4],
|
||||
"dropout": 0.0,
|
||||
"num_res_blocks": 2,
|
||||
"out_ch": 2,
|
||||
"resolution": 256,
|
||||
"timestep_conditioning": has_timestep,
|
||||
"z_channels": 8,
|
||||
}
|
||||
return config
|
||||
|
||||
|
||||
def infer_vae_encoder_config(weights: Dict[str, mx.array]) -> dict:
|
||||
"""Return VAE encoder config (architecture is consistent across versions)."""
|
||||
return {
|
||||
"convolution_dimensions": 3,
|
||||
"encoder_blocks": [
|
||||
["res_x", {"num_layers": 4}],
|
||||
["compress_space_res", {"multiplier": 2}],
|
||||
["res_x", {"num_layers": 6}],
|
||||
["compress_time_res", {"multiplier": 2}],
|
||||
["res_x", {"num_layers": 6}],
|
||||
["compress_all_res", {"multiplier": 2}],
|
||||
["res_x", {"num_layers": 2}],
|
||||
["compress_all_res", {"multiplier": 2}],
|
||||
["res_x", {"num_layers": 2}],
|
||||
],
|
||||
"encoder_spatial_padding_mode": "zeros",
|
||||
"in_channels": 3,
|
||||
"latent_log_var": "uniform",
|
||||
"norm_layer": "pixel_norm",
|
||||
"out_channels": 128,
|
||||
"patch_size": 4,
|
||||
}
|
||||
|
||||
|
||||
def infer_audio_vae_config(weights: Dict[str, mx.array]) -> dict:
|
||||
"""Return audio VAE config."""
|
||||
return {
|
||||
"attn_resolutions": [],
|
||||
"attn_type": "vanilla",
|
||||
"causality_axis": "height",
|
||||
"ch": 128,
|
||||
"ch_mult": [1, 2, 4],
|
||||
"dropout": 0.0,
|
||||
"give_pre_end": False,
|
||||
"is_causal": True,
|
||||
"mel_bins": 64,
|
||||
"mel_hop_length": 160,
|
||||
"mid_block_add_attention": False,
|
||||
"norm_type": "pixel",
|
||||
"num_res_blocks": 2,
|
||||
"out_ch": 2,
|
||||
"resamp_with_conv": True,
|
||||
"resolution": 256,
|
||||
"sample_rate": 16000,
|
||||
"tanh_out": False,
|
||||
"z_channels": 8,
|
||||
}
|
||||
|
||||
|
||||
def infer_vocoder_config(weights: Dict[str, mx.array]) -> dict:
|
||||
"""Infer vocoder config from weights."""
|
||||
# Check for bwe_generator (LTX-2.3 BigVGAN vocoder)
|
||||
has_bwe = any(k.startswith("bwe_generator") for k in weights)
|
||||
|
||||
if has_bwe:
|
||||
return {
|
||||
"type": "bigvgan",
|
||||
"has_bwe_generator": True,
|
||||
}
|
||||
|
||||
return {
|
||||
"output_sample_rate": 24000,
|
||||
"resblock_dilation_sizes": [[1, 3, 5], [1, 3, 5], [1, 3, 5]],
|
||||
"resblock_kernel_sizes": [3, 7, 11],
|
||||
"stereo": True,
|
||||
"upsample_initial_channel": 1024,
|
||||
"upsample_kernel_sizes": [16, 15, 8, 4, 4],
|
||||
"upsample_rates": [6, 5, 2, 2, 2],
|
||||
}
|
||||
|
||||
|
||||
# ─── Main ─────────────────────────────────────────────────────────────────────
|
||||
|
||||
|
||||
def convert(source: str, output_path: Path, variant: str = "distilled"):
|
||||
"""Convert monolithic safetensors to modular directory layout.
|
||||
|
||||
Args:
|
||||
source: HF repo ID (e.g. "Lightricks/LTX-2"), local directory, or file path.
|
||||
output_path: Output directory for the modular layout.
|
||||
variant: "distilled" or "dev".
|
||||
"""
|
||||
source_path = resolve_source(source, variant)
|
||||
|
||||
print(f"Loading monolithic weights from {source_path.name}...")
|
||||
all_weights = mx.load(str(source_path))
|
||||
total_keys = len(all_weights)
|
||||
print(f" Loaded {total_keys} keys")
|
||||
|
||||
# Route keys to components
|
||||
print("\nExtracting components...")
|
||||
|
||||
# 1. Transformer
|
||||
print(" [1/6] Transformer...")
|
||||
transformer_weights = sanitize_transformer(all_weights)
|
||||
num_shards = save_sharded(transformer_weights, output_path / "transformer")
|
||||
config = infer_transformer_config(transformer_weights)
|
||||
save_config(config, output_path / "transformer")
|
||||
t_params = sum(v.size for v in transformer_weights.values())
|
||||
print(f" {len(transformer_weights)} keys, {t_params:,} params, {num_shards} shards")
|
||||
|
||||
# 2. VAE Decoder
|
||||
print(" [2/6] VAE Decoder...")
|
||||
vae_decoder_weights = sanitize_vae_decoder(all_weights)
|
||||
save_single(vae_decoder_weights, output_path / "vae" / "decoder")
|
||||
config = infer_vae_decoder_config(vae_decoder_weights, variant)
|
||||
save_config(config, output_path / "vae" / "decoder")
|
||||
d_params = sum(v.size for v in vae_decoder_weights.values())
|
||||
print(f" {len(vae_decoder_weights)} keys, {d_params:,} params")
|
||||
|
||||
# 3. VAE Encoder
|
||||
print(" [3/6] VAE Encoder...")
|
||||
vae_encoder_weights = sanitize_vae_encoder(all_weights)
|
||||
save_single(vae_encoder_weights, output_path / "vae" / "encoder")
|
||||
config = infer_vae_encoder_config(vae_encoder_weights)
|
||||
save_config(config, output_path / "vae" / "encoder")
|
||||
e_params = sum(v.size for v in vae_encoder_weights.values())
|
||||
print(f" {len(vae_encoder_weights)} keys, {e_params:,} params")
|
||||
|
||||
# 4. Audio VAE Decoder
|
||||
print(" [4/6] Audio VAE Decoder...")
|
||||
audio_decoder_weights = sanitize_audio_decoder(all_weights)
|
||||
save_single(audio_decoder_weights, output_path / "audio_vae")
|
||||
config = infer_audio_vae_config(audio_decoder_weights)
|
||||
save_config(config, output_path / "audio_vae")
|
||||
a_params = sum(v.size for v in audio_decoder_weights.values())
|
||||
print(f" {len(audio_decoder_weights)} keys, {a_params:,} params")
|
||||
|
||||
# 5. Vocoder
|
||||
print(" [5/6] Vocoder...")
|
||||
vocoder_weights = sanitize_vocoder(all_weights)
|
||||
save_single(vocoder_weights, output_path / "vocoder")
|
||||
config = infer_vocoder_config(vocoder_weights)
|
||||
save_config(config, output_path / "vocoder")
|
||||
v_params = sum(v.size for v in vocoder_weights.values())
|
||||
print(f" {len(vocoder_weights)} keys, {v_params:,} params")
|
||||
|
||||
# 6. Text Projections
|
||||
print(" [6/6] Text Projections...")
|
||||
text_proj_weights = extract_text_projections(all_weights)
|
||||
tp_dir = output_path / "text_projections"
|
||||
tp_dir.mkdir(parents=True, exist_ok=True)
|
||||
mx.save_safetensors(str(tp_dir / "model.safetensors"), text_proj_weights)
|
||||
tp_params = sum(v.size for v in text_proj_weights.values())
|
||||
print(f" {len(text_proj_weights)} keys, {tp_params:,} params")
|
||||
|
||||
# 7. Copy upscaler files
|
||||
print("\nCopying upscaler files...")
|
||||
source_dir = source_path.parent
|
||||
is_hf_repo = "/" in source and not Path(source).exists()
|
||||
upscaler_files = []
|
||||
|
||||
if is_hf_repo:
|
||||
from huggingface_hub import list_repo_files
|
||||
|
||||
upscaler_files = [
|
||||
f for f in list_repo_files(source) if UPSCALER_PATTERN.match(f)
|
||||
]
|
||||
else:
|
||||
upscaler_files = [
|
||||
f.name for f in source_dir.iterdir()
|
||||
if f.is_file() and UPSCALER_PATTERN.match(f.name)
|
||||
]
|
||||
|
||||
if not upscaler_files:
|
||||
print(" No upscaler files found")
|
||||
|
||||
for upscaler_file in sorted(upscaler_files):
|
||||
dest = output_path / upscaler_file
|
||||
if dest.exists():
|
||||
print(f" {upscaler_file}: already exists, skipping")
|
||||
continue
|
||||
|
||||
local_candidate = source_dir / upscaler_file
|
||||
if local_candidate.is_file():
|
||||
shutil.copy2(str(local_candidate), str(dest))
|
||||
print(f" {upscaler_file}: copied")
|
||||
elif is_hf_repo:
|
||||
from huggingface_hub import hf_hub_download
|
||||
|
||||
print(f" {upscaler_file}: downloading from {source}...")
|
||||
downloaded = hf_hub_download(repo_id=source, filename=upscaler_file)
|
||||
shutil.copy2(downloaded, str(dest))
|
||||
print(f" {upscaler_file}: done")
|
||||
else:
|
||||
print(f" {upscaler_file}: not found, skipping")
|
||||
|
||||
# 8. Link text_encoder and tokenizer directories
|
||||
print("\nLinking text encoder & tokenizer...")
|
||||
for subdir in ["text_encoder", "tokenizer"]:
|
||||
dest = output_path / subdir
|
||||
if dest.exists():
|
||||
print(f" {subdir}/: already exists, skipping")
|
||||
continue
|
||||
|
||||
local_candidate = source_dir / subdir
|
||||
if local_candidate.is_dir():
|
||||
# Resolve through symlinks to get the real directory
|
||||
real_path = local_candidate.resolve()
|
||||
dest.symlink_to(real_path)
|
||||
print(f" {subdir}/: symlinked to {real_path}")
|
||||
elif is_hf_repo:
|
||||
from huggingface_hub import list_repo_files, snapshot_download
|
||||
|
||||
# Only download if the subdir exists in the repo
|
||||
repo_files = list_repo_files(source)
|
||||
if any(f.startswith(f"{subdir}/") for f in repo_files):
|
||||
print(f" {subdir}/: downloading from {source}...")
|
||||
snapshot_download(
|
||||
repo_id=source,
|
||||
allow_patterns=f"{subdir}/*",
|
||||
local_dir=str(output_path),
|
||||
)
|
||||
print(f" {subdir}/: done")
|
||||
else:
|
||||
print(f" {subdir}/: not in repo, skipping")
|
||||
else:
|
||||
print(f" {subdir}/: not found in source, skipping")
|
||||
|
||||
# Summary
|
||||
all_converted = (
|
||||
len(transformer_weights)
|
||||
+ len(vae_decoder_weights)
|
||||
+ len(vae_encoder_weights)
|
||||
+ len(audio_decoder_weights)
|
||||
+ len(vocoder_weights)
|
||||
+ len(text_proj_weights)
|
||||
)
|
||||
print(f"\nDone! Converted {all_converted}/{total_keys} keys")
|
||||
if all_converted < total_keys:
|
||||
# Find unconverted keys
|
||||
converted_prefixes = set()
|
||||
for key in all_weights:
|
||||
if key.startswith(TRANSFORMER_PREFIX):
|
||||
converted_prefixes.add(key)
|
||||
elif key.startswith(VAE_DECODER_PREFIX) or key.startswith(VAE_STATS_PREFIX):
|
||||
converted_prefixes.add(key)
|
||||
elif key.startswith(VAE_ENCODER_PREFIX):
|
||||
converted_prefixes.add(key)
|
||||
elif key.startswith(AUDIO_DECODER_PREFIX) or key.startswith(AUDIO_STATS_PREFIX):
|
||||
converted_prefixes.add(key)
|
||||
elif key.startswith(AUDIO_ENCODER_PREFIX):
|
||||
converted_prefixes.add(key)
|
||||
elif key.startswith(VOCODER_PREFIX):
|
||||
converted_prefixes.add(key)
|
||||
elif key.startswith(TEXT_PROJ_PREFIX):
|
||||
converted_prefixes.add(key)
|
||||
elif key.startswith(VIDEO_CONNECTOR_PREFIX) or key.startswith(AUDIO_CONNECTOR_PREFIX):
|
||||
converted_prefixes.add(key)
|
||||
|
||||
skipped = set(all_weights.keys()) - converted_prefixes
|
||||
if skipped:
|
||||
print(f" Skipped {len(skipped)} keys:")
|
||||
for k in sorted(skipped)[:20]:
|
||||
print(f" {k}")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
parser = argparse.ArgumentParser(
|
||||
description="Convert monolithic LTX-2/2.3 safetensors to modular MLX layout"
|
||||
)
|
||||
parser.add_argument(
|
||||
"--source",
|
||||
type=str,
|
||||
required=True,
|
||||
help="HF repo ID (e.g. Lightricks/LTX-2, Lightricks/LTX-2.3), local directory, or direct safetensors file path",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--output",
|
||||
type=str,
|
||||
required=True,
|
||||
help="Output directory for modular layout",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--variant",
|
||||
type=str,
|
||||
choices=["distilled", "dev"],
|
||||
default="distilled",
|
||||
help="Model variant (affects VAE decoder config and which file to download)",
|
||||
)
|
||||
args = parser.parse_args()
|
||||
|
||||
convert(args.source, Path(args.output), variant=args.variant)
|
||||
Reference in New Issue
Block a user