diff --git a/mlx_video/models/ltx_2/README.md b/mlx_video/models/ltx_2/README.md index f84400e..5578fac 100644 --- a/mlx_video/models/ltx_2/README.md +++ b/mlx_video/models/ltx_2/README.md @@ -155,6 +155,32 @@ uv run mlx_video.upscale --input video.mp4 --output upscaled.mp4 --refine --prom | `--audio-start-time` | 0.0 | Start time in seconds for audio file | | `--tiling` | `auto` | VAE tiling mode: `auto`, `none`, `aggressive`, `conservative` | | `--stream` | false | Stream frames as they decode | +| `--spatial-upscaler` | auto (x2) | Spatial upscaler file for two-stage pipelines (see below) | + +### Spatial Upscalers (LTX-2.3) + +LTX-2.3 ships with multiple spatial upscaler variants. Use `--spatial-upscaler` to select one: + +| Variant | Scale | Output (from 256x256) | Architecture | +|---------|-------|-----------------------|--------------| +| `ltx-2.3-spatial-upscaler-x2-1.0.safetensors` (default) | 2.0x | 512x512 | Conv2d + PixelShuffle(2) | +| `ltx-2.3-spatial-upscaler-x2-1.1.safetensors` | 2.0x | 512x512 | Same arch, newer weights | +| `ltx-2.3-spatial-upscaler-x1.5-1.0.safetensors` | 1.5x | 384x384 | Conv2d + PixelShuffle(3) + BlurDownsample | + +```bash +# Default (x2-1.0, auto-detected) +uv run mlx_video.generate --prompt "A sunset" --model-repo ./LTX-2.3-distilled + +# x2-1.1 (newer weights) +uv run mlx_video.generate --prompt "A sunset" --model-repo ./LTX-2.3-distilled \ + --spatial-upscaler ltx-2.3-spatial-upscaler-x2-1.1.safetensors + +# x1.5 (smaller output, faster) +uv run mlx_video.generate --prompt "A sunset" --model-repo ./LTX-2.3-distilled \ + --spatial-upscaler ltx-2.3-spatial-upscaler-x1.5-1.0.safetensors +``` + +> **Note:** Stage 1 always runs at half the target resolution. With x1.5, the final output is 75% of `--width`/`--height` (e.g., 512 target -> 256 stage 1 -> 384 output). With x2, the output matches the target exactly. ### Dev / Dev-Two-Stage @@ -189,8 +215,8 @@ HQ defaults: 15 steps (vs 30), `cfg-rescale` 0.45 (vs 0.7), STG disabled. Uses t ### Distilled Pipeline (default) 1. **Stage 1**: Generate at half resolution with 8 denoising steps (fixed sigmas) -2. **Upsample**: 2x spatial upsampling via LatentUpsampler -3. **Stage 2**: Refine at full resolution with 3 denoising steps +2. **Upsample**: Spatial upsampling via LatentUpsampler (x2 or x1.5, selectable via `--spatial-upscaler`) +3. **Stage 2**: Refine at upsampled resolution with 3 denoising steps 4. **Decode**: VAE decoder converts latents to RGB video ### Dev Pipeline @@ -199,14 +225,14 @@ HQ defaults: 15 steps (vs 30), `cfg-rescale` 0.45 (vs 0.7), STG disabled. Uses t ### Dev Two-Stage Pipeline 1. **Stage 1**: Dev denoising at half resolution with CFG -2. **Upsample**: 2x spatial upsampling via LatentUpsampler -3. **Stage 2**: Distilled refinement at full resolution with LoRA weights (3 steps, no CFG) +2. **Upsample**: Spatial upsampling via LatentUpsampler (x2 or x1.5) +3. **Stage 2**: Distilled refinement at upsampled resolution with LoRA weights (3 steps, no CFG) 4. **Decode**: VAE decoder converts latents to RGB video ### Dev Two-Stage HQ Pipeline 1. **Stage 1**: res_2s denoising at half resolution with CFG + LoRA@0.25 (15 steps, 2 evals/step) -2. **Upsample**: 2x spatial upsampling via LatentUpsampler -3. **Stage 2**: res_2s refinement at full resolution with LoRA@0.5 (3 steps, no CFG) +2. **Upsample**: Spatial upsampling via LatentUpsampler (x2 or x1.5) +3. **Stage 2**: res_2s refinement at upsampled resolution with LoRA@0.5 (3 steps, no CFG) 4. **Decode**: VAE decoder converts latents to RGB video The res_2s sampler uses an exponential Rosenbrock-type Runge-Kutta integrator with SDE noise injection, producing higher quality results than Euler at the same compute budget (~30 total model evaluations). diff --git a/mlx_video/models/ltx_2/generate.py b/mlx_video/models/ltx_2/generate.py index c7df2dc..81b815f 100644 --- a/mlx_video/models/ltx_2/generate.py +++ b/mlx_video/models/ltx_2/generate.py @@ -1461,6 +1461,7 @@ def generate_video( lora_strength_stage_2: Optional[float] = None, audio_file: Optional[str] = None, audio_start_time: float = 0.0, + spatial_upscaler: Optional[str] = None, ): """Generate video using LTX-2 models. @@ -1557,10 +1558,35 @@ def generate_video( model_path = get_model_path(model_repo) text_encoder_path = model_path if text_encoder_repo is None else get_model_path(text_encoder_repo) + # Resolve spatial upscaler path for two-stage pipelines + upscaler_path = None + upscaler_scale = 2.0 + if is_two_stage: + if spatial_upscaler is not None: + # User-specified upscaler file + upscaler_path = model_path / spatial_upscaler if not Path(spatial_upscaler).is_absolute() else Path(spatial_upscaler) + if not upscaler_path.exists(): + # Try as a filename within model_path + upscaler_path = model_path / spatial_upscaler + # Detect scale from filename + if "x1.5" in str(upscaler_path): + upscaler_scale = 1.5 + elif "x2" in str(upscaler_path): + upscaler_scale = 2.0 + else: + # Auto-detect: prefer x2 upscaler + upscaler_files = sorted(model_path.glob("*spatial-upscaler-x2*.safetensors")) + if upscaler_files: + upscaler_path = upscaler_files[0] + upscaler_scale = 2.0 + # Calculate latent dimensions if is_two_stage: + # Stage 1 always at half resolution (matches PyTorch) stage1_h, stage1_w = height // 2 // 32, width // 2 // 32 - stage2_h, stage2_w = height // 32, width // 32 + # Stage 2 resolution = stage 1 * upscaler scale + stage2_h = int(stage1_h * upscaler_scale) + stage2_w = int(stage1_w * upscaler_scale) else: latent_h, latent_w = height // 32, width // 32 latent_frames = 1 + (num_frames - 1) // 8 @@ -1697,13 +1723,15 @@ def generate_video( with console.status("[blue]🖼️ Loading VAE encoder and encoding image...[/]", spinner="dots"): vae_encoder = VideoEncoder.from_pretrained(model_path / "vae" / "encoder") - input_image = load_image(image, height=height // 2, width=width // 2, dtype=model_dtype) - stage1_image_tensor = prepare_image_for_encoding(input_image, height // 2, width // 2, dtype=model_dtype) + s1_h, s1_w = stage1_h * 32, stage1_w * 32 + input_image = load_image(image, height=s1_h, width=s1_w, dtype=model_dtype) + stage1_image_tensor = prepare_image_for_encoding(input_image, s1_h, s1_w, dtype=model_dtype) stage1_image_latent = vae_encoder(stage1_image_tensor) mx.eval(stage1_image_latent) - input_image = load_image(image, height=height, width=width, dtype=model_dtype) - stage2_image_tensor = prepare_image_for_encoding(input_image, height, width, dtype=model_dtype) + s2_h, s2_w = stage2_h * 32, stage2_w * 32 + input_image = load_image(image, height=s2_h, width=s2_w, dtype=model_dtype) + stage2_image_tensor = prepare_image_for_encoding(input_image, s2_h, s2_w, dtype=model_dtype) stage2_image_latent = vae_encoder(stage2_image_tensor) mx.eval(stage2_image_latent) @@ -1712,7 +1740,7 @@ def generate_video( console.print("[green]✓[/] VAE encoder loaded and image encoded") # Stage 1 - console.print(f"\n[bold yellow]⚡ Stage 1:[/] Generating at {width//2}x{height//2} (8 steps)") + console.print(f"\n[bold yellow]⚡ Stage 1:[/] Generating at {stage1_w*32}x{stage1_h*32} (8 steps)") mx.random.seed(seed) positions = create_position_grid(1, latent_frames, stage1_h, stage1_w) @@ -1757,11 +1785,10 @@ def generate_video( ) # Upsample latents - with console.status("[magenta]🔍 Upsampling latents 2x...[/]", spinner="dots"): - upscaler_files = sorted(model_path.glob("*spatial-upscaler-x2*.safetensors")) - if not upscaler_files: + with console.status(f"[magenta]🔍 Upsampling latents {upscaler_scale}x...[/]", spinner="dots"): + if upscaler_path is None or not upscaler_path.exists(): raise FileNotFoundError(f"No spatial upscaler found in {model_path}") - upsampler = load_upsampler(str(upscaler_files[0])) + upsampler, upscaler_scale = load_upsampler(str(upscaler_path)) mx.eval(upsampler.parameters()) vae_decoder = VideoDecoder.from_pretrained(str(model_path / "vae" / "decoder")) @@ -1774,7 +1801,7 @@ def generate_video( console.print("[green]✓[/] Latents upsampled") # Stage 2 - console.print(f"\n[bold yellow]⚡ Stage 2:[/] Refining at {width}x{height} (3 steps)") + console.print(f"\n[bold yellow]⚡ Stage 2:[/] Refining at {stage2_w*32}x{stage2_h*32} (3 steps)") positions = create_position_grid(1, latent_frames, stage2_h, stage2_w) mx.eval(positions) @@ -1916,13 +1943,15 @@ def generate_video( with console.status("[blue]🖼️ Loading VAE encoder and encoding image...[/]", spinner="dots"): vae_encoder = VideoEncoder.from_pretrained(model_path / "vae" / "encoder") - input_image = load_image(image, height=height // 2, width=width // 2, dtype=model_dtype) - stage1_image_tensor = prepare_image_for_encoding(input_image, height // 2, width // 2, dtype=model_dtype) + s1_h, s1_w = stage1_h * 32, stage1_w * 32 + input_image = load_image(image, height=s1_h, width=s1_w, dtype=model_dtype) + stage1_image_tensor = prepare_image_for_encoding(input_image, s1_h, s1_w, dtype=model_dtype) stage1_image_latent = vae_encoder(stage1_image_tensor) mx.eval(stage1_image_latent) - input_image = load_image(image, height=height, width=width, dtype=model_dtype) - stage2_image_tensor = prepare_image_for_encoding(input_image, height, width, dtype=model_dtype) + s2_h, s2_w = stage2_h * 32, stage2_w * 32 + input_image = load_image(image, height=s2_h, width=s2_w, dtype=model_dtype) + stage2_image_tensor = prepare_image_for_encoding(input_image, s2_h, s2_w, dtype=model_dtype) stage2_image_latent = vae_encoder(stage2_image_tensor) mx.eval(stage2_image_latent) @@ -1930,12 +1959,12 @@ def generate_video( mx.clear_cache() console.print("[green]✓[/] VAE encoder loaded and image encoded") - # Stage 1: Dev denoising at half resolution with CFG + # Stage 1: Dev denoising at reduced resolution with CFG sigmas = ltx2_scheduler(steps=num_inference_steps) mx.eval(sigmas) console.print(f"[dim]Stage 1 sigma schedule: {sigmas[0].item():.4f} → {sigmas[-2].item():.4f} → {sigmas[-1].item():.4f}[/]") - console.print(f"\n[bold yellow]⚡ Stage 1:[/] Dev generating at {width//2}x{height//2} ({num_inference_steps} steps, CFG={cfg_scale}, rescale={cfg_rescale})") + console.print(f"\n[bold yellow]⚡ Stage 1:[/] Dev generating at {stage1_w*32}x{stage1_h*32} ({num_inference_steps} steps, CFG={cfg_scale}, rescale={cfg_rescale})") mx.random.seed(seed) positions = create_position_grid(1, latent_frames, stage1_h, stage1_w) @@ -1989,12 +2018,11 @@ def generate_video( mx.eval(audio_latents) - # Upsample latents 2x - with console.status("[magenta]🔍 Upsampling latents 2x...[/]", spinner="dots"): - upscaler_files = sorted(model_path.glob("*spatial-upscaler-x2*.safetensors")) - if not upscaler_files: + # Upsample latents + with console.status(f"[magenta]🔍 Upsampling latents {upscaler_scale}x...[/]", spinner="dots"): + if upscaler_path is None or not upscaler_path.exists(): raise FileNotFoundError(f"No spatial upscaler found in {model_path}") - upsampler = load_upsampler(str(upscaler_files[0])) + upsampler, upscaler_scale = load_upsampler(str(upscaler_path)) mx.eval(upsampler.parameters()) vae_decoder = VideoDecoder.from_pretrained(str(model_path / "vae" / "decoder")) @@ -2091,13 +2119,15 @@ def generate_video( with console.status("[blue]Loading VAE encoder and encoding image...[/]", spinner="dots"): vae_encoder = VideoEncoder.from_pretrained(model_path / "vae" / "encoder") - input_image = load_image(image, height=height // 2, width=width // 2, dtype=model_dtype) - stage1_image_tensor = prepare_image_for_encoding(input_image, height // 2, width // 2, dtype=model_dtype) + s1_h, s1_w = stage1_h * 32, stage1_w * 32 + input_image = load_image(image, height=s1_h, width=s1_w, dtype=model_dtype) + stage1_image_tensor = prepare_image_for_encoding(input_image, s1_h, s1_w, dtype=model_dtype) stage1_image_latent = vae_encoder(stage1_image_tensor) mx.eval(stage1_image_latent) - input_image = load_image(image, height=height, width=width, dtype=model_dtype) - stage2_image_tensor = prepare_image_for_encoding(input_image, height, width, dtype=model_dtype) + s2_h, s2_w = stage2_h * 32, stage2_w * 32 + input_image = load_image(image, height=s2_h, width=s2_w, dtype=model_dtype) + stage2_image_tensor = prepare_image_for_encoding(input_image, s2_h, s2_w, dtype=model_dtype) stage2_image_latent = vae_encoder(stage2_image_tensor) mx.eval(stage2_image_latent) @@ -2118,14 +2148,14 @@ def generate_video( with console.status(f"[blue]Merging distilled LoRA (stage 1, strength={hq_lora_strength_s1})...[/]", spinner="dots"): load_and_merge_lora(transformer, lora_path, strength=hq_lora_strength_s1) - # Stage 1: res_2s denoising at half resolution with CFG + # Stage 1: res_2s denoising at reduced resolution with CFG # HQ passes actual token count to scheduler (unlike regular dev-two-stage) num_tokens = latent_frames * stage1_h * stage1_w sigmas = ltx2_scheduler(steps=hq_steps, num_tokens=num_tokens) mx.eval(sigmas) console.print(f"[dim]Stage 1 sigma schedule: {sigmas[0].item():.4f} -> {sigmas[-2].item():.4f} -> {sigmas[-1].item():.4f} (tokens={num_tokens})[/]") - console.print(f"\n[bold yellow]Stage 1:[/] res_2s at {width//2}x{height//2} ({hq_steps} steps, CFG={cfg_scale}, rescale={hq_cfg_rescale})") + console.print(f"\n[bold yellow]Stage 1:[/] res_2s at {stage1_w*32}x{stage1_h*32} ({hq_steps} steps, CFG={cfg_scale}, rescale={hq_cfg_rescale})") mx.random.seed(seed) positions = create_position_grid(1, latent_frames, stage1_h, stage1_w) @@ -2179,12 +2209,11 @@ def generate_video( mx.eval(audio_latents) - # Upsample latents 2x - with console.status("[magenta]Upsampling latents 2x...[/]", spinner="dots"): - upscaler_files = sorted(model_path.glob("*spatial-upscaler-x2*.safetensors")) - if not upscaler_files: + # Upsample latents + with console.status(f"[magenta]Upsampling latents {upscaler_scale}x...[/]", spinner="dots"): + if upscaler_path is None or not upscaler_path.exists(): raise FileNotFoundError(f"No spatial upscaler found in {model_path}") - upsampler = load_upsampler(str(upscaler_files[0])) + upsampler, upscaler_scale = load_upsampler(str(upscaler_path)) mx.eval(upsampler.parameters()) vae_decoder = VideoDecoder.from_pretrained(str(model_path / "vae" / "decoder")) @@ -2204,7 +2233,7 @@ def generate_video( load_and_merge_lora(transformer, lora_path, strength=additional_strength) # Stage 2: res_2s refinement at full resolution (no CFG) - console.print(f"\n[bold yellow]Stage 2:[/] res_2s refining at {width}x{height} (3 steps, no CFG)") + console.print(f"\n[bold yellow]Stage 2:[/] res_2s refining at {stage2_w*32}x{stage2_h*32} (3 steps, no CFG)") positions = create_position_grid(1, latent_frames, stage2_h, stage2_w) mx.eval(positions) @@ -2509,6 +2538,9 @@ Examples: parser.add_argument("--lora-strength", type=float, default=1.0, help="LoRA merge strength (dev-two-stage pipeline, default 1.0)") parser.add_argument("--lora-strength-stage-1", type=float, default=0.25, help="LoRA strength for HQ stage 1 (default 0.25)") parser.add_argument("--lora-strength-stage-2", type=float, default=0.5, help="LoRA strength for HQ stage 2 (default 0.5)") + parser.add_argument("--spatial-upscaler", type=str, default=None, + help="Spatial upscaler filename (e.g. ltx-2.3-spatial-upscaler-x1.5-1.0.safetensors). " + "Auto-detects x2 by default. Use this to select x1.5 or a specific version.") args = parser.parse_args() pipeline_map = { @@ -2559,6 +2591,7 @@ Examples: lora_strength_stage_2=args.lora_strength_stage_2, audio_file=args.audio_file, audio_start_time=args.audio_start_time, + spatial_upscaler=args.spatial_upscaler, ) diff --git a/mlx_video/models/ltx_2/upsampler.py b/mlx_video/models/ltx_2/upsampler.py index 1180664..9ede781 100644 --- a/mlx_video/models/ltx_2/upsampler.py +++ b/mlx_video/models/ltx_2/upsampler.py @@ -115,65 +115,135 @@ class GroupNorm3d(nn.Module): class PixelShuffle2D(nn.Module): - """Pixel shuffle for 2D spatial upsampling.""" + """Pixel shuffle for 2D spatial upsampling with per-axis factors.""" - def __init__(self, upscale_factor: int = 2): + def __init__(self, upscale_factor_h: int = 2, upscale_factor_w: int = 2): super().__init__() - self.upscale_factor = upscale_factor + self.rh = upscale_factor_h + self.rw = upscale_factor_w def __call__(self, x: mx.array) -> mx.array: - # x: (N, H, W, C) where C = out_channels * upscale_factor^2 + # x: (N, H, W, C) where C = out_channels * rh * rw n, h, w, c = x.shape - r = self.upscale_factor - out_c = c // (r * r) + rh, rw = self.rh, self.rw + out_c = c // (rh * rw) - # Reshape: (N, H, W, out_c, r, r) - x = mx.reshape(x, (n, h, w, out_c, r, r)) + # Reshape: (N, H, W, out_c, rh, rw) + x = mx.reshape(x, (n, h, w, out_c, rh, rw)) - # Permute: (N, H, r, W, r, out_c) + # Permute: (N, H, rh, W, rw, out_c) x = mx.transpose(x, (0, 1, 4, 2, 5, 3)) - # Reshape: (N, H*r, W*r, out_c) - x = mx.reshape(x, (n, h * r, w * r, out_c)) + # Reshape: (N, H*rh, W*rw, out_c) + x = mx.reshape(x, (n, h * rh, w * rw, out_c)) return x +class BlurDownsample(nn.Module): + """Anti-aliased downsampling with a fixed 5x5 binomial blur kernel. + + PyTorch source uses a depthwise conv with the binomial kernel. + The kernel weight is stored as (1, 1, 5, 5) and loaded via safetensors. + """ + + def __init__(self, stride: int = 2): + super().__init__() + self.stride = stride + # 5x5 binomial (1,4,6,4,1) kernel, normalized + # This will be overwritten by loaded weights if available + k = mx.array([1.0, 4.0, 6.0, 4.0, 1.0]) + kernel_2d = mx.outer(k, k) + kernel_2d = kernel_2d / kernel_2d.sum() + # MLX conv2d weight: (O, H, W, I) — we use (1, 5, 5, 1) for per-channel + self.kernel = kernel_2d.reshape(1, 5, 5, 1) + + def __call__(self, x: mx.array) -> mx.array: + # x: (N, H, W, C) channels-last + n, h, w, c = x.shape + + # Pad with edge replication (2 on each side for 5x5 kernel) + x = mx.pad(x, [(0, 0), (2, 2), (2, 2), (0, 0)], mode="edge") + + # Apply blur per-channel: reshape so each channel is a separate "batch" + # (N, H+4, W+4, C) -> (N*C, H+4, W+4, 1) + x = mx.transpose(x, (0, 3, 1, 2)) # (N, C, H+4, W+4) + x = mx.reshape(x, (n * c, h + 4, w + 4, 1)) + + # Depthwise conv: (N*C, H+4, W+4, 1) * (1, 5, 5, 1) -> (N*C, H_out, W_out, 1) + x = mx.conv2d(x, self.kernel, stride=(self.stride, self.stride)) + + _, h_out, w_out, _ = x.shape + # Reshape back: (N*C, H_out, W_out, 1) -> (N, C, H_out, W_out) -> (N, H_out, W_out, C) + x = mx.reshape(x, (n, c, h_out, w_out)) + x = mx.transpose(x, (0, 2, 3, 1)) + + return x + + +class SpatialUpsampler2x(nn.Module): + """Standard 2x spatial upsampler: Conv2d + PixelShuffle(2).""" + + def __init__(self, mid_channels: int = 1024): + super().__init__() + self.scale = 2.0 + # Sequential: conv (index 0) + pixel shuffle + # Weight key: upsampler.0.weight -> mapped to upsampler.conv.weight in sanitize + self.conv = nn.Conv2d(mid_channels, 4 * mid_channels, kernel_size=3, padding=1) + self.pixel_shuffle = PixelShuffle2D(2, 2) + + def __call__(self, x: mx.array) -> mx.array: + # x: (N, D, H, W, C) + n, d, h, w, c = x.shape + x = mx.reshape(x, (n * d, h, w, c)) + x = self.conv(x) + x = self.pixel_shuffle(x) + x = mx.reshape(x, (n, d, h * 2, w * 2, c)) + return x + + class SpatialRationalResampler(nn.Module): + """Rational spatial resampler for non-integer scale factors (e.g., 1.5x). - def __init__(self, mid_channels: int = 1024, scale: float = 2.0): + For scale=1.5: upsample 3x via PixelShuffle, then downsample 2x via BlurDownsample. + Rational fraction: 1.5 = 3/2. + """ + + def __init__(self, mid_channels: int = 1024, scale: float = 1.5): super().__init__() self.scale = scale - # 2D conv: mid_channels -> 4*mid_channels for pixel shuffle - self.conv = nn.Conv2d(mid_channels, 4 * mid_channels, kernel_size=3, padding=1) + # Rational fraction for 1.5: numerator=3, denominator=2 + num, den = _rational_for_scale(scale) + self.num = num + self.den = den - # Blur kernel for antialiasing - self.blur_down_kernel = mx.ones((1, 1, 5, 5)) / 25.0 - - self.pixel_shuffle = PixelShuffle2D(2) + # Conv2d: mid_channels -> num^2 * mid_channels for PixelShuffle(num) + self.conv = nn.Conv2d(mid_channels, num * num * mid_channels, kernel_size=3, padding=1) + self.pixel_shuffle = PixelShuffle2D(num, num) + self.blur_down = BlurDownsample(stride=den) def __call__(self, x: mx.array) -> mx.array: - # x: (N, D, H, W, C) - channels last 3D format - + # x: (N, D, H, W, C) n, d, h, w, c = x.shape - - # Process frame by frame - # Reshape to (N*D, H, W, C) for 2D operations x = mx.reshape(x, (n * d, h, w, c)) - # Apply 2D conv x = self.conv(x) + x = self.pixel_shuffle(x) # H*num, W*num + x = self.blur_down(x) # H*num/den, W*num/den - # Pixel shuffle for 2x upscaling - x = self.pixel_shuffle(x) - - # Reshape back to (N, D, H*2, W*2, C) - x = mx.reshape(x, (n, d, h * 2, w * 2, c)) - + _, h_out, w_out, _ = x.shape + x = mx.reshape(x, (n, d, h_out, w_out, c)) return x +def _rational_for_scale(scale: float) -> Tuple[int, int]: + """Convert a float scale to a rational fraction (numerator, denominator).""" + from fractions import Fraction + frac = Fraction(scale).limit_denominator(10) + return frac.numerator, frac.denominator + + class ResBlock3D(nn.Module): def __init__(self, channels: int): @@ -201,17 +271,19 @@ class ResBlock3D(nn.Module): class LatentUpsampler(nn.Module): - def __init__( self, in_channels: int = 128, mid_channels: int = 1024, num_blocks_per_stage: int = 4, + spatial_scale: float = 2.0, + rational_resampler: bool = False, ): super().__init__() self.in_channels = in_channels self.mid_channels = mid_channels + self.spatial_scale = spatial_scale # Initial projection self.initial_conv = Conv3d(in_channels, mid_channels, kernel_size=3, padding=1) @@ -221,7 +293,10 @@ class LatentUpsampler(nn.Module): self.res_blocks = {i: ResBlock3D(mid_channels) for i in range(num_blocks_per_stage)} # Upsampler: 2D spatial upsampling (frame-by-frame) - self.upsampler = SpatialRationalResampler(mid_channels=mid_channels, scale=2.0) + if rational_resampler: + self.upsampler = SpatialRationalResampler(mid_channels=mid_channels, scale=spatial_scale) + else: + self.upsampler = SpatialUpsampler2x(mid_channels=mid_channels) # Post-upsample ResBlocks - use dict with int keys for MLX parameter tracking self.post_upsample_res_blocks = {i: ResBlock3D(mid_channels) for i in range(num_blocks_per_stage)} @@ -230,14 +305,14 @@ class LatentUpsampler(nn.Module): self.final_conv = Conv3d(mid_channels, in_channels, kernel_size=3, padding=1) def __call__(self, latent: mx.array, debug: bool = False) -> mx.array: - """Upsample latents by 2x spatially. + """Upsample latents spatially. Args: latent: Input tensor of shape (B, C, F, H, W) - channels first debug: If True, print intermediate values for debugging Returns: - Upsampled tensor of shape (B, C, F, H*2, W*2) - channels first + Upsampled tensor of shape (B, C, F, H*scale, W*scale) - channels first """ def debug_stats(name, t): if debug: @@ -250,41 +325,27 @@ class LatentUpsampler(nn.Module): # Convert from channels first (B, C, F, H, W) to channels last (B, F, H, W, C) x = mx.transpose(latent, (0, 2, 3, 4, 1)) - if debug: - debug_stats("After transpose to channels-last", x) # Initial conv x = self.initial_conv(x) - if debug: - debug_stats("After initial_conv", x) x = self.initial_norm(x) - if debug: - debug_stats("After initial_norm", x) x = nn.silu(x) - if debug: - debug_stats("After silu", x) # Pre-upsample blocks for i in sorted(self.res_blocks.keys()): x = self.res_blocks[i](x) - if debug: - debug_stats(f"After res_blocks[{i}]", x) # Upsample (2D spatial, frame-by-frame) x = self.upsampler(x) if debug: - debug_stats("After upsampler (spatial 2x)", x) + debug_stats(f"After upsampler (spatial {self.spatial_scale}x)", x) # Post-upsample blocks for i in sorted(self.post_upsample_res_blocks.keys()): x = self.post_upsample_res_blocks[i](x) - if debug: - debug_stats(f"After post_upsample_res_blocks[{i}]", x) # Final conv x = self.final_conv(x) - if debug: - debug_stats("After final_conv", x) # Convert back to channels first (B, C, F, H, W) x = mx.transpose(x, (0, 4, 1, 2, 3)) @@ -315,33 +376,49 @@ def upsample_latents( return latent -def load_upsampler(weights_path: str) -> LatentUpsampler: +def load_upsampler(weights_path: str) -> Tuple[LatentUpsampler, float]: """Load upsampler from safetensors weights. + Auto-detects whether the weights are for x2 or x1.5 upscaling based on + the upsampler conv output channels: + - x2: upsampler.0.weight shape [4*mid, mid, 3, 3] (4096 out channels) + - x1.5: upsampler.conv.weight shape [9*mid, mid, 3, 3] (9216 out channels) + Args: weights_path: Path to upsampler weights file Returns: - Loaded LatentUpsampler model + Tuple of (LatentUpsampler model, spatial_scale) """ print(f"Loading spatial upsampler from {weights_path}...") raw_weights = mx.load(weights_path) - # Check weight shapes to determine mid_channels - # res_blocks.0.conv1.weight should be (mid_channels, mid_channels, 3, 3, 3) + # Detect mid_channels from res_blocks sample_key = "res_blocks.0.conv1.weight" if sample_key in raw_weights: mid_channels = raw_weights[sample_key].shape[0] else: - mid_channels = 1024 # default + mid_channels = 1024 - print(f" Detected mid_channels: {mid_channels}") + # Detect upsampler type from conv output channels + # x2 uses sequential: upsampler.0.weight (4*mid out channels) + # x1.5 uses named: upsampler.conv.weight (9*mid out channels) + upsampler.blur_down.kernel + rational_resampler = "upsampler.blur_down.kernel" in raw_weights + if rational_resampler: + # x1.5: conv out = 9 * mid_channels (3^2 * mid for PixelShuffle(3)) + spatial_scale = 1.5 + else: + spatial_scale = 2.0 + + print(f" Detected: mid_channels={mid_channels}, scale={spatial_scale}x, rational={rational_resampler}") # Create model upsampler = LatentUpsampler( in_channels=128, mid_channels=mid_channels, num_blocks_per_stage=4, + spatial_scale=spatial_scale, + rational_resampler=rational_resampler, ) # Sanitize weights - convert from PyTorch to MLX format @@ -349,7 +426,7 @@ def load_upsampler(weights_path: str) -> LatentUpsampler: for key, value in raw_weights.items(): new_key = key - # LTX-2.3 upsampler uses sequential indexing: upsampler.0.* -> upsampler.conv.* + # x2 upsampler uses sequential indexing: upsampler.0.* -> upsampler.conv.* if key.startswith("upsampler.0."): new_key = key.replace("upsampler.0.", "upsampler.conv.") @@ -358,7 +435,7 @@ def load_upsampler(weights_path: str) -> LatentUpsampler: value = mx.transpose(value, (0, 2, 3, 4, 1)) # Conv2d weights: PyTorch (O, I, H, W) -> MLX (O, H, W, I) - if "weight" in new_key and value.ndim == 4: + if ("weight" in new_key or "kernel" in new_key) and value.ndim == 4: value = mx.transpose(value, (0, 2, 3, 1)) sanitized[new_key] = value @@ -368,4 +445,4 @@ def load_upsampler(weights_path: str) -> LatentUpsampler: print(f" Loaded {len(sanitized)} weights") - return upsampler + return upsampler, spatial_scale