Function main

scripts/convert_sd3_to_diffusers.py:268–347

(args)

def main(args):
    original_ckpt = load_original_checkpoint(args.checkpoint_path)
    original_dtype = next(iter(original_ckpt.values())).dtype

    # Initialize dtype with a default value
    dtype = None

    if args.dtype is None:
        dtype = original_dtype
    elif args.dtype == "fp16":
        dtype = torch.float16
    elif args.dtype == "bf16":
        dtype = torch.bfloat16
    elif args.dtype == "fp32":
        dtype = torch.float32
    else:
        raise ValueError(f"Unsupported dtype: {args.dtype}")

    if dtype != original_dtype:
        print(
            f"Checkpoint dtype {original_dtype} does not match requested dtype {dtype}. This can lead to unexpected results, proceed with caution."
        )

    num_layers = list(set(int(k.split(".", 2)[1]) for k in original_ckpt if "joint_blocks" in k))[-1] + 1  # noqa: C401

    caption_projection_dim = get_caption_projection_dim(original_ckpt)

    # () for sd3.0; (0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12) for sd3.5
    attn2_layers = get_attn2_layers(original_ckpt)

    # sd3.5 use qk norm("rms_norm")
    has_qk_norm = any("ln_q" in key for key in original_ckpt.keys())

    # sd3.5 2b use pox_embed_max_size=384 and sd3.0 and sd3.5 8b use 192
    pos_embed_max_size = get_pos_embed_max_size(original_ckpt)

    converted_transformer_state_dict = convert_sd3_transformer_checkpoint_to_diffusers(
        original_ckpt, num_layers, caption_projection_dim, attn2_layers, has_qk_norm
    )

    with CTX():
        transformer = SD3Transformer2DModel(
            sample_size=128,
            patch_size=2,
            in_channels=16,
            joint_attention_dim=4096,
            num_layers=num_layers,
            caption_projection_dim=caption_projection_dim,
            num_attention_heads=num_layers,
            pos_embed_max_size=pos_embed_max_size,
            qk_norm="rms_norm" if has_qk_norm else None,
            dual_attention_layers=attn2_layers,
        )
    if is_accelerate_available():
        load_model_dict_into_meta(transformer, converted_transformer_state_dict)
    else:
        transformer.load_state_dict(converted_transformer_state_dict, strict=True)