File size: 5,989 Bytes
bbae6ac | 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 | """Load the tiny LTX-2.5 split checkpoint into a Diffusers pipeline."""
import json
import tempfile
from pathlib import Path
import torch
from diffusers import (
AutoencoderKLLTX2Audio,
AutoencoderKLLTX2Video,
FlowMatchEulerDiscreteScheduler,
LTX2Pipeline,
LTX2VideoTransformer3DModel,
)
from diffusers.loaders.single_file_utils import (
convert_ltx2_audio_vae_to_diffusers,
convert_ltx2_transformer_to_diffusers,
convert_ltx2_vae_to_diffusers,
)
from diffusers.pipelines.ltx2.connectors import LTX2TextConnectors
from diffusers.pipelines.ltx2.vocoder import LTX2Vocoder
from safetensors.torch import load_file
from transformers import (
AutoTokenizer,
Gemma4UnifiedConfig,
Gemma4UnifiedForConditionalGeneration,
Gemma4UnifiedTextConfig,
Gemma4UnifiedVisionConfig,
)
FILES = {
"transformer": "diffusion_models/ltx-2.5-22b-distilled-transformer-bf16.safetensors",
"text_encoder": "text_encoders/gemma4-12b-with-proj-ltx-2.5-bf16.safetensors",
"vae": "vae/ltx-2.5-video-vae-bf16.safetensors",
"audio_vae": "vae/ltx-2.5-audio-vae-bf16.safetensors",
}
ASSETS = {
"tokenizer_json": "tokenizer.json",
"hf_asset__tokenizer_config.json": "tokenizer_config.json",
"hf_asset__chat_template.jinja": "chat_template.jinja",
}
# Diffusers name -> Comfy single-file name for the tensors the converters drop.
CONNECTOR_PREFIX = {
"video_connector.transformer_blocks.": "video_embeddings_connector.transformer_1d_blocks.",
"audio_connector.transformer_blocks.": "audio_embeddings_connector.transformer_1d_blocks.",
"video_connector.learnable_registers": "video_embeddings_connector.learnable_registers",
"audio_connector.learnable_registers": "audio_embeddings_connector.learnable_registers",
}
TEXT_PREFIX = {
"model.language_model.": "model.",
"model.embed_vision.": "vision_model.",
}
PROJECTIONS = {
"video_text_proj_in.weight": "text_embedding_projection.video_aggregate_embed.weight",
"audio_text_proj_in.weight": "text_embedding_projection.audio_aggregate_embed.weight",
}
def _connector_name(name):
for local, source in CONNECTOR_PREFIX.items():
if name.startswith(local):
name = source + name[len(local) :]
break
return name.replace("norm_q", "q_norm").replace("norm_k", "k_norm")
def _text_name(name):
for local, source in TEXT_PREFIX.items():
if name.startswith(local):
return source + name[len(local) :]
return None
def load_transformer(config, state):
model = LTX2VideoTransformer3DModel(**config["transformer"])
plain = {name: value for name, value in state.items() if "embeddings_connector" not in name}
model.load_state_dict(convert_ltx2_transformer_to_diffusers(plain))
return model
def load_connectors(config, transformer_state, text_state):
model = LTX2TextConnectors(**config["connectors"])
weights = {}
for name in model.state_dict():
if name in PROJECTIONS:
weights[name] = text_state[PROJECTIONS[name]]
else:
weights[name] = transformer_state["model.diffusion_model." + _connector_name(name)]
model.load_state_dict(weights, strict=True)
return model
def load_text_encoder(config, text_state):
text = Gemma4UnifiedTextConfig(
**{key: value for key, value in config["text_encoder"].items() if key != "tie_word_embeddings"}
)
gemma = Gemma4UnifiedConfig(
text_config=text.to_dict(),
vision_config=Gemma4UnifiedVisionConfig(
patch_size=16, pooling_kernel_size=2, mm_embed_dim=8, mm_posemb_size=16, output_proj_dims=8
).to_dict(),
)
gemma.tie_word_embeddings = True
model = Gemma4UnifiedForConditionalGeneration(gemma)
weights = {}
for name in model.state_dict():
source = _text_name(name)
if source is not None and source in text_state:
weights[name] = text_state[source]
model.load_state_dict(weights, strict=False)
return model
def load_tokenizer(text_state):
with tempfile.TemporaryDirectory() as directory:
for tensor_name, file_name in ASSETS.items():
Path(directory, file_name).write_bytes(text_state[tensor_name].numpy().tobytes())
return AutoTokenizer.from_pretrained(directory)
def load_pipeline(model_id, device=None):
model_id = Path(model_id)
config = json.loads((model_id / "config.json").read_text())
transformer_state = load_file(str(model_id / FILES["transformer"]))
text_state = load_file(str(model_id / FILES["text_encoder"]))
transformer = load_transformer(config, transformer_state)
connectors = load_connectors(config, transformer_state, text_state)
vae = AutoencoderKLLTX2Video(**config["vae"])
vae.load_state_dict(convert_ltx2_vae_to_diffusers(load_file(str(model_id / FILES["vae"]))))
audio_state = load_file(str(model_id / FILES["audio_vae"]))
audio_vae = AutoencoderKLLTX2Audio(**config["audio_vae"])
audio_vae.load_state_dict(
convert_ltx2_audio_vae_to_diffusers(
{name: value for name, value in audio_state.items() if not name.startswith("vocoder.")}
)
)
vocoder = LTX2Vocoder(**config["vocoder"])
vocoder.load_state_dict(
{
name.removeprefix("vocoder."): value
for name, value in audio_state.items()
if name.startswith("vocoder.")
}
)
text_encoder = load_text_encoder(config, text_state)
pipeline = LTX2Pipeline(
scheduler=FlowMatchEulerDiscreteScheduler(),
vae=vae,
audio_vae=audio_vae,
text_encoder=text_encoder,
tokenizer=load_tokenizer(text_state),
connectors=connectors,
transformer=transformer,
vocoder=vocoder,
)
for module in (transformer, vae, audio_vae, text_encoder, connectors, vocoder):
module.eval()
pipeline.to(device or ("cuda" if torch.cuda.is_available() else "cpu"))
return pipeline
|