def create_model(source_model_id: str, source_revision: str, save_folder: str) -> str: import json from pathlib import Path import torch from diffusers import AutoencoderKLLTX2Audio, AutoencoderKLLTX2Video, LTX2VideoTransformer3DModel from diffusers.pipelines.ltx2.connectors import LTX2TextConnectors from diffusers.pipelines.ltx2.vocoder import LTX2Vocoder from huggingface_hub import hf_hub_download from safetensors.torch import save_file from transformers import ( Gemma4UnifiedConfig, Gemma4UnifiedForConditionalGeneration, Gemma4UnifiedTextConfig, Gemma4UnifiedVisionConfig, set_seed, ) tokenizer_source = "google/gemma-4-12b-it" tokenizer_revision = "43c9d014b231474fa7212dbb48a762566c2b255d" folder = Path(save_folder) folder.mkdir(parents=True, exist_ok=True) for name in ("diffusion_models", "text_encoders", "vae"): (folder / name).mkdir(exist_ok=True) set_seed(42) # The split single-file layout carries no config; record the reduced # dimensions so the components can be rebuilt from the published tensors. config = { "transformer": { "in_channels": 128, "out_channels": 128, "audio_in_channels": 8, "audio_out_channels": 8, "num_layers": 4, "num_attention_heads": 2, "attention_head_dim": 32, "audio_num_attention_heads": 2, "audio_attention_head_dim": 32, "cross_attention_dim": 64, "audio_cross_attention_dim": 64, "caption_channels": 8, "use_prompt_embeddings": False, "cross_attn_mod": True, "audio_cross_attn_mod": True, "use_prompt_adaln_single": False, "use_keyframes_abs_pos_embedding": True, }, "connectors": { "caption_channels": 8, "text_proj_in_factor": 3, "video_connector_num_attention_heads": 2, "video_connector_attention_head_dim": 32, "video_connector_num_layers": 2, "video_connector_num_learnable_registers": 4, "audio_connector_num_attention_heads": 2, "audio_connector_attention_head_dim": 32, "audio_connector_num_layers": 2, "audio_connector_num_learnable_registers": 4, "video_hidden_dim": 64, "audio_hidden_dim": 64, "per_modality_projections": True, }, "vae": { "in_channels": 3, "out_channels": 3, "latent_channels": 128, "block_out_channels": [16, 32], "decoder_block_out_channels": [16], "layers_per_block": [1, 1, 1], "decoder_layers_per_block": [1, 1], "spatio_temporal_scaling": [True, True], "decoder_spatio_temporal_scaling": [True], "downsample_type": ["spatiotemporal", "spatiotemporal"], "upsample_type": ["spatiotemporal"], "upsample_residual": [True], "upsample_factor": [2], "patch_size": 4, "patch_size_t": 1, }, "audio_vae": { "base_channels": 8, "output_channels": 2, "ch_mult": [1, 2], "num_res_blocks": 1, "in_channels": 2, "resolution": 64, "latent_channels": 8, "mel_bins": 4, }, "vocoder": { "in_channels": 8, "hidden_channels": 32, "out_channels": 2, "upsample_kernel_sizes": [8, 8], "upsample_factors": [4, 4], "resnet_kernel_sizes": [3], "resnet_dilations": [[1, 1, 1]], }, "text_encoder": { "vocab_size": 262144, "hidden_size": 8, "intermediate_size": 16, "num_hidden_layers": 2, "num_attention_heads": 1, "num_key_value_heads": 1, "head_dim": 32, "max_position_embeddings": 4096, "tie_word_embeddings": True, }, } (folder / "config.json").write_text(json.dumps(config, indent=2) + "\n") transformer_renames = [ ("audio_proj_in", "audio_patchify_proj"), ("proj_in", "patchify_proj"), ("av_cross_attn_video_scale_shift", "av_ca_video_scale_shift_adaln_single"), ("av_cross_attn_video_a2v_gate", "av_ca_a2v_gate_adaln_single"), ("av_cross_attn_audio_scale_shift", "av_ca_audio_scale_shift_adaln_single"), ("av_cross_attn_audio_v2a_gate", "av_ca_v2a_gate_adaln_single"), ("video_a2v_cross_attn_scale_shift_table", "scale_shift_table_a2v_ca_video"), ("audio_a2v_cross_attn_scale_shift_table", "scale_shift_table_a2v_ca_audio"), ("norm_q", "q_norm"), ("norm_k", "k_norm"), ("audio_time_embed", "audio_adaln_single"), ("time_embed", "adaln_single"), ] vae_renames = [ ("down_blocks.1", "down_blocks.0.downsamplers.0"), ("down_blocks.2", "down_blocks.1"), ("down_blocks.3", "down_blocks.1.downsamplers.0"), ("down_blocks.4", "down_blocks.2"), ("down_blocks.5", "down_blocks.2.downsamplers.0"), ("down_blocks.6", "down_blocks.3"), ("down_blocks.7", "down_blocks.3.downsamplers.0"), ("down_blocks.8", "mid_block"), ("up_blocks.0", "mid_block"), ("up_blocks.1", "up_blocks.0.upsamplers.0"), ("up_blocks.2", "up_blocks.0"), ("up_blocks.3", "up_blocks.1.upsamplers.0"), ("up_blocks.4", "up_blocks.1"), ("up_blocks.5", "up_blocks.2.upsamplers.0"), ("up_blocks.6", "up_blocks.2"), ("resnets", "res_blocks"), ("latents_mean", "per_channel_statistics.mean-of-means"), ("latents_std", "per_channel_statistics.std-of-means"), ] def to_comfy(state, renames): output = {} for key, value in state.items(): for comfy, diffusers in reversed(renames): key = key.replace(diffusers, comfy) output[key] = value return output def cast(state): # The source stores modulation tables in FP32 and every other weight in BF16. return {k: v.float() if "scale_shift_table" in k else v.to(torch.bfloat16) for k, v in state.items()} transformer = LTX2VideoTransformer3DModel(**config["transformer"]) connectors = LTX2TextConnectors(**config["connectors"]) state = { f"model.diffusion_model.{k}": v for k, v in to_comfy(cast(transformer.state_dict()), transformer_renames).items() } connector_weights = cast(connectors.state_dict()) for key, value in connector_weights.items(): if key.startswith("video_connector.transformer_blocks."): name = key.replace( "video_connector.transformer_blocks.", "video_embeddings_connector.transformer_1d_blocks." ) elif key.startswith("audio_connector.transformer_blocks."): name = key.replace( "audio_connector.transformer_blocks.", "audio_embeddings_connector.transformer_1d_blocks." ) else: continue state[f"model.diffusion_model.{name.replace('norm_q', 'q_norm').replace('norm_k', 'k_norm')}"] = value for side in ("video", "audio"): state[f"model.diffusion_model.{side}_embeddings_connector.learnable_registers"] = connector_weights[ f"{side}_connector.learnable_registers" ] save_file(state, str(folder / "diffusion_models/ltx-2.5-22b-distilled-transformer-bf16.safetensors")) vae = AutoencoderKLLTX2Video(**config["vae"]) save_file(cast(to_comfy(vae.state_dict(), vae_renames)), str(folder / "vae/ltx-2.5-video-vae-bf16.safetensors")) audio_vae = AutoencoderKLLTX2Audio(**config["audio_vae"]) vocoder = LTX2Vocoder(**config["vocoder"]) audio_state = {f"audio_vae.{k}": v for k, v in cast(audio_vae.state_dict()).items()} audio_state.update({f"vocoder.{k}": v for k, v in cast(vocoder.state_dict()).items()}) save_file(audio_state, str(folder / "vae/ltx-2.5-audio-vae-bf16.safetensors")) text = Gemma4UnifiedTextConfig(**{k: v for k, v in config["text_encoder"].items() if k != "tie_word_embeddings"}) vision = Gemma4UnifiedVisionConfig( patch_size=16, pooling_kernel_size=2, mm_embed_dim=8, mm_posemb_size=16, output_proj_dims=8 ) text_config = Gemma4UnifiedConfig(text_config=text.to_dict(), vision_config=vision.to_dict()) text_config.tie_word_embeddings = True text_encoder = Gemma4UnifiedForConditionalGeneration(text_config) text_state = {} for key, value in cast(text_encoder.state_dict()).items(): if key == "lm_head.weight": continue # The source names the text tower `model.*` and the vision tower `vision_model.*`. key = key.replace("model.language_model.", "model.", 1) key = key.replace("model.embed_vision.", "vision_model.", 1) text_state[key] = value # The source keeps the connector projections and the two modality projectors # inside the text encoder file. projection = cast(connectors.state_dict()) text_state["text_embedding_projection.video_aggregate_embed.weight"] = projection["video_text_proj_in.weight"] text_state["text_embedding_projection.audio_aggregate_embed.weight"] = projection["audio_text_proj_in.weight"] hidden = config["text_encoder"]["hidden_size"] text_state["multi_modal_projector.embedding_projection.weight"] = torch.randn(hidden, hidden).to(torch.bfloat16) text_state["audio_projector.embedding_projection.weight"] = torch.randn( hidden, config["connectors"]["audio_hidden_dim"] ).to(torch.bfloat16) # The source embeds the tokenizer and processor assets as U8 tensors; keep # that layout instead of publishing a second copy in a tokenizer folder. assets = { "tokenizer_json": "tokenizer.json", "hf_asset__tokenizer_config.json": "tokenizer_config.json", "hf_asset__chat_template.jinja": "chat_template.jinja", "hf_asset__generation_config.json": "generation_config.json", "hf_asset__processor_config.json": "processor_config.json", } for tensor_name, file_name in assets.items(): data = Path(hf_hub_download(tokenizer_source, file_name, revision=tokenizer_revision)).read_bytes() text_state[tensor_name] = torch.frombuffer(bytearray(data), dtype=torch.uint8) save_file(text_state, str(folder / "text_encoders/gemma4-12b-with-proj-ltx-2.5-bf16.safetensors")) loader_source = '''"""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 ''' (folder / "ltx2_loader.py").write_text(loader_source) return repr(transformer) create_model( source_model_id="Lightricks/LTX-2.5", source_revision="49b093a39c19169a8a81721dcc541e07e945232f", save_folder="/tmp/ltx-2-5", )