Skip to content

Add MPS (Apple Silicon) support #155

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
imperatormk wants to merge 1 commit into
base: main
Choose a base branch
from
Open

Add MPS (Apple Silicon) support #155

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
63 changes: 51 additions & 12 deletions demos/cli.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -7,6 +7,14 @@
import numpy as np
import torch

try:
from mps_conv3d import patch_conv3d, is_available as conv3d_available
if conv3d_available():
patch_conv3d()
print("✓ MPS Conv3D patched")
except ImportError:
pass

from genmo.lib.progress import progress_bar
from genmo.lib.utils import save_video
from genmo.mochi_preview.pipelines import (
Expand All @@ -15,35 +23,65 @@
MochiMultiGPUPipeline,
MochiSingleGPUPipeline,
T5ModelFactory,
get_device,
linear_quadratic_schedule,
)

pipeline = None
model_dir_path = None
lora_path = None
num_gpus = torch.cuda.device_count()
quantize_nf4 = False
attention_mode = None
# Check for available GPUs - works with CUDA, falls back to MPS/CPU
if torch.cuda.is_available():
num_gpus = torch.cuda.device_count()
elif torch.backends.mps.is_available():
num_gpus = 1 # MPS is single-device
print("Using MPS (Apple Silicon)")
else:
num_gpus = 0 # CPU mode
print("No GPU detected, using CPU")
cpu_offload = False


def configure_model(model_dir_path_, lora_path_, cpu_offload_, fast_model_=False):
global model_dir_path, lora_path, cpu_offload
def configure_model(model_dir_path_, lora_path_, cpu_offload_, quantize_nf4_=False, attention_mode_=None):
global model_dir_path, lora_path, cpu_offload, quantize_nf4, attention_mode
model_dir_path = model_dir_path_
lora_path = lora_path_
cpu_offload = cpu_offload_
quantize_nf4 = quantize_nf4_
attention_mode = attention_mode_


def load_model():
global num_gpus, pipeline, model_dir_path, lora_path
global num_gpus, pipeline, model_dir_path, lora_path, quantize_nf4, attention_mode
if pipeline is None:
MOCHI_DIR = model_dir_path
print(f"Launching with {num_gpus} GPUs. If you want to force single GPU mode use CUDA_VISIBLE_DEVICES=0.")
klass = MochiSingleGPUPipeline if num_gpus == 1 else MochiMultiGPUPipeline
device = get_device()
if device.type == "cuda":
print(f"Launching with {num_gpus} GPUs. If you want to force single GPU mode use CUDA_VISIBLE_DEVICES=0.")
elif device.type == "mps":
print("Launching on MPS (Apple Silicon)")
if quantize_nf4:
print("NF4 quantization enabled (~5GB model size)")
else:
print("Launching on CPU")
# Multi-GPU only supported on CUDA
klass = MochiSingleGPUPipeline if (num_gpus <= 1 or device.type != "cuda") else MochiMultiGPUPipeline
# Check for local T5 weights, else use HuggingFace
t5_local = f"{MOCHI_DIR}/../t5"
t5_path = t5_local if os.path.exists(t5_local) else None
if t5_path:
print(f"Using local T5: {t5_path}")

