convert_mochi_to_diffusers.py

importargparse
fromcontextlibimportnullcontext

importtorch
fromaccelerateimportinit_empty_weights
fromsafetensors.torchimportload_file
fromtransformersimportT5EncoderModel, T5Tokenizer

fromdiffusersimportAutoencoderKLMochi, FlowMatchEulerDiscreteScheduler, MochiPipeline, MochiTransformer3DModel
fromdiffusers.utils.import_utilsimportis_accelerate_available


CTX=init_empty_weightsifis_accelerate_available() elsenullcontext

TOKENIZER_MAX_LENGTH=256

parser=argparse.ArgumentParser()
parser.add_argument("--transformer_checkpoint_path", default=None, type=str)
parser.add_argument("--vae_encoder_checkpoint_path", default=None, type=str)
parser.add_argument("--vae_decoder_checkpoint_path", default=None, type=str)
parser.add_argument("--output_path", required=True, type=str)
parser.add_argument("--push_to_hub", action="store_true", default=False, help="Whether to push to HF Hub after saving")
parser.add_argument("--text_encoder_cache_dir", type=str, default=None, help="Path to text encoder cache directory")
parser.add_argument("--dtype", type=str, default=None)

args=parser.parse_args()


# This is specific to `AdaLayerNormContinuous`:
# Diffusers implementation split the linear projection into the scale, shift while Mochi split it into shift, scale
defswap_scale_shift(weight, dim):
shift, scale=weight.chunk(2, dim=0)
new_weight=torch.cat([scale, shift], dim=0)
returnnew_weight


defswap_proj_gate(weight):
proj, gate=weight.chunk(2, dim=0)
new_weight=torch.cat([gate, proj], dim=0)
returnnew_weight


defconvert_mochi_transformer_checkpoint_to_diffusers(ckpt_path):
original_state_dict=load_file(ckpt_path, device="cpu")
new_state_dict= {}

# Convert patch_embed
new_state_dict["patch_embed.proj.weight"] =original_state_dict.pop("x_embedder.proj.weight")
new_state_dict["patch_embed.proj.bias"] =original_state_dict.pop("x_embedder.proj.bias")

# Convert time_embed
new_state_dict["time_embed.timestep_embedder.linear_1.weight"] =original_state_dict.pop("t_embedder.mlp.0.weight")
new_state_dict["time_embed.timestep_embedder.linear_1.bias"] =original_state_dict.pop("t_embedder.mlp.0.bias")
new_state_dict["time_embed.timestep_embedder.linear_2.weight"] =original_state_dict.pop("t_embedder.mlp.2.weight")
new_state_dict["time_embed.timestep_embedder.linear_2.bias"] =original_state_dict.pop("t_embedder.mlp.2.bias")
new_state_dict["time_embed.pooler.to_kv.weight"] =original_state_dict.pop("t5_y_embedder.to_kv.weight")
new_state_dict["time_embed.pooler.to_kv.bias"] =original_state_dict.pop("t5_y_embedder.to_kv.bias")
new_state_dict["time_embed.pooler.to_q.weight"] =original_state_dict.pop("t5_y_embedder.to_q.weight")
new_state_dict["time_embed.pooler.to_q.bias"] =original_state_dict.pop("t5_y_embedder.to_q.bias")
new_state_dict["time_embed.pooler.to_out.weight"] =original_state_dict.pop("t5_y_embedder.to_out.weight")
new_state_dict["time_embed.pooler.to_out.bias"] =original_state_dict.pop("t5_y_embedder.to_out.bias")
new_state_dict["time_embed.caption_proj.weight"] =original_state_dict.pop("t5_yproj.weight")
new_state_dict["time_embed.caption_proj.bias"] =original_state_dict.pop("t5_yproj.bias")

# Convert transformer blocks
num_layers=48
foriinrange(num_layers):
block_prefix=f"transformer_blocks.{i}."
old_prefix=f"blocks.{i}."

# norm1
new_state_dict[block_prefix+"norm1.linear.weight"] =original_state_dict.pop(old_prefix+"mod_x.weight")
new_state_dict[block_prefix+"norm1.linear.bias"] =original_state_dict.pop(old_prefix+"mod_x.bias")
ifi<num_layers-1:
new_state_dict[block_prefix+"norm1_context.linear.weight"] =original_state_dict.pop(
old_prefix+"mod_y.weight"
 )
