ltx-2-5 / .meta /create_model.py

Upload folder using huggingface_hub

bbae6ac verified 5 days ago

16.9 kB

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

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