convert_ltx_to_diffusers.py

importargparse
frompathlibimportPath
fromtypingimportAny, Dict

importtorch
fromaccelerateimportinit_empty_weights
fromsafetensors.torchimportload_file
fromtransformersimportT5EncoderModel, T5Tokenizer

fromdiffusersimportAutoencoderKLLTXVideo, FlowMatchEulerDiscreteScheduler, LTXPipeline, LTXVideoTransformer3DModel


defremove_keys_(key: str, state_dict: Dict[str, Any]):
state_dict.pop(key)


TOKENIZER_MAX_LENGTH=128

TRANSFORMER_KEYS_RENAME_DICT= {
"patchify_proj": "proj_in",
"adaln_single": "time_embed",
"q_norm": "norm_q",
"k_norm": "norm_k",
}

TRANSFORMER_SPECIAL_KEYS_REMAP= {
"vae": remove_keys_,
}

VAE_KEYS_RENAME_DICT= {
# decoder
"up_blocks.0": "mid_block",
"up_blocks.1": "up_blocks.0",
"up_blocks.2": "up_blocks.1.upsamplers.0",
"up_blocks.3": "up_blocks.1",
"up_blocks.4": "up_blocks.2.conv_in",
"up_blocks.5": "up_blocks.2.upsamplers.0",
"up_blocks.6": "up_blocks.2",
"up_blocks.7": "up_blocks.3.conv_in",
"up_blocks.8": "up_blocks.3.upsamplers.0",
"up_blocks.9": "up_blocks.3",
# encoder
"down_blocks.0": "down_blocks.0",
"down_blocks.1": "down_blocks.0.downsamplers.0",
"down_blocks.2": "down_blocks.0.conv_out",
"down_blocks.3": "down_blocks.1",
"down_blocks.4": "down_blocks.1.downsamplers.0",
"down_blocks.5": "down_blocks.1.conv_out",
"down_blocks.6": "down_blocks.2",
"down_blocks.7": "down_blocks.2.downsamplers.0",
"down_blocks.8": "down_blocks.3",
"down_blocks.9": "mid_block",
# common
"conv_shortcut": "conv_shortcut.conv",
"res_blocks": "resnets",
"norm3.norm": "norm3",
"per_channel_statistics.mean-of-means": "latents_mean",
"per_channel_statistics.std-of-means": "latents_std",
}

VAE_091_RENAME_DICT= {
# decoder
"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",
"up_blocks.7": "up_blocks.3.upsamplers.0",
"up_blocks.8": "up_blocks.3",
# common
"last_time_embedder": "time_embedder",
"last_scale_shift_table": "scale_shift_table",
}

VAE_095_RENAME_DICT= {
# decoder
"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",
"up_blocks.7": "up_blocks.3.upsamplers.0",
"up_blocks.8": "up_blocks.3",
# encoder
"down_blocks.0": "down_blocks.0",
"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",
# common
"last_time_embedder": "time_embedder",
"last_scale_shift_table": "scale_shift_table",
}

VAE_SPECIAL_KEYS_REMAP= {
"per_channel_statistics.channel": remove_keys_,
"per_channel_statistics.mean-of-means": remove_keys_,
"per_channel_statistics.mean-of-stds": remove_keys_,
"model.diffusion_model": remove_keys_,
}


defget_state_dict(saved_dict: Dict[str, Any]) ->Dict[str, Any]:
state_dict=saved_dict
if"model"insaved_dict.keys():
state_dict=state_dict["model"]
if"module"insaved_dict.keys():
state_dict=state_dict["module"]
if"state_dict"insaved_dict.keys():
state_dict=state_dict["state_dict"]
returnstate_dict


defupdate_state_dict_inplace(state_dict: Dict[str, Any], old_key: str, new_key: str) ->Dict[str, Any]:
state_dict[new_key] =state_dict.pop(old_key)


defconvert_transformer(
ckpt_path: str,
dtype: torch.dtype,
version: str="0.9.0",
):
PREFIX_KEY="model.diffusion_model."

original_state_dict=get_state_dict(load_file(ckpt_path))
config= {}
ifversion=="0.9.5":
config["_use_causal_rope_fix"] =True
withinit_empty_weights():
transformer=LTXVideoTransformer3DModel(**config)

forkeyinlist(original_state_dict.keys()):
new_key=key[:]
ifnew_key.startswith(PREFIX_KEY):
new_key=key[len(PREFIX_KEY) :]
forreplace_key, rename_keyinTRANSFORMER_KEYS_RENAME_DICT.items():
new_key=new_key.replace(replace_key, rename_key)
update_state_dict_inplace(original_state_dict, key, new_key)

forkeyinlist(original_state_dict.keys()):
forspecial_key, handler_fn_inplaceinTRANSFORMER_SPECIAL_KEYS_REMAP.items():
ifspecial_keynotinkey:
continue
handler_fn_inplace(key, original_state_dict)

transformer.load_state_dict(original_state_dict, strict=True, assign=True)
returntransformer


defconvert_vae(ckpt_path: str, config, dtype: torch.dtype):
PREFIX_KEY="vae."

original_state_dict=get_state_dict(load_file(ckpt_path))
withinit_empty_weights():
vae=AutoencoderKLLTXVideo(**config)

forkeyinlist(original_state_dict.keys()):
new_key=key[:]
ifnew_key.startswith(PREFIX_KEY):
new_key=key[len(PREFIX_KEY) :]
forreplace_key, rename_keyinVAE_KEYS_RENAME_DICT.items():
new_key=new_key.replace(replace_key, rename_key)
update_state_dict_inplace(original_state_dict, key, new_key)

forkeyinlist(original_state_dict.keys()):
forspecial_key, handler_fn_inplaceinVAE_SPECIAL_KEYS_REMAP.items():
ifspecial_keynotinkey:
continue
handler_fn_inplace(key, original_state_dict)

vae.load_state_dict(original_state_dict, strict=True, assign=True)
returnvae


defget_vae_config(version: str) ->Dict[str, Any]:
ifversion=="0.9.0":
config= {
"in_channels": 3,
"out_channels": 3,
"latent_channels": 128,
"block_out_channels": (128, 256, 512, 512),
"down_block_types": (
"LTXVideoDownBlock3D",
"LTXVideoDownBlock3D",
"LTXVideoDownBlock3D",
"LTXVideoDownBlock3D",
 ),
"decoder_block_out_channels": (128, 256, 512, 512),
"layers_per_block": (4, 3, 3, 3, 4),
"decoder_layers_per_block": (4, 3, 3, 3, 4),
"spatio_temporal_scaling": (True, True, True, False),
"decoder_spatio_temporal_scaling": (True, True, True, False),
"decoder_inject_noise": (False, False, False, False, False),
"downsample_type": ("conv", "conv", "conv", "conv"),
"upsample_residual": (False, False, False, False),
"upsample_factor": (1, 1, 1, 1),
"patch_size": 4,
"patch_size_t": 1,
"resnet_norm_eps": 1e-6,
"scaling_factor": 1.0,
"encoder_causal": True,
"decoder_causal": False,
"timestep_conditioning": False,
 }
elifversion=="0.9.1":
config= {
"in_channels": 3,
"out_channels": 3,
"latent_channels": 128,
"block_out_channels": (128, 256, 512, 512),
"down_block_types": (
"LTXVideoDownBlock3D",
"LTXVideoDownBlock3D",
"LTXVideoDownBlock3D",
"LTXVideoDownBlock3D",
 ),
"decoder_block_out_channels": (256, 512, 1024),
"layers_per_block": (4, 3, 3, 3, 4),
"decoder_layers_per_block": (5, 6, 7, 8),
"spatio_temporal_scaling": (True, True, True, False),
"decoder_spatio_temporal_scaling": (True, True, True),
"decoder_inject_noise": (True, True, True, False),
"downsample_type": ("conv", "conv", "conv", "conv"),
"upsample_residual": (True, True, True),
"upsample_factor": (2, 2, 2),
"timestep_conditioning": True,
"patch_size": 4,
"patch_size_t": 1,
"resnet_norm_eps": 1e-6,
"scaling_factor": 1.0,
"encoder_causal": True,
"decoder_causal": False,
 }