kwargs = dict(
text_encoder_factory=T5ModelFactory(),
text_encoder_factory=T5ModelFactory(model_dir=t5_path),
dit_factory=DitModelFactory(
model_path=f"{MOCHI_DIR}/dit.safetensors",
lora_path=lora_path,
model_dtype="bf16",
quantize_nf4=quantize_nf4,
attention_mode=attention_mode,
),
decoder_factory=DecoderModelFactory(
model_path=f"{MOCHI_DIR}/decoder.safetensors",
Expand Down Expand Up @@ -93,8 +131,7 @@ def generate_video(
"sigma_schedule": sigma_schedule,
"cfg_schedule": cfg_schedule,
"num_inference_steps": num_inference_steps,
# We *need* flash attention to batch cfg
# and it's only worth doing in a high-memory regime (assume multiple GPUs)
# Batched CFG requires flash attention (B=2 not supported by SDPA path)
"batch_cfg": False,
"prompt": prompt,
"negative_prompt": negative_prompt,
Expand Down Expand Up @@ -145,14 +182,16 @@ def generate_video(
@click.option("--model_dir", required=True, help="Path to the model directory.")
@click.option("--lora_path", required=False, help="Path to the lora file.")
@click.option("--cpu_offload", is_flag=True, help="Whether to offload model to CPU")
@click.option("--quantize", is_flag=True, help="Use NF4 quantization (4-bit) for MPS - reduces memory from ~20GB to ~5GB")
@click.option("--attention-mode", type=click.Choice(["flash", "mps_flash", "sdpa", "sage"]), default=None, help="Attention mode (auto-detected if not specified)")
@click.option("--out_dir", default="outputs", help="Output directory for generated videos")
@click.option("--threshold-noise", default=0.025, help="threshold noise")
@click.option("--linear-steps", default=None, type=int, help="linear steps")
def generate_cli(
prompt, sweep_file, negative_prompt, width, height, num_frames, seed, cfg_scale, num_steps,
model_dir, lora_path, cpu_offload, out_dir, threshold_noise, linear_steps
prompt, sweep_file, negative_prompt, width, height, num_frames, seed, cfg_scale, num_steps,
model_dir, lora_path, cpu_offload, quantize, attention_mode, out_dir, threshold_noise, linear_steps
):
configure_model(model_dir, lora_path, cpu_offload)
configure_model(model_dir, lora_path, cpu_offload, quantize, attention_mode)

if sweep_file:
with open(sweep_file, 'r') as f:
Expand Down
103 changes: 103 additions & 0 deletions scripts/quantize_diffusers_transformer.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,103 @@
#!/usr/bin/env python
"""
Quantize Diffusers MochiTransformer3DModel to NF4 and save to disk.

Usage:
python scripts/quantize_diffusers_transformer.py /path/to/transformer

Output:
/path/to/transformer_nf4.pt (~5GB instead of ~20GB)

The quantized model can be loaded with:
model = torch.load("transformer_nf4.pt", weights_only=False)
"""
import argparse
import os
import sys
import torch
import gc


def main():
parser = argparse.ArgumentParser(description="Quantize Diffusers MochiTransformer3DModel to NF4")
parser.add_argument("input", help="Path to transformer directory (with model files)")
parser.add_argument("--output", "-o", help="Output path (default: input_nf4.pt)")
parser.add_argument("--dtype", choices=["bf16", "fp16"], default="bf16",
help="Compute dtype (default: bf16, use fp16 for faster MPS)")
parser.add_argument("--force", "-f", action="store_true",
help="Overwrite existing output file")
args = parser.parse_args()

input_path = args.input
dtype_suffix = "_nf4_fp16.pt" if args.dtype == "fp16" else "_nf4.pt"
output_path = args.output or os.path.join(os.path.dirname(input_path), f"transformer{dtype_suffix}")
compute_dtype = torch.float16 if args.dtype == "fp16" else torch.bfloat16

if not os.path.exists(input_path):
print(f"Error: Input path not found: {input_path}")
sys.exit(1)

if os.path.exists(output_path) and not args.force:
print(f"Output already exists: {output_path}")
print("Use --force to overwrite")
sys.exit(0)

print(f"Input: {input_path}")
print(f"Output: {output_path}")
print(f"Dtype: {args.dtype}")
print()

from diffusers import MochiTransformer3DModel
from mps_bitsandbytes import quantize_model, BitsAndBytesConfig

print("Loading Diffusers MochiTransformer3DModel...")
transformer = MochiTransformer3DModel.from_pretrained(
input_path,
variant="bf16",
torch_dtype=torch.bfloat16,
local_files_only=True,
)
num_params = sum(p.numel() for p in transformer.parameters())
print(f" Loaded: {num_params:,} params")

# Check size before quantization
param_bytes = sum(p.numel() * p.element_size() for p in transformer.parameters())
print(f" Size before: {param_bytes / 1e9:.2f} GB")

print()
print("Quantizing to NF4...")
config = BitsAndBytesConfig(
load_in_4bit=True,
bnb_4bit_quant_type="nf4",
bnb_4bit_compute_dtype=compute_dtype,
)

transformer_q = quantize_model(transformer, quantization_config=config)

# Check size after
total_size = 0
for param in transformer_q.parameters():
total_size += param.numel() * param.element_size()
for buf in transformer_q.buffers():
total_size += buf.numel() * buf.element_size()
print(f" Size after: {total_size / 1e9:.2f} GB")
print(f" Savings: {(1 - total_size / param_bytes) * 100:.1f}%")

print()
print(f"Saving to {output_path}...")
torch.save(transformer_q, output_path)

# Verify file size
file_size = os.path.getsize(output_path)
print(f" File size: {file_size / 1e9:.2f} GB")

print()
print("Done!")
print()
print("To use:")
print(f" model = torch.load('{output_path}', weights_only=False)")
print(f" model = model.to('mps') # or 'cuda'")


if __name__ == "__main__":
main()
127 changes: 127 additions & 0 deletions scripts/quantize_dit.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,127 @@
#!/usr/bin/env python
"""
Quantize Mochi DiT (Genmo) weights to NF4 and save to disk.

Usage:
python scripts/quantize_dit.py /path/to/weights/dit.safetensors

Output:
/path/to/weights/dit_nf4.pt (~5GB instead of ~20GB)

The quantized model can be loaded with:
model = torch.load("dit_nf4.pt", weights_only=False)
"""
import argparse
import sys
import os

# Add src to path
sys.path.insert(0, os.path.join(os.path.dirname(__file__), '..', 'src'))

import torch
from safetensors.torch import load_file


def main():
parser = argparse.ArgumentParser(description="Quantize Mochi DiT to NF4")
parser.add_argument("input", help="Path to dit.safetensors")
parser.add_argument("--output", "-o", help="Output path (default: input_nf4.pt)")
parser.add_argument("--dtype", choices=["bf16", "fp16"], default="bf16",
help="Compute dtype (default: bf16, use fp16 for faster MPS)")
parser.add_argument("--force", "-f", action="store_true",
help="Overwrite existing output file")
args = parser.parse_args()

input_path = args.input
dtype_suffix = "_nf4_fp16.pt" if args.dtype == "fp16" else "_nf4.pt"
output_path = args.output or input_path.replace(".safetensors", dtype_suffix)
compute_dtype = torch.float16 if args.dtype == "fp16" else torch.bfloat16

if not os.path.exists(input_path):
print(f"Error: Input file not found: {input_path}")
sys.exit(1)

if os.path.exists(output_path) and not args.force:
print(f"Output already exists: {output_path}")
print("Use --force to overwrite")
sys.exit(0)

print(f"Input: {input_path}")
print(f"Output: {output_path}")
print(f"Dtype: {args.dtype}")
print()

# Import after path setup
from mps_bitsandbytes import quantize_model, BitsAndBytesConfig
from mps_bitsandbytes.integration import get_memory_footprint
from genmo.mochi_preview.dit.joint_model.asymm_models_joint import AsymmDiTJoint

print("Creating model skeleton...")
model = torch.nn.utils.skip_init(AsymmDiTJoint,
depth=48,
patch_size=2,
num_heads=24,
hidden_size_x=3072,
hidden_size_y=1536,
mlp_ratio_x=4.0,
mlp_ratio_y=4.0,
in_channels=12,
qk_norm=True,
qkv_bias=False,
out_bias=True,
patch_embed_bias=True,
timestep_mlp_bias=True,
timestep_scale=1000.0,
t5_feat_dim=4096,
t5_token_length=256,
rope_theta=10000.0,
attention_mode="sdpa", # Will be overridden at runtime based on device
)

print("Loading weights from safetensors...")
sd = load_file(input_path)
model.load_state_dict(sd)
del sd # Free memory

print()
print("Quantizing to NF4...")
config = BitsAndBytesConfig(
load_in_4bit=True,
bnb_4bit_quant_type="nf4",
bnb_4bit_compute_dtype=compute_dtype,
)

# Skip patch embed conv and normalization layers (keep full precision)
modules_to_skip = ["x_embedder.proj", "norm", "final_layer"]

model = quantize_model(
model,
quantization_config=config,
modules_to_not_convert=modules_to_skip,
)

# Report stats
stats = get_memory_footprint(model)
print(f" Original (fp16): {stats['fp16_size_gb']:.2f} GB")
print(f" Quantized (nf4): {stats['actual_size_gb']:.2f} GB")
print(f" Savings: {stats['savings_pct']:.1f}%")
print()

print(f"Saving to {output_path}...")
# Save entire model (not just state_dict) so we skip skeleton creation on load
torch.save(model, output_path)

# Check file size
size_gb = os.path.getsize(output_path) / (1024**3)
print(f" File size: {size_gb:.2f} GB")
print()
print("Done!")
print()
print("To use with Genmo pipeline:")
print(f" model = torch.load('{output_path}', weights_only=False)")
print()
print("Or run demos/cli.py with --quantize flag (auto-detects cached weights)")


if __name__ == "__main__":
main()
Loading