Function main

scripts/convert_cogview4_to_diffusers.py:166–250

(args)

def main(args):
    if 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}")

    transformer = None
    vae = None

    if args.transformer_checkpoint_path is not None:
        converted_transformer_state_dict = convert_cogview4_transformer_checkpoint_to_diffusers(
            args.transformer_checkpoint_path
        )
        transformer = CogView4Transformer2DModel(
            patch_size=2,
            in_channels=16,
            num_layers=28,
            attention_head_dim=128,
            num_attention_heads=32,
            out_channels=16,
            text_embed_dim=4096,
            time_embed_dim=512,
            condition_dim=256,
            pos_embed_max_size=128,
        )
        transformer.load_state_dict(converted_transformer_state_dict, strict=True)
        if dtype is not None:
            # Original checkpoint data type will be preserved
            transformer = transformer.to(dtype=dtype)

    if args.vae_checkpoint_path is not None:
        vae_config = {
            "in_channels": 3,
            "out_channels": 3,
            "down_block_types": ("DownEncoderBlock2D",) * 4,
            "up_block_types": ("UpDecoderBlock2D",) * 4,
            "block_out_channels": (128, 512, 1024, 1024),
            "layers_per_block": 3,
            "act_fn": "silu",
            "latent_channels": 16,
            "norm_num_groups": 32,
            "sample_size": 1024,
            "scaling_factor": 1.0,
            "shift_factor": 0.0,
            "force_upcast": True,
            "use_quant_conv": False,
            "use_post_quant_conv": False,
            "mid_block_add_attention": False,
        }
        converted_vae_state_dict = convert_cogview4_vae_checkpoint_to_diffusers(args.vae_checkpoint_path, vae_config)
        vae = AutoencoderKL(**vae_config)
        vae.load_state_dict(converted_vae_state_dict, strict=True)
        if dtype is not None:
            vae = vae.to(dtype=dtype)