def main(args):
    """
    Main function to convert CogView4 checkpoints to Diffusers format.

    Args:
        args (argparse.Namespace): Parsed command-line arguments.
    """
    # Determine the desired data type
    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

    # Convert Transformer checkpoint if provided
    if args.transformer_checkpoint_path is not None:
        converted_transformer_state_dict = convert_megatron_transformer_checkpoint_to_diffusers(
            ckpt_path=args.transformer_checkpoint_path,
            num_layers=args.num_layers,
            num_heads=args.num_heads,
            hidden_size=args.hidden_size,
        )
        transformer = CogView4Transformer2DModel(
            patch_size=2,
            in_channels=32 if args.control else 16,
            num_layers=args.num_layers,
            attention_head_dim=args.attention_head_dim,
            num_attention_heads=args.num_heads,
            out_channels=16,
            text_embed_dim=args.hidden_size,
            time_embed_dim=args.time_embed_dim,
            condition_dim=args.condition_dim,
            pos_embed_max_size=args.pos_embed_max_size,
        )

        transformer.load_state_dict(converted_transformer_state_dict, strict=True)

        # Convert to the specified dtype
        if dtype is not None:
            transformer = transformer.to(dtype=dtype)

    # Convert VAE checkpoint if provided
    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)

    # Load the text encoder and tokenizer
    text_encoder_id = "THUDM/glm-4-9b-hf"
    tokenizer = PreTrainedTokenizerFast.from_pretrained(text_encoder_id)
    text_encoder = GlmModel.from_pretrained(
        text_encoder_id,
        cache_dir=args.text_encoder_cache_dir,
        torch_dtype=torch.bfloat16 if args.dtype == "bf16" else torch.float32,
    )
    for param in text_encoder.parameters():
        param.data = param.data.contiguous()

    # Initialize the scheduler
    scheduler = FlowMatchEulerDiscreteScheduler(
        base_shift=0.25, max_shift=0.75, base_image_seq_len=256, use_dynamic_shifting=True, time_shift_type="linear"
    )

    # Create the pipeline
    if args.control:
        pipe = CogView4ControlPipeline(
            tokenizer=tokenizer,
            text_encoder=text_encoder,
            vae=vae,
            transformer=transformer,
            scheduler=scheduler,
        )
    else:
        pipe = CogView4Pipeline(
            tokenizer=tokenizer,
            text_encoder=text_encoder,
            vae=vae,
            transformer=transformer,
            scheduler=scheduler,
        )

    # Save the converted pipeline
    pipe.save_pretrained(
        args.output_path,
        safe_serialization=True,
        max_shard_size="5GB",
        push_to_hub=args.push_to_hub,
    )

    Main function to convert CogView4 checkpoints to Diffusers format.

    Args:
        args (argparse.Namespace): Parsed command-line arguments.
    
scripts/convert_cogview4_to_diffusers_megatron.py