VAE_KEYS_RENAME_DICT.update(VAE_091_RENAME_DICT)
elifversion=="0.9.5":
config= {
"in_channels": 3,
"out_channels": 3,
"latent_channels": 128,
"block_out_channels": (128, 256, 512, 1024, 2048),
"down_block_types": (
"LTXVideo095DownBlock3D",
"LTXVideo095DownBlock3D",
"LTXVideo095DownBlock3D",
"LTXVideo095DownBlock3D",
 ),
"decoder_block_out_channels": (256, 512, 1024),
"layers_per_block": (4, 6, 6, 2, 2),
"decoder_layers_per_block": (5, 5, 5, 5),
"spatio_temporal_scaling": (True, True, True, True),
"decoder_spatio_temporal_scaling": (True, True, True),
"decoder_inject_noise": (False, False, False, False),
"downsample_type": ("spatial", "temporal", "spatiotemporal", "spatiotemporal"),
"upsample_residual": (True, True, True),
"upsample_factor": (2, 2, 2),
"timestep_conditioning": True,
"patch_size": 4,
"patch_size_t": 1,
"resnet_norm_eps": 1e-6,
"scaling_factor": 1.0,
"encoder_causal": True,
"decoder_causal": False,
"spatial_compression_ratio": 32,
"temporal_compression_ratio": 8,
 }
VAE_KEYS_RENAME_DICT.update(VAE_095_RENAME_DICT)
returnconfig


defget_args():
parser=argparse.ArgumentParser()
parser.add_argument(
"--transformer_ckpt_path", type=str, default=None, help="Path to original transformer checkpoint"
 )
parser.add_argument("--vae_ckpt_path", type=str, default=None, help="Path to original vae checkpoint")
parser.add_argument(
"--text_encoder_cache_dir", type=str, default=None, help="Path to text encoder cache directory"
 )
parser.add_argument(
"--typecast_text_encoder",
action="store_true",
default=False,
help="Whether or not to apply fp16/bf16 precision to text_encoder",
 )
parser.add_argument("--save_pipeline", action="store_true")
parser.add_argument("--output_path", type=str, required=True, help="Path where converted model should be saved")
parser.add_argument("--dtype", default="fp32", help="Torch dtype to save the model in.")
parser.add_argument(
"--version", type=str, default="0.9.0", choices=["0.9.0", "0.9.1", "0.9.5"], help="Version of the LTX model"
 )
returnparser.parse_args()


DTYPE_MAPPING= {
"fp32": torch.float32,
"fp16": torch.float16,
"bf16": torch.bfloat16,
}

VARIANT_MAPPING= {
"fp32": None,
"fp16": "fp16",
"bf16": "bf16",
}


if__name__=="__main__":
args=get_args()

transformer=None
dtype=DTYPE_MAPPING[args.dtype]
variant=VARIANT_MAPPING[args.dtype]
output_path=Path(args.output_path)

ifargs.save_pipeline:
assertargs.transformer_ckpt_pathisnotNoneandargs.vae_ckpt_pathisnotNone

ifargs.transformer_ckpt_pathisnotNone:
transformer: LTXVideoTransformer3DModel=convert_transformer(args.transformer_ckpt_path, dtype)
ifnotargs.save_pipeline:
transformer.save_pretrained(
output_path/"transformer", safe_serialization=True, max_shard_size="5GB", variant=variant
 )

ifargs.vae_ckpt_pathisnotNone:
config=get_vae_config(args.version)
vae: AutoencoderKLLTXVideo=convert_vae(args.vae_ckpt_path, config, dtype)
ifnotargs.save_pipeline:
vae.save_pretrained(output_path/"vae", safe_serialization=True, max_shard_size="5GB", variant=variant)

ifargs.save_pipeline:
text_encoder_id="google/t5-v1_1-xxl"
tokenizer=T5Tokenizer.from_pretrained(text_encoder_id, model_max_length=TOKENIZER_MAX_LENGTH)
text_encoder=T5EncoderModel.from_pretrained(text_encoder_id, cache_dir=args.text_encoder_cache_dir)

ifargs.typecast_text_encoder:
text_encoder=text_encoder.to(dtype=dtype)

# Apparently, the conversion does not work anymore without this :shrug:
forparamintext_encoder.parameters():
param.data=param.data.contiguous()

ifargs.version=="0.9.5":
scheduler=FlowMatchEulerDiscreteScheduler(use_dynamic_shifting=False)
else:
scheduler=FlowMatchEulerDiscreteScheduler(
use_dynamic_shifting=True,
base_shift=0.95,
max_shift=2.05,
base_image_seq_len=1024,
max_image_seq_len=4096,
shift_terminal=0.1,
 )

pipe=LTXPipeline(
scheduler=scheduler,
vae=vae,
text_encoder=text_encoder,
tokenizer=tokenizer,
transformer=transformer,
 )

pipe.save_pretrained(args.output_path, safe_serialization=True, variant=variant, max_shard_size="5GB")