diff --git a/demos/cli.py b/demos/cli.py index 7e8a213..a64405c 100755 --- a/demos/cli.py +++ b/demos/cli.py @@ -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 ( @@ -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", @@ -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, @@ -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: diff --git a/scripts/quantize_diffusers_transformer.py b/scripts/quantize_diffusers_transformer.py new file mode 100644 index 0000000..f2a1f16 --- /dev/null +++ b/scripts/quantize_diffusers_transformer.py @@ -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() diff --git a/scripts/quantize_dit.py b/scripts/quantize_dit.py new file mode 100644 index 0000000..efd69e2 --- /dev/null +++ b/scripts/quantize_dit.py @@ -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() diff --git a/src/genmo/lib/attn_imports.py b/src/genmo/lib/attn_imports.py index 4d4a8c1..2cefa66 100644 --- a/src/genmo/lib/attn_imports.py +++ b/src/genmo/lib/attn_imports.py @@ -3,6 +3,34 @@ import torch +# ============================================================================ +# MPS Flash Attention (Apple Silicon) +# ============================================================================ +try: + from mps_flash_attn import flash_attention as mps_flash_attn + from mps_flash_attn import is_available as mps_flash_available + HAS_MPS_FLASH = mps_flash_available() + try: + from mps_flash_attn import quantize_kv_nf4 as mps_quantize_kv_nf4 + from mps_flash_attn import flash_attention_nf4 as mps_flash_attn_nf4 + HAS_MPS_FLASH_NF4 = True + except ImportError: + mps_quantize_kv_nf4 = None + mps_flash_attn_nf4 = None + HAS_MPS_FLASH_NF4 = False + if HAS_MPS_FLASH: + nf4_status = " (with NF4 KV cache)" if HAS_MPS_FLASH_NF4 else "" + print(f"✓ MPS Flash Attention available{nf4_status}") +except ImportError: + mps_flash_attn = None + mps_quantize_kv_nf4 = None + mps_flash_attn_nf4 = None + HAS_MPS_FLASH = False + HAS_MPS_FLASH_NF4 = False + +# ============================================================================ +# CUDA Flash Attention +# ============================================================================ try: from flash_attn import flash_attn_varlen_func as flash_varlen_attn except ImportError: @@ -15,13 +43,20 @@ from torch.nn.attention import SDPBackend, sdpa_kernel +# Device-agnostic backend selection training_backends = [SDPBackend.FLASH_ATTENTION, SDPBackend.EFFICIENT_ATTENTION] eval_backends = list(training_backends) -if torch.cuda.get_device_properties(0).major >= 9.0: - # Enable fast CuDNN attention on Hopper. - # This gives NaN on the backward pass for some reason, - # so only use it for evaluation. - eval_backends.append(SDPBackend.CUDNN_ATTENTION) + +# Only check CUDA properties if CUDA is available +if torch.cuda.is_available(): + try: + if torch.cuda.get_device_properties(0).major >= 9.0: + # Enable fast CuDNN attention on Hopper. + # This gives NaN on the backward pass for some reason, + # so only use it for evaluation. + eval_backends.append(SDPBackend.CUDNN_ATTENTION) + except Exception: + pass # No CUDA device available @contextmanager def sdpa_attn_ctx(training: bool = False): diff --git a/src/genmo/lib/utils.py b/src/genmo/lib/utils.py index b8ebbd6..1fb95a3 100644 --- a/src/genmo/lib/utils.py +++ b/src/genmo/lib/utils.py @@ -4,7 +4,10 @@ import time import numpy as np -from moviepy.editor import ImageSequenceClip +try: + from moviepy.editor import ImageSequenceClip # moviepy 1.x +except ImportError: + from moviepy import ImageSequenceClip # moviepy 2.x from PIL import Image from genmo.lib.progress import get_new_progress_bar @@ -46,6 +49,27 @@ def save_video(final_frames, output_path, fps=30): assert final_frames.ndim == 4 and final_frames.shape[3] == 3, f"invalid shape: {final_frames} (need t h w c)" if final_frames.dtype != np.uint8: final_frames = (final_frames * 255).astype(np.uint8) + + # Fast path: use ffmpeg directly if available + try: + import shutil + if shutil.which('ffmpeg'): + with tempfile.TemporaryDirectory() as tmpdir: + # Save frames as images + for i, frame in enumerate(final_frames): + Image.fromarray(frame).save(os.path.join(tmpdir, f'{i:04d}.png')) + # Encode with ffmpeg + subprocess.run([ + 'ffmpeg', '-y', '-framerate', str(fps), + '-i', os.path.join(tmpdir, '%04d.png'), + '-c:v', 'libx264', '-pix_fmt', 'yuv420p', + '-crf', '18', output_path + ], capture_output=True) + return + except Exception: + pass + + # Fallback: moviepy ImageSequenceClip(list(final_frames), fps=fps).write_videofile(output_path) diff --git a/src/genmo/mochi_preview/dit/joint_model/asymm_models_joint.py b/src/genmo/mochi_preview/dit/joint_model/asymm_models_joint.py index b7a482c..ea46206 100644 --- a/src/genmo/mochi_preview/dit/joint_model/asymm_models_joint.py +++ b/src/genmo/mochi_preview/dit/joint_model/asymm_models_joint.py @@ -142,7 +142,7 @@ def prepare_qkv( # Process visual features x = modulated_rmsnorm(x, scale_x) # (B, M, dim_x) where M = N / cp_group_size qkv_x = self.qkv_x(x) # (B, M, 3 * dim_x) - assert qkv_x.dtype == torch.bfloat16 + assert qkv_x.dtype in (torch.bfloat16, torch.float16), f"Expected bf16/fp16, got {qkv_x.dtype}" qkv_x = cp.all_to_all_collect_tokens(qkv_x, self.num_heads) # (3, B, N, local_h, head_dim) @@ -211,6 +211,17 @@ def sdpa_attention(self, q, k, v): def sage_attention(self, q, k, v): return sage_attn(q, k, v, attn_mask=None, dropout_p=0.0, is_causal=False) + def mps_flash_attention(self, q, k, v): + """MPS Flash Attention using mps-flash-attn library.""" + try: + from mps_flash_attn import flash_attention + # mps_flash_attn expects (B, H, S, D) format + out = flash_attention(q, k, v, scale=self.softmax_scale) + return out + except ImportError: + # Fallback to SDPA if mps-flash-attn not available + return self.sdpa_attention(q, k, v) + def run_attention( self, q: torch.Tensor, # (total <= B * (N + L), num_heads, head_dim) @@ -244,6 +255,8 @@ def run_attention( if self.attention_mode == "sdpa": out = self.sdpa_attention(q, k, v) # (B, local_heads, seq_len, head_dim) + elif self.attention_mode == "mps_flash": + out = self.mps_flash_attention(q, k, v) # (B, local_heads, seq_len, head_dim) elif self.attention_mode == "sage": out = self.sage_attention(q, k, v) # (B, local_heads, seq_len, head_dim) else: @@ -509,11 +522,25 @@ def __init__( self.linear = nn.Linear(hidden_size, patch_size * patch_size * out_channels, device=device) def forward(self, x, c): + # MPS requires input dtype to match weight dtype + # mod/linear are float32, input may be bf16/fp16 + input_dtype = x.dtype + weight_dtype = self.mod.weight.dtype + c = F.silu(c) + if c.dtype != weight_dtype: + c = c.to(weight_dtype) shift, scale = self.mod(c).chunk(2, dim=1) - x = modulate(self.norm_final(x), shift, scale) + + x_norm = self.norm_final(x) + if x_norm.dtype != weight_dtype: + x_norm = x_norm.to(weight_dtype) + if shift.dtype != x_norm.dtype: + shift = shift.to(x_norm.dtype) + scale = scale.to(x_norm.dtype) + x = modulate(x_norm, shift, scale) x = self.linear(x) - return x + return x.to(input_dtype) class AsymmDiTJoint(nn.Module): @@ -611,7 +638,13 @@ def embed_x(self, x: torch.Tensor) -> torch.Tensor: Returns: x: (B, C=3072, N) tensor of visual tokens with positional embedding. """ - return self.x_embedder(x) # Convert BcTHW to BCN + # Conv2d is float32, input may be bf16/fp16 - cast to match, then cast back + input_dtype = x.dtype + conv_dtype = self.x_embedder.proj.weight.dtype + if x.dtype != conv_dtype: + x = x.to(conv_dtype) + x = self.x_embedder(x) # Convert BcTHW to BCN + return x.to(input_dtype) @torch.compile(disable=not COMPILE_MMDIT_BLOCK) def prepare( diff --git a/src/genmo/mochi_preview/dit/joint_model/rope_mixed.py b/src/genmo/mochi_preview/dit/joint_model/rope_mixed.py index f2952bd..607f133 100644 --- a/src/genmo/mochi_preview/dit/joint_model/rope_mixed.py +++ b/src/genmo/mochi_preview/dit/joint_model/rope_mixed.py @@ -80,9 +80,11 @@ def compute_mixed_rotation( freqs_cos: [N, num_heads, num_freqs] - cosine components freqs_sin: [N, num_heads, num_freqs] - sine components """ - with torch.autocast("cuda", enabled=False): - assert freqs.ndim == 3 - freqs_sum = torch.einsum("Nd,dhf->Nhf", pos.to(freqs), freqs) - freqs_cos = torch.cos(freqs_sum) - freqs_sin = torch.sin(freqs_sum) - return freqs_cos, freqs_sin + # Compute in float32 for precision (MPS needs this, autocast doesn't work on MPS) + assert freqs.ndim == 3 + pos_f32 = pos.to(torch.float32) + freqs_f32 = freqs.to(torch.float32) + freqs_sum = torch.einsum("Nd,dhf->Nhf", pos_f32, freqs_f32) + freqs_cos = torch.cos(freqs_sum) + freqs_sin = torch.sin(freqs_sum) + return freqs_cos.to(freqs.dtype), freqs_sin.to(freqs.dtype) diff --git a/src/genmo/mochi_preview/dit/joint_model/temporal_rope.py b/src/genmo/mochi_preview/dit/joint_model/temporal_rope.py index a8276db..1c38237 100644 --- a/src/genmo/mochi_preview/dit/joint_model/temporal_rope.py +++ b/src/genmo/mochi_preview/dit/joint_model/temporal_rope.py @@ -19,16 +19,17 @@ def apply_rotary_emb_qk_real( Returns: torch.Tensor: The input tensor with rotary embeddings applied. """ - assert xqk.dtype == torch.bfloat16 - # Split the last dimension into even and odd parts - xqk_even = xqk[..., 0::2] - xqk_odd = xqk[..., 1::2] + # Do RoPE math in float32 like Diffusers does (bf16 loses precision on MPS) + orig_dtype = xqk.dtype + xqk_even = xqk[..., 0::2].float() + xqk_odd = xqk[..., 1::2].float() + freqs_cos = freqs_cos.float() + freqs_sin = freqs_sin.float() - # Apply rotation - cos_part = (xqk_even * freqs_cos - xqk_odd * freqs_sin).type_as(xqk) - sin_part = (xqk_even * freqs_sin + xqk_odd * freqs_cos).type_as(xqk) + # Apply rotation in float32 + cos_part = xqk_even * freqs_cos - xqk_odd * freqs_sin + sin_part = xqk_even * freqs_sin + xqk_odd * freqs_cos - # Interleave the results back into the original shape - out = torch.stack([cos_part, sin_part], dim=-1).flatten(-2) - assert out.dtype == torch.bfloat16 + # Interleave and cast back to original dtype + out = torch.stack([cos_part, sin_part], dim=-1).flatten(-2).to(orig_dtype) return out diff --git a/src/genmo/mochi_preview/pipelines.py b/src/genmo/mochi_preview/pipelines.py index d9ec1d8..509ade1 100644 --- a/src/genmo/mochi_preview/pipelines.py +++ b/src/genmo/mochi_preview/pipelines.py @@ -7,7 +7,12 @@ from typing import Any, Dict, List, Literal, Optional, Union, cast import numpy as np -import ray +try: + import ray + HAS_RAY = True +except ImportError: + ray = None + HAS_RAY = False import torch import torch.distributed as dist import torch.nn as nn @@ -41,6 +46,20 @@ ) from genmo.mochi_preview.vae.vae_stats import dit_latents_to_vae_latents +# ============================================================================ +# MPS NF4 Quantization Support +# ============================================================================ +try: + from mps_bitsandbytes import quantize_model, BitsAndBytesConfig + from mps_bitsandbytes.nn import Linear4bit + HAS_MPS_QUANTIZATION = True + print("✓ MPS NF4 quantization available") +except ImportError: + HAS_MPS_QUANTIZATION = False + quantize_model = None + BitsAndBytesConfig = None + Linear4bit = None + def load_to_cpu(p, weights_only=True): if p.endswith(".safetensors"): @@ -50,6 +69,109 @@ def load_to_cpu(p, weights_only=True): return torch.load(p, map_location="cpu", weights_only=weights_only) +# ============================================================================ +# MPS / Apple Silicon Support +# ============================================================================ +def get_device(): + """Get the best available device (CUDA > MPS > CPU).""" + if torch.cuda.is_available(): + return torch.device("cuda:0") + elif torch.backends.mps.is_available(): + return torch.device("mps") + else: + return torch.device("cpu") + + +def get_autocast_context(device, dtype=torch.bfloat16): + """Get autocast context for the given device.""" + if device.type == "cuda": + return torch.autocast("cuda", dtype=dtype) + elif device.type == "mps": + return torch.autocast("mps", dtype=dtype) + else: + # CPU doesn't support bfloat16 autocast well, use float32 + return torch.autocast("cpu", dtype=torch.float32, enabled=False) + + +def get_attention_mode(device): + """Get the best attention mode for the device.""" + if device.type == "cuda": + # Try flash attention first, fall back to sdpa + try: + from flash_attn import flash_attn_varlen_func + return "flash" + except ImportError: + return "sdpa" + elif device.type == "mps": + try: + from mps_flash_attn import is_available + if is_available(): + return "mps_flash" + except ImportError: + pass + return "sdpa" + else: + return "sdpa" + + +def quantize_dit_nf4(model: nn.Module, compute_dtype=torch.bfloat16): + """ + Quantize DiT model to NF4 (4-bit) for memory-efficient inference. + + Reduces ~10B params from ~20GB (bf16) to ~5GB (nf4). + Skips the patch embedding Conv2d (small, keep full precision). + """ + if not HAS_MPS_QUANTIZATION: + print("Warning: mps_bitsandbytes not available, skipping quantization") + return model + + print("Quantizing DiT to NF4...") + config = BitsAndBytesConfig( + load_in_4bit=True, + bnb_4bit_quant_type="nf4", + bnb_4bit_compute_dtype=compute_dtype, + ) + + # Skip the patch embed conv and any normalization layers + modules_to_skip = ["x_embedder.proj", "norm", "final_layer"] + + model = quantize_model( + model, + quantization_config=config, + modules_to_not_convert=modules_to_skip, + ) + + # Report memory savings + try: + from mps_bitsandbytes.integration import get_memory_footprint + stats = get_memory_footprint(model) + print(f" Model size: {stats['actual_size_gb']:.2f} GB (saved {stats['savings_pct']:.1f}%)") + except Exception: + pass + + return model + + +def save_quantized_dit(model: nn.Module, path: str): + """Save quantized DiT model to disk for fast loading later.""" + print(f"Saving quantized model to {path}...") + torch.save(model, path) + print(f" Saved!") + + +def load_quantized_dit(model: nn.Module, path: str, device=None): + """Load pre-quantized DiT model from disk (skips skeleton + quantization).""" + print(f"Loading pre-quantized model from {path}...") + + # Load entire model directly (includes Linear4bit modules) + model = torch.load(path, map_location="cpu", weights_only=False) + + if device is not None: + model = model.to(device) + print(f" Loaded!") + return model + + def linear_quadratic_schedule(num_steps, threshold_noise, linear_steps=None): if linear_steps is None: linear_steps = num_steps // 2 @@ -87,7 +209,8 @@ def setup_fsdp_sync(model, device_id, *, param_dtype, auto_wrap_policy) -> FSDP: sync_module_states=True, use_orig_params=True, ) - torch.cuda.synchronize() + if torch.cuda.is_available(): + torch.cuda.synchronize() return model @@ -108,7 +231,7 @@ def __init__(self, model_dir=None): super().__init__() self.model_dir = model_dir or T5_MODEL - def get_model(self, *, local_rank, device_id, world_size): + def get_model(self, *, local_rank, device_id, world_size, target_device=None): super().get_model(local_rank=local_rank, device_id=device_id, world_size=world_size) model = T5EncoderModel.from_pretrained(self.model_dir) if world_size > 1: @@ -123,6 +246,9 @@ def get_model(self, *, local_rank, device_id, world_size): }, ), ) + elif target_device is not None: + # MPS or explicit device path + model = model.to(target_device) elif isinstance(device_id, int): model = model.to(torch.device(f"cuda:{device_id}")) # type: ignore return model.eval() @@ -134,19 +260,20 @@ def __init__( model_path: str, model_dtype: str, lora_path: Optional[str] = None, - attention_mode: Optional[str] = None + attention_mode: Optional[str] = None, + quantize_nf4: bool = False, ): - # Infer attention mode if not specified + # Infer attention mode if not specified - use device-aware selection if attention_mode is None: - from genmo.lib.attn_imports import flash_varlen_attn # type: ignore - attention_mode = "sdpa" if flash_varlen_attn is None else "flash" + attention_mode = get_attention_mode(get_device()) print(f"Attention mode: {attention_mode}") super().__init__( model_path=model_path, lora_path=lora_path, model_dtype=model_dtype, - attention_mode=attention_mode + attention_mode=attention_mode, + quantize_nf4=quantize_nf4, ) def get_model( @@ -155,6 +282,7 @@ def get_model( local_rank, device_id, world_size, + target_device=None, model_kwargs=None, patch_model_fns=None, strict_load=True, @@ -242,6 +370,29 @@ def get_model( lambda_fn=lambda m: m in model.blocks, ), ) + elif target_device is not None: + # MPS or explicit device path - quantize before moving to device + if self.kwargs.get("quantize_nf4", False): + # Check for cached quantized weights - prefer FP16 version (faster on MPS) + model_path = self.kwargs["model_path"] + quant_cache_fp16 = model_path.replace(".safetensors", "_nf4_fp16.pt") + quant_cache_bf16 = model_path.replace(".safetensors", "_nf4.pt") + + if os.path.exists(quant_cache_fp16): + print(f"Found cached NF4+FP16 weights: {quant_cache_fp16}") + model = load_quantized_dit(model, quant_cache_fp16, device=target_device) + elif os.path.exists(quant_cache_bf16): + print(f"Found cached NF4+BF16 weights: {quant_cache_bf16}") + model = load_quantized_dit(model, quant_cache_bf16, device=target_device) + else: + # Quantize on-the-fly with FP16 for MPS speed + compute_dtype = torch.float16 if target_device.type == "mps" else torch.bfloat16 + model = quantize_dit_nf4(model, compute_dtype=compute_dtype) + model = model.to(target_device) + # Save FP16 version for next time + save_quantized_dit(model, quant_cache_fp16) + else: + model = model.to(target_device) elif isinstance(device_id, int): model = model.to(torch.device(f"cuda:{device_id}")) return model.eval() @@ -251,7 +402,7 @@ class DecoderModelFactory(ModelFactory): def __init__(self, *, model_path: str): super().__init__(model_path=model_path) - def get_model(self, *, local_rank=0, device_id=0, world_size=1): + def get_model(self, *, local_rank=0, device_id=0, world_size=1, target_device=None): # TODO(ved): Set flag for torch.compile # TODO(ved): Use skip_init @@ -272,7 +423,12 @@ def get_model(self, *, local_rank=0, device_id=0, world_size=1): # VAE is not FSDP-wrapped state_dict = load_file(self.kwargs["model_path"]) decoder.load_state_dict(state_dict, strict=True) - device = torch.device(f"cuda:{device_id}") if isinstance(device_id, int) else "cpu" + if target_device is not None: + device = target_device + elif isinstance(device_id, int): + device = torch.device(f"cuda:{device_id}") + else: + device = "cpu" decoder.eval().to(device) return decoder @@ -281,7 +437,7 @@ class EncoderModelFactory(ModelFactory): def __init__(self, *, model_path: str): super().__init__(model_path=model_path) - def get_model(self, *, local_rank=0, device_id=0, world_size=1): + def get_model(self, *, local_rank=0, device_id=0, world_size=1, target_device=None): # TODO(ved): Set flag for torch.compile # TODO(ved): Use skip_init @@ -303,7 +459,12 @@ def get_model(self, *, local_rank=0, device_id=0, world_size=1): ) state_dict = load_file(self.kwargs["model_path"]) encoder.load_state_dict(state_dict, strict=True) - device = torch.device(f"cuda:{device_id}") if isinstance(device_id, int) else "cpu" + if target_device is not None: + device = target_device + elif isinstance(device_id, int): + device = torch.device(f"cuda:{device_id}") + else: + device = "cpu" encoder.eval().to(device) return encoder @@ -451,19 +612,40 @@ def sample_model(device, dit, conditioning, **args): cond_batched["packed_indices"] = compute_packed_indices(device, cond_batched["y_mask"][0], num_latents) z = repeat(z, "b ... -> (repeat b) ...", repeat=2) + # Detect model dtype from Linear4bit compute_dtype (for quantized models) + model_dtype = torch.bfloat16 # Default fallback + if HAS_MPS_QUANTIZATION and Linear4bit is not None: + for module in dit.modules(): + if isinstance(module, Linear4bit): + model_dtype = module.compute_dtype + break + + # Cast conditioning tensors (T5 features are float32) to model dtype + def cast_cond_to_dtype(cond_dict, dtype): + if cond_dict is None: + return + if "y_feat" in cond_dict: + cond_dict["y_feat"] = [t.to(dtype) for t in cond_dict["y_feat"]] + + cast_cond_to_dtype(cond_text, model_dtype) + cast_cond_to_dtype(cond_null, model_dtype) + cast_cond_to_dtype(cond_batched, model_dtype) + def model_fn(*, z, sigma, cfg_scale): + # Cast z to model dtype for forward pass, keep original float32 for accumulation + z_input = z.to(model_dtype) + if cond_batched: - with torch.autocast("cuda", dtype=torch.bfloat16): - out = dit(z, sigma, **cond_batched) + out = dit(z_input, sigma, **cond_batched) out_cond, out_uncond = torch.chunk(out, chunks=2, dim=0) else: nonlocal cond_text, cond_null - with torch.autocast("cuda", dtype=torch.bfloat16): - out_cond = dit(z, sigma, **cond_text) - out_uncond = dit(z, sigma, **cond_null) - assert out_cond.shape == out_uncond.shape - out_uncond = out_uncond.to(z) - out_cond = out_cond.to(z) + out_cond = dit(z_input, sigma, **cond_text) + out_uncond = dit(z_input, sigma, **cond_null) + + # CFG in float32 for precision (critical for MPS!) + out_cond = out_cond.to(torch.float32) + out_uncond = out_uncond.to(torch.float32) return out_uncond + cfg_scale * (out_cond - out_uncond) # Euler sampler w/ customizable sigma schedule & cfg scale @@ -472,9 +654,16 @@ def model_fn(*, z, sigma, cfg_scale): dsigma = sigma - sigma_schedule[i + 1] # `pred` estimates `z_0 - eps`. + # Sigma tensor must be float32 and match z dtype for precision + sigma_tensor = torch.full( + [B] if cond_text else [B * 2], + sigma, + device=z.device, + dtype=torch.float32 + ) pred = model_fn( z=z, - sigma=torch.full([B] if cond_text else [B * 2], sigma, device=z.device), + sigma=sigma_tensor, cfg_scale=cfg_schedule[i], ) assert pred.dtype == torch.float32 @@ -507,6 +696,20 @@ def t5_tokenizer(model_dir=None): return T5Tokenizer.from_pretrained(model_dir or T5_MODEL, legacy=False) +def get_max_memory_fn(device): + """Get a memory reporting function appropriate for the device.""" + if device.type == "cuda": + return lambda: print(f"Max memory reserved: {torch.cuda.max_memory_reserved() / 1024**3:.2f} GB") + elif device.type == "mps": + # MPS doesn't have direct memory query, but we can try + try: + return lambda: print(f"MPS memory allocated: {torch.mps.current_allocated_memory() / 1024**3:.2f} GB") + except AttributeError: + return lambda: print("MPS memory stats not available") + else: + return lambda: print("CPU mode - no GPU memory tracking") + + class MochiSingleGPUPipeline: def __init__( self, @@ -518,34 +721,83 @@ def __init__( decode_type: str = "full", decode_args: Optional[Dict[str, Any]] = None, fast_init=True, - strict_load=True + strict_load=True, + lazy_load=False, # NEW: load models on-demand for memory-constrained devices ): - self.device = torch.device("cuda:0") + self.device = get_device() + print(f"Using device: {self.device}") self.tokenizer = t5_tokenizer(text_encoder_factory.model_dir) - t = Timer() self.cpu_offload = cpu_offload self.decode_args = decode_args or {} self.decode_type = decode_type - init_id = "cpu" if cpu_offload else 0 - with t("load_text_encoder"): - self.text_encoder = text_encoder_factory.get_model( - local_rank=0, - device_id=init_id, - world_size=1, - ) - with t("load_dit"): - self.dit = dit_factory.get_model(local_rank=0, device_id=init_id, world_size=1, fast_init=fast_init, strict_load=strict_load) # type: ignore - with t("load_vae"): - self.decoder = decoder_factory.get_model(local_rank=0, device_id=init_id, world_size=1) - t.print_stats() + self.fast_init = fast_init + self.strict_load = strict_load + + # Store factories for lazy loading + self._text_encoder_factory = text_encoder_factory + self._dit_factory = dit_factory + self._decoder_factory = decoder_factory + self._lazy_load = lazy_load or (self.device.type == "mps") # Auto-enable for MPS + + # For MPS/CPU, we use target_device; for CUDA we use device_id + if self.device.type == "cuda": + self._init_id = "cpu" if cpu_offload else 0 + self._target_device = None + else: + self._init_id = "cpu" + self._target_device = None if cpu_offload else self.device + + if self._lazy_load: + print("Lazy loading enabled - models will be loaded on-demand") + self.text_encoder = None + self.dit = None + self.decoder = None + else: + t = Timer() + with t("load_text_encoder"): + self.text_encoder = self._load_text_encoder() + with t("load_dit"): + self.dit = self._load_dit() + with t("load_vae"): + self.decoder = self._load_decoder() + t.print_stats() + + def _load_text_encoder(self): + return self._text_encoder_factory.get_model( + local_rank=0, + device_id=self._init_id, + world_size=1, + target_device=self._target_device, + ) + + def _load_dit(self): + return self._dit_factory.get_model( + local_rank=0, + device_id=self._init_id, + world_size=1, + target_device=self._target_device, + fast_init=self.fast_init, + strict_load=self.strict_load + ) + + def _load_decoder(self): + return self._decoder_factory.get_model( + local_rank=0, + device_id=self._init_id, + world_size=1, + target_device=self._target_device, + ) def __call__(self, batch_cfg, prompt, negative_prompt, **kwargs): with torch.inference_mode(): - print_max_memory = lambda: print( - f"Max memory reserved: {torch.cuda.max_memory_reserved() / 1024**3:.2f} GB" - ) + print_max_memory = get_max_memory_fn(self.device) print_max_memory() + # Lazy load T5 if needed + if self.text_encoder is None: + print("Loading T5 encoder...") + self.text_encoder = self._load_text_encoder() + with move_to_device(self.text_encoder, self.device): conditioning = get_conditioning( tokenizer=self.tokenizer, @@ -555,13 +807,46 @@ def __call__(self, batch_cfg, prompt, negative_prompt, **kwargs): prompt=prompt, negative_prompt=negative_prompt, ) + # Free T5 memory - it's not needed after encoding! + del self.text_encoder + self.text_encoder = None + if self.device.type == "mps": + torch.mps.empty_cache() + elif self.device.type == "cuda": + torch.cuda.empty_cache() + print("T5 encoder freed from memory") print_max_memory() + # Lazy load DiT if needed + if self.dit is None: + print("Loading DiT...") + self.dit = self._load_dit() + with move_to_device(self.dit, self.device): latents = sample_model(self.device, self.dit, conditioning, **kwargs) print_max_memory() + # Free DiT memory before VAE decode on memory-constrained devices + if self.device.type == "mps": + latents = latents.cpu() # Move to CPU before freeing DiT + del self.dit + self.dit = None + del conditioning # Free T5 embeddings + import gc + gc.collect() + torch.mps.empty_cache() + + # Lazy load VAE decoder if needed + if self.decoder is None: + print("Loading VAE decoder...") + self.decoder = self._load_decoder() + with move_to_device(self.decoder, self.device): + import time as _time + _vae_start = _time.time() + print("Starting VAE decode...") + # Move latents back to device for decoding + latents = latents.to(self.device) if self.decode_type == "tiled_full": frames = decode_latents_tiled_full( self.decoder, latents, **self.decode_args) @@ -571,6 +856,9 @@ def __call__(self, batch_cfg, prompt, negative_prompt, **kwargs): num_tiles_w=4, num_tiles_h=2) else: frames = decode_latents(self.decoder, latents) + if self.device.type == "mps": + torch.mps.synchronize() + print(f"VAE decode took {_time.time() - _vae_start:.1f}s") print_max_memory() return frames.cpu().numpy() @@ -644,6 +932,8 @@ def __init__( decoder_factory: ModelFactory, world_size: int, ): + if not HAS_RAY: + raise ImportError("ray is required for multi-GPU mode: pip install ray") ray.init() RemoteClass = ray.remote(MultiGPUContext) self.ctxs = [ diff --git a/src/genmo/mochi_preview/vae/models.py b/src/genmo/mochi_preview/vae/models.py index 568bcab..cb953a8 100644 --- a/src/genmo/mochi_preview/vae/models.py +++ b/src/genmo/mochi_preview/vae/models.py @@ -1015,7 +1015,16 @@ def decode_latents(decoder, z): assert z.ndim == 5 cp_rank, cp_size = cp.get_cp_rank_size() z = z.tensor_split(cp_size, dim=2)[cp_rank] # split along temporal dim - with torch.autocast("cuda", dtype=torch.bfloat16): + device = z.device + + if device.type == "cuda": + ctx = torch.autocast("cuda", dtype=torch.bfloat16) + elif device.type == "mps": + # MPS: use float32 - bfloat16 autocast causes corruption + ctx = torch.autocast("mps", dtype=torch.float32, enabled=False) + else: + ctx = torch.autocast("cpu", dtype=torch.float32, enabled=False) + with ctx: samples = decoder(z) samples = gather_all_frames(samples) return normalize_decoded_frames(samples)