new_state_dict[block_prefix+"norm1_context.linear.bias"] =original_state_dict.pop(
old_prefix+"mod_y.bias"
 )
else:
new_state_dict[block_prefix+"norm1_context.linear_1.weight"] =original_state_dict.pop(
old_prefix+"mod_y.weight"
 )
new_state_dict[block_prefix+"norm1_context.linear_1.bias"] =original_state_dict.pop(
old_prefix+"mod_y.bias"
 )

# Visual attention
qkv_weight=original_state_dict.pop(old_prefix+"attn.qkv_x.weight")
q, k, v=qkv_weight.chunk(3, dim=0)

new_state_dict[block_prefix+"attn1.to_q.weight"] =q
new_state_dict[block_prefix+"attn1.to_k.weight"] =k
new_state_dict[block_prefix+"attn1.to_v.weight"] =v
new_state_dict[block_prefix+"attn1.norm_q.weight"] =original_state_dict.pop(
old_prefix+"attn.q_norm_x.weight"
 )
new_state_dict[block_prefix+"attn1.norm_k.weight"] =original_state_dict.pop(
old_prefix+"attn.k_norm_x.weight"
 )
new_state_dict[block_prefix+"attn1.to_out.0.weight"] =original_state_dict.pop(
old_prefix+"attn.proj_x.weight"
 )
new_state_dict[block_prefix+"attn1.to_out.0.bias"] =original_state_dict.pop(old_prefix+"attn.proj_x.bias")

# Context attention
qkv_weight=original_state_dict.pop(old_prefix+"attn.qkv_y.weight")
q, k, v=qkv_weight.chunk(3, dim=0)

new_state_dict[block_prefix+"attn1.add_q_proj.weight"] =q
new_state_dict[block_prefix+"attn1.add_k_proj.weight"] =k
new_state_dict[block_prefix+"attn1.add_v_proj.weight"] =v
new_state_dict[block_prefix+"attn1.norm_added_q.weight"] =original_state_dict.pop(
old_prefix+"attn.q_norm_y.weight"
 )
new_state_dict[block_prefix+"attn1.norm_added_k.weight"] =original_state_dict.pop(
old_prefix+"attn.k_norm_y.weight"
 )
ifi<num_layers-1:
new_state_dict[block_prefix+"attn1.to_add_out.weight"] =original_state_dict.pop(
old_prefix+"attn.proj_y.weight"
 )
new_state_dict[block_prefix+"attn1.to_add_out.bias"] =original_state_dict.pop(
old_prefix+"attn.proj_y.bias"
 )

# MLP
new_state_dict[block_prefix+"ff.net.0.proj.weight"] =swap_proj_gate(
original_state_dict.pop(old_prefix+"mlp_x.w1.weight")
 )
new_state_dict[block_prefix+"ff.net.2.weight"] =original_state_dict.pop(old_prefix+"mlp_x.w2.weight")
ifi<num_layers-1:
new_state_dict[block_prefix+"ff_context.net.0.proj.weight"] =swap_proj_gate(
original_state_dict.pop(old_prefix+"mlp_y.w1.weight")
 )
new_state_dict[block_prefix+"ff_context.net.2.weight"] =original_state_dict.pop(
old_prefix+"mlp_y.w2.weight"
 )

# Output layers
new_state_dict["norm_out.linear.weight"] =swap_scale_shift(
original_state_dict.pop("final_layer.mod.weight"), dim=0
 )
new_state_dict["norm_out.linear.bias"] =swap_scale_shift(original_state_dict.pop("final_layer.mod.bias"), dim=0)
new_state_dict["proj_out.weight"] =original_state_dict.pop("final_layer.linear.weight")
new_state_dict["proj_out.bias"] =original_state_dict.pop("final_layer.linear.bias")

new_state_dict["pos_frequencies"] =original_state_dict.pop("pos_frequencies")

print("Remaining Keys:", original_state_dict.keys())

returnnew_state_dict


defconvert_mochi_vae_state_dict_to_diffusers(encoder_ckpt_path, decoder_ckpt_path):
encoder_state_dict=load_file(encoder_ckpt_path, device="cpu")
decoder_state_dict=load_file(decoder_ckpt_path, device="cpu")
new_state_dict= {}

# ==== Decoder =====
prefix="decoder."

# Convert conv_in
new_state_dict[f"{prefix}conv_in.weight"] =decoder_state_dict.pop("blocks.0.0.weight")
new_state_dict[f"{prefix}conv_in.bias"] =decoder_state_dict.pop("blocks.0.0.bias")

# Convert block_in (MochiMidBlock3D)
foriinrange(3): # layers_per_block[-1] = 3
new_state_dict[f"{prefix}block_in.resnets.{i}.norm1.norm_layer.weight"] =decoder_state_dict.pop(
f"blocks.0.{i+1}.stack.0.weight"
 )
new_state_dict[f"{prefix}block_in.resnets.{i}.norm1.norm_layer.bias"] =decoder_state_dict.pop(
f"blocks.0.{i+1}.stack.0.bias"
 )
new_state_dict[f"{prefix}block_in.resnets.{i}.conv1.conv.weight"] =decoder_state_dict.pop(
f"blocks.0.{i+1}.stack.2.weight"
 )
new_state_dict[f"{prefix}block_in.resnets.{i}.conv1.conv.bias"] =decoder_state_dict.pop(
f"blocks.0.{i+1}.stack.2.bias"
 )
new_state_dict[f"{prefix}block_in.resnets.{i}.norm2.norm_layer.weight"] =decoder_state_dict.pop(
f"blocks.0.{i+1}.stack.3.weight"
 )
new_state_dict[f"{prefix}block_in.resnets.{i}.norm2.norm_layer.bias"] =decoder_state_dict.pop(
f"blocks.0.{i+1}.stack.3.bias"
 )
new_state_dict[f"{prefix}block_in.resnets.{i}.conv2.conv.weight"] =decoder_state_dict.pop(
f"blocks.0.{i+1}.stack.5.weight"
 )
new_state_dict[f"{prefix}block_in.resnets.{i}.conv2.conv.bias"] =decoder_state_dict.pop(
f"blocks.0.{i+1}.stack.5.bias"
 )

# Convert up_blocks (MochiUpBlock3D)
down_block_layers= [6, 4, 3] # layers_per_block[-2], layers_per_block[-3], layers_per_block[-4]
forblockinrange(3):
foriinrange(down_block_layers[block]):
new_state_dict[f"{prefix}up_blocks.{block}.resnets.{i}.norm1.norm_layer.weight"] =decoder_state_dict.pop(
f"blocks.{block+1}.blocks.{i}.stack.0.weight"
 )
new_state_dict[f"{prefix}up_blocks.{block}.resnets.{i}.norm1.norm_layer.bias"] =decoder_state_dict.pop(
f"blocks.{block+1}.blocks.{i}.stack.0.bias"
 )
new_state_dict[f"{prefix}up_blocks.{block}.resnets.{i}.conv1.conv.weight"] =decoder_state_dict.pop(
f"blocks.{block+1}.blocks.{i}.stack.2.weight"
 )
new_state_dict[f"{prefix}up_blocks.{block}.resnets.{i}.conv1.conv.bias"] =decoder_state_dict.pop(
f"blocks.{block+1}.blocks.{i}.stack.2.bias"
 )
new_state_dict[f"{prefix}up_blocks.{block}.resnets.{i}.norm2.norm_layer.weight"] =decoder_state_dict.pop(
f"blocks.{block+1}.blocks.{i}.stack.3.weight"
 )
new_state_dict[f"{prefix}up_blocks.{block}.resnets.{i}.norm2.norm_layer.bias"] =decoder_state_dict.pop(
f"blocks.{block+1}.blocks.{i}.stack.3.bias"
 )
new_state_dict[f"{prefix}up_blocks.{block}.resnets.{i}.conv2.conv.weight"] =decoder_state_dict.pop(
f"blocks.{block+1}.blocks.{i}.stack.5.weight"
 )
new_state_dict[f"{prefix}up_blocks.{block}.resnets.{i}.conv2.conv.bias"] =decoder_state_dict.pop(
f"blocks.{block+1}.blocks.{i}.stack.5.bias"
 )
new_state_dict[f"{prefix}up_blocks.{block}.proj.weight"] =decoder_state_dict.pop(
f"blocks.{block+1}.proj.weight"
 )
new_state_dict[f"{prefix}up_blocks.{block}.proj.bias"] =decoder_state_dict.pop(
f"blocks.{block+1}.proj.bias"
 )

# Convert block_out (MochiMidBlock3D)
foriinrange(3): # layers_per_block[0] = 3
new_state_dict[f"{prefix}block_out.resnets.{i}.norm1.norm_layer.weight"] =decoder_state_dict.pop(
f"blocks.4.{i}.stack.0.weight"
 )
new_state_dict[f"{prefix}block_out.resnets.{i}.norm1.norm_layer.bias"] =decoder_state_dict.pop(
f"blocks.4.{i}.stack.0.bias"
 )
new_state_dict[f"{prefix}block_out.resnets.{i}.conv1.conv.weight"] =decoder_state_dict.pop(
f"blocks.4.{i}.stack.2.weight"
 )
new_state_dict[f"{prefix}block_out.resnets.{i}.conv1.conv.bias"] =decoder_state_dict.pop(
f"blocks.4.{i}.stack.2.bias"
 )
new_state_dict[f"{prefix}block_out.resnets.{i}.norm2.norm_layer.weight"] =decoder_state_dict.pop(
f"blocks.4.{i}.stack.3.weight"
 )
new_state_dict[f"{prefix}block_out.resnets.{i}.norm2.norm_layer.bias"] =decoder_state_dict.pop(
f"blocks.4.{i}.stack.3.bias"
 )
new_state_dict[f"{prefix}block_out.resnets.{i}.conv2.conv.weight"] =decoder_state_dict.pop(
f"blocks.4.{i}.stack.5.weight"
 )
new_state_dict[f"{prefix}block_out.resnets.{i}.conv2.conv.bias"] =decoder_state_dict.pop(
f"blocks.4.{i}.stack.5.bias"
 )

# Convert proj_out (Conv1x1 ~= nn.Linear)
new_state_dict[f"{prefix}proj_out.weight"] =decoder_state_dict.pop("output_proj.weight")
new_state_dict[f"{prefix}proj_out.bias"] =decoder_state_dict.pop("output_proj.bias")

print("Remaining Decoder Keys:", decoder_state_dict.keys())

# ==== Encoder =====
prefix="encoder."

new_state_dict[f"{prefix}proj_in.weight"] =encoder_state_dict.pop("layers.0.weight")
new_state_dict[f"{prefix}proj_in.bias"] =encoder_state_dict.pop("layers.0.bias")

# Convert block_in (MochiMidBlock3D)
foriinrange(3): # layers_per_block[0] = 3
new_state_dict[f"{prefix}block_in.resnets.{i}.norm1.norm_layer.weight"] =encoder_state_dict.pop(
f"layers.{i+1}.stack.0.weight"
 )
new_state_dict[f"{prefix}block_in.resnets.{i}.norm1.norm_layer.bias"] =encoder_state_dict.pop(
f"layers.{i+1}.stack.0.bias"
 )
new_state_dict[f"{prefix}block_in.resnets.{i}.conv1.conv.weight"] =encoder_state_dict.pop(
f"layers.{i+1}.stack.2.weight"
 )
new_state_dict[f"{prefix}block_in.resnets.{i}.conv1.conv.bias"] =encoder_state_dict.pop(
f"layers.{i+1}.stack.2.bias"
 )
new_state_dict[f"{prefix}block_in.resnets.{i}.norm2.norm_layer.weight"] =encoder_state_dict.pop(
f"layers.{i+1}.stack.3.weight"
 )
new_state_dict[f"{prefix}block_in.resnets.{i}.norm2.norm_layer.bias"] =encoder_state_dict.pop(
f"layers.{i+1}.stack.3.bias"
 )
new_state_dict[f"{prefix}block_in.resnets.{i}.conv2.conv.weight"] =encoder_state_dict.pop(
f"layers.{i+1}.stack.5.weight"
 )
new_state_dict[f"{prefix}block_in.resnets.{i}.conv2.conv.bias"] =encoder_state_dict.pop(
f"layers.{i+1}.stack.5.bias"
 )

# Convert down_blocks (MochiDownBlock3D)
down_block_layers= [3, 4, 6] # layers_per_block[1], layers_per_block[2], layers_per_block[3]
forblockinrange(3):
new_state_dict[f"{prefix}down_blocks.{block}.conv_in.conv.weight"] =encoder_state_dict.pop(
f"layers.{block+4}.layers.0.weight"
 )
new_state_dict[f"{prefix}down_blocks.{block}.conv_in.conv.bias"] =encoder_state_dict.pop(
f"layers.{block+4}.layers.0.bias"
 )

foriinrange(down_block_layers[block]):
# Convert resnets
new_state_dict[f"{prefix}down_blocks.{block}.resnets.{i}.norm1.norm_layer.weight"] = (
encoder_state_dict.pop(f"layers.{block+4}.layers.{i+1}.stack.0.weight")
 )
new_state_dict[f"{prefix}down_blocks.{block}.resnets.{i}.norm1.norm_layer.bias"] =encoder_state_dict.pop(
f"layers.{block+4}.layers.{i+1}.stack.0.bias"
 )
new_state_dict[f"{prefix}down_blocks.{block}.resnets.{i}.conv1.conv.weight"] =encoder_state_dict.pop(
f"layers.{block+4}.layers.{i+1}.stack.2.weight"
 )
new_state_dict[f"{prefix}down_blocks.{block}.resnets.{i}.conv1.conv.bias"] =encoder_state_dict.pop(
f"layers.{block+4}.layers.{i+1}.stack.2.bias"
 )
new_state_dict[f"{prefix}down_blocks.{block}.resnets.{i}.norm2.norm_layer.weight"] = (
encoder_state_dict.pop(f"layers.{block+4}.layers.{i+1}.stack.3.weight")
 )
new_state_dict[f"{prefix}down_blocks.{block}.resnets.{i}.norm2.norm_layer.bias"] =encoder_state_dict.pop(
f"layers.{block+4}.layers.{i+1}.stack.3.bias"
 )
new_state_dict[f"{prefix}down_blocks.{block}.resnets.{i}.conv2.conv.weight"] =encoder_state_dict.pop(
f"layers.{block+4}.layers.{i+1}.stack.5.weight"
 )
new_state_dict[f"{prefix}down_blocks.{block}.resnets.{i}.conv2.conv.bias"] =encoder_state_dict.pop(
f"layers.{block+4}.layers.{i+1}.stack.5.bias"
 )

# Convert attentions
qkv_weight=encoder_state_dict.pop(f"layers.{block+4}.layers.{i+1}.attn_block.attn.qkv.weight")
q, k, v=qkv_weight.chunk(3, dim=0)

new_state_dict[f"{prefix}down_blocks.{block}.attentions.{i}.to_q.weight"] =q
new_state_dict[f"{prefix}down_blocks.{block}.attentions.{i}.to_k.weight"] =k
new_state_dict[f"{prefix}down_blocks.{block}.attentions.{i}.to_v.weight"] =v
new_state_dict[f"{prefix}down_blocks.{block}.attentions.{i}.to_out.0.weight"] =encoder_state_dict.pop(
f"layers.{block+4}.layers.{i+1}.attn_block.attn.out.weight"
 )
new_state_dict[f"{prefix}down_blocks.{block}.attentions.{i}.to_out.0.bias"] =encoder_state_dict.pop(
f"layers.{block+4}.layers.{i+1}.attn_block.attn.out.bias"
 )
new_state_dict[f"{prefix}down_blocks.{block}.norms.{i}.norm_layer.weight"] =encoder_state_dict.pop(
f"layers.{block+4}.layers.{i+1}.attn_block.norm.weight"
 )
new_state_dict[f"{prefix}down_blocks.{block}.norms.{i}.norm_layer.bias"] =encoder_state_dict.pop(
f"layers.{block+4}.layers.{i+1}.attn_block.norm.bias"
 )

# Convert block_out (MochiMidBlock3D)
foriinrange(3): # layers_per_block[-1] = 3
# Convert resnets
new_state_dict[f"{prefix}block_out.resnets.{i}.norm1.norm_layer.weight"] =encoder_state_dict.pop(
f"layers.{i+7}.stack.0.weight"
 )
new_state_dict[f"{prefix}block_out.resnets.{i}.norm1.norm_layer.bias"] =encoder_state_dict.pop(
f"layers.{i+7}.stack.0.bias"
 )
new_state_dict[f"{prefix}block_out.resnets.{i}.conv1.conv.weight"] =encoder_state_dict.pop(
f"layers.{i+7}.stack.2.weight"
 )
