n0p/mlx-video
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Add custom spatial upscaling support to LTX-2 video generation; introduce spatial_upscaler parameter and enhance resolution handling for two-stage pipelines

Browse Source
This commit is contained in:
Prince Canuma
2026-03-17 02:23:47 +01:00
parent cc302d79b0
commit 57f66bcae2
3 changed files with 234 additions and 98 deletions
Download Patch File Download Diff File Expand all files Collapse all files

38
mlx_video/models/ltx_2/README.md
View File

@@ -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).

101
mlx_video/models/ltx_2/generate.py
View File

@@ -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,
)

193
mlx_video/models/ltx_2/upsampler.py
View File

@@ -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