| """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", |
| } |
| |
| 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 |
|
|
