Function main

scripts/convert_ltx2_to_diffusers.py:1049–1189

(args)

def main(args):
    vae_dtype = DTYPE_MAPPING[args.vae_dtype]
    audio_vae_dtype = DTYPE_MAPPING[args.audio_vae_dtype]
    dit_dtype = DTYPE_MAPPING[args.dit_dtype]
    vocoder_dtype = DTYPE_MAPPING[args.vocoder_dtype]
    text_encoder_dtype = DTYPE_MAPPING[args.text_encoder_dtype]

    combined_ckpt = None
    load_combined_models = any(
        [
            args.vae,
            args.audio_vae,
            args.dit,
            args.vocoder,
            args.connectors,
            args.full_pipeline,
            args.upsample_pipeline,
        ]
    )
    if args.combined_filename is not None and load_combined_models:
        combined_ckpt = load_original_checkpoint(args, filename=args.combined_filename)

    if args.vae or args.full_pipeline or args.upsample_pipeline:
        if args.vae_filename is not None:
            original_vae_ckpt = load_hub_or_local_checkpoint(filename=args.vae_filename)
        elif combined_ckpt is not None:
            original_vae_ckpt = get_model_state_dict_from_combined_ckpt(combined_ckpt, args.vae_prefix)
        vae = convert_ltx2_video_vae(
            original_vae_ckpt, version=args.version, timestep_conditioning=args.timestep_conditioning
        )
        if not args.full_pipeline and not args.upsample_pipeline:
            vae.to(vae_dtype).save_pretrained(os.path.join(args.output_path, "vae"))

    if args.audio_vae or args.full_pipeline:
        if args.audio_vae_filename is not None:
            original_audio_vae_ckpt = load_hub_or_local_checkpoint(filename=args.audio_vae_filename)
        elif combined_ckpt is not None:
            original_audio_vae_ckpt = get_model_state_dict_from_combined_ckpt(combined_ckpt, args.audio_vae_prefix)
        audio_vae = convert_ltx2_audio_vae(original_audio_vae_ckpt, version=args.version)
        if not args.full_pipeline:
            audio_vae.to(audio_vae_dtype).save_pretrained(os.path.join(args.output_path, "audio_vae"))

    if args.dit or args.full_pipeline:
        if args.dit_filename is not None:
            original_dit_ckpt = load_hub_or_local_checkpoint(filename=args.dit_filename)
        elif combined_ckpt is not None:
            original_dit_ckpt = get_model_state_dict_from_combined_ckpt(combined_ckpt, args.dit_prefix)
        transformer = convert_ltx2_transformer(original_dit_ckpt, version=args.version)
        if not args.full_pipeline:
            transformer.to(dit_dtype).save_pretrained(os.path.join(args.output_path, "transformer"))

    if args.connectors or args.full_pipeline:
        if args.dit_filename is not None:
            original_connectors_ckpt = load_hub_or_local_checkpoint(filename=args.dit_filename)
        elif combined_ckpt is not None:
            original_connectors_ckpt = get_model_state_dict_from_combined_ckpt(combined_ckpt, args.dit_prefix)
        connectors = convert_ltx2_connectors(original_connectors_ckpt, version=args.version)
        if not args.full_pipeline: