Mochi 1 Transformer 42

🚀 Distilled Mochi Transformer

This repository houses a distilled transformer for GenmoAI's Mochi-1. Compared to the original transformer with 48 blocks, this one consists of 42 blocks, offering a more streamlined option.

✨ Features

• Optimized Architecture: The distilled transformer reduces the number of blocks from 48 to 42, potentially improving efficiency.
• Iterative Block Dropping: Blocks are dropped based on the minimum Mean Squared Error (MSE) of the latent after each block.
• Neighboring Block Training: After dropping a block, the neighboring blocks are trained for 1K steps to maintain performance.

📚 Documentation

Training details

We analyzed the MSE of the latent after each block and iteratively dropped the blocks with the minimum MSE value.

After each block drop, we trained the neighboring blocks (one before and one after the deleted block) for 1K steps.

🚀 Try it here: Interactive Demo

💻 Usage Examples

Basic Usage

import torch
from diffusers import MochiPipeline, MochiTransformer3DModel
from diffusers.utils import export_to_video

transformer = MochiTransformer3DModel.from_pretrained(
    "NimVideo/mochi-1-transformer-42",
    torch_dtype=torch.bfloat16,
)
pipe = MochiPipeline.from_pretrained(
    "genmo/mochi-1-preview", 
    transformer=transformer,
    variant="bf16", 
    torch_dtype=torch.bfloat16
)

pipe.enable_model_cpu_offload()
pipe.enable_vae_tiling()

prompt = "Close-up of a chameleon's eye, with its scaly skin changing color. Ultra high resolution 4k."
frames = pipe(prompt, num_frames=85).frames[0]

export_to_video(frames, "mochi.mp4", fps=30)

📄 License

This project is licensed under the Apache 2.0 License.

Acknowledgements

Original code and models mochi.

Contacts

Issues should be raised directly in the repository.