new_state_dict[f"{prefix}block_out.resnets.{i}.conv1.conv.bias"] =encoder_state_dict.pop(
f"layers.{i+7}.stack.2.bias"
 )
new_state_dict[f"{prefix}block_out.resnets.{i}.norm2.norm_layer.weight"] =encoder_state_dict.pop(
f"layers.{i+7}.stack.3.weight"
 )
new_state_dict[f"{prefix}block_out.resnets.{i}.norm2.norm_layer.bias"] =encoder_state_dict.pop(
f"layers.{i+7}.stack.3.bias"
 )
new_state_dict[f"{prefix}block_out.resnets.{i}.conv2.conv.weight"] =encoder_state_dict.pop(
f"layers.{i+7}.stack.5.weight"
 )
new_state_dict[f"{prefix}block_out.resnets.{i}.conv2.conv.bias"] =encoder_state_dict.pop(
f"layers.{i+7}.stack.5.bias"
 )

# Convert attentions
qkv_weight=encoder_state_dict.pop(f"layers.{i+7}.attn_block.attn.qkv.weight")
q, k, v=qkv_weight.chunk(3, dim=0)

new_state_dict[f"{prefix}block_out.attentions.{i}.to_q.weight"] =q
new_state_dict[f"{prefix}block_out.attentions.{i}.to_k.weight"] =k
new_state_dict[f"{prefix}block_out.attentions.{i}.to_v.weight"] =v
new_state_dict[f"{prefix}block_out.attentions.{i}.to_out.0.weight"] =encoder_state_dict.pop(
f"layers.{i+7}.attn_block.attn.out.weight"
 )
new_state_dict[f"{prefix}block_out.attentions.{i}.to_out.0.bias"] =encoder_state_dict.pop(
f"layers.{i+7}.attn_block.attn.out.bias"
 )
new_state_dict[f"{prefix}block_out.norms.{i}.norm_layer.weight"] =encoder_state_dict.pop(
f"layers.{i+7}.attn_block.norm.weight"
 )
new_state_dict[f"{prefix}block_out.norms.{i}.norm_layer.bias"] =encoder_state_dict.pop(
f"layers.{i+7}.attn_block.norm.bias"
 )

# Convert output layers
new_state_dict[f"{prefix}norm_out.norm_layer.weight"] =encoder_state_dict.pop("output_norm.weight")
new_state_dict[f"{prefix}norm_out.norm_layer.bias"] =encoder_state_dict.pop("output_norm.bias")
new_state_dict[f"{prefix}proj_out.weight"] =encoder_state_dict.pop("output_proj.weight")

print("Remaining Encoder Keys:", encoder_state_dict.keys())

returnnew_state_dict


defmain(args):
ifargs.dtypeisNone:
dtype=None
ifargs.dtype=="fp16":
dtype=torch.float16
elifargs.dtype=="bf16":
dtype=torch.bfloat16
elifargs.dtype=="fp32":
dtype=torch.float32
else:
raiseValueError(f"Unsupported dtype: {args.dtype}")

transformer=None
vae=None

ifargs.transformer_checkpoint_pathisnotNone:
converted_transformer_state_dict=convert_mochi_transformer_checkpoint_to_diffusers(
args.transformer_checkpoint_path
 )
transformer=MochiTransformer3DModel()
transformer.load_state_dict(converted_transformer_state_dict, strict=True)
ifdtypeisnotNone:
transformer=transformer.to(dtype=dtype)

ifargs.vae_encoder_checkpoint_pathisnotNoneandargs.vae_decoder_checkpoint_pathisnotNone:
vae=AutoencoderKLMochi(latent_channels=12, out_channels=3)
converted_vae_state_dict=convert_mochi_vae_state_dict_to_diffusers(
args.vae_encoder_checkpoint_path, args.vae_decoder_checkpoint_path
 )
vae.load_state_dict(converted_vae_state_dict, strict=True)
ifdtypeisnotNone:
vae=vae.to(dtype=dtype)

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)

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

pipe=MochiPipeline(
scheduler=FlowMatchEulerDiscreteScheduler(invert_sigmas=True),
vae=vae,
text_encoder=text_encoder,
tokenizer=tokenizer,
transformer=transformer,
 )
pipe.save_pretrained(args.output_path, safe_serialization=True, max_shard_size="5GB", push_to_hub=args.push_to_hub)


if__name__=="__main__":
main(args)