ltx2_upsampler

class PixelShuffleND(nn.Module):
    """N-dimensional pixel shuffle for upsampling."""
def __init__(self, dims: int, upscale_factors: Tuple[int, int, int] = (2, 2, 2)) -> None:
    super().__init__()
    if dims not in (1, 2, 3):
        raise ValueError("dims must be 1, 2, or 3")
    self.dims = dims
    self.upscale_factors = upscale_factors

def forward(self, x: torch.Tensor) -> torch.Tensor:
    if self.dims == 3:
        return rearrange(
            x,
            "b (c p1 p2 p3) d h w -> b c (d p1) (h p2) (w p3)",
            p1=self.upscale_factors[0],
            p2=self.upscale_factors[1],
            p3=self.upscale_factors[2],
        )
    if self.dims == 2:
        return rearrange(
            x,
            "b (c p1 p2) h w -> b c (h p1) (w p2)",
            p1=self.upscale_factors[0],
            p2=self.upscale_factors[1],
        )
    if self.dims == 1:
        return rearrange(
            x,
            "b (c p1) f h w -> b c (f p1) h w",
            p1=self.upscale_factors[0],
        )
    raise ValueError(f&quot;Unsupported dims: &#123;self.dims&#125;&quot;)</code></pre></details>
fastvideo.models.upsamplers.ltx2_upsampler.PixelShuffleND.forward
forward(x: torch.Tensor) -> torch.Tensor
Source
def forward(self, x: torch.Tensor) -> torch.Tensor:
    if self.dims == 3:
        return rearrange(
            x,
            "b (c p1 p2 p3) d h w -> b c (d p1) (h p2) (w p3)",
            p1=self.upscale_factors[0],
            p2=self.upscale_factors[1],
            p3=self.upscale_factors[2],
        )
    if self.dims == 2:
        return rearrange(
            x,
            "b (c p1 p2) h w -> b c (h p1) (w p2)",
            p1=self.upscale_factors[0],
            p2=self.upscale_factors[1],
        )
    if self.dims == 1:
        return rearrange(
            x,
            "b (c p1) f h w -> b c (f p1) h w",
            p1=self.upscale_factors[0],
        )
    raise ValueError(f"Unsupported dims: {self.dims}")
fastvideo.models.upsamplers.ltx2_upsampler.BlurDownsample
BlurDownsample(dims: int, stride: int, kernel_size: int = 5)
Anti-aliased spatial downsampling by integer stride using a fixed separable binomial kernel.
Applies only on H,W. Works for dims=2 or dims=3 (per-frame).
Source
class BlurDownsample(nn.Module):
    """
    Anti-aliased spatial downsampling by integer stride using a fixed separable binomial kernel.
    Applies only on H,W. Works for dims=2 or dims=3 (per-frame).
    """
def __init__(self, dims: int, stride: int, kernel_size: int = 5) -&gt; None:
    super().__init__()
    if dims not in (2, 3):
        raise ValueError(&quot;dims must be 2 or 3&quot;)
    if stride &lt; 1:
        raise ValueError(&quot;stride must be &gt;= 1&quot;)
    if kernel_size &lt; 3 or kernel_size % 2 != 1:
        raise ValueError(&quot;kernel_size must be an odd integer &gt;= 3&quot;)

    self.dims = dims
    self.stride = stride
    self.kernel_size = kernel_size

    k = torch.tensor(&#91;math.comb(kernel_size - 1, idx) for idx in range(kernel_size)&#93;)
    k2d = k&#91;:, None&#93; @ k&#91;None, :&#93;
    k2d = (k2d / k2d.sum()).float()
    self.register_buffer(&quot;kernel&quot;, k2d&#91;None, None, :, :&#93;)

def forward(self, x: torch.Tensor) -&gt; torch.Tensor:
    if self.stride == 1:
        return x
    if self.dims == 2:
        return self._apply_2d(x)
    b, _, f, _, _ = x.shape
    x = rearrange(x, &quot;b c f h w -&gt; (b f) c h w&quot;)
    x = self._apply_2d(x)
    h2, w2 = x.shape&#91;-2:&#93;
    return rearrange(x, &quot;(b f) c h w -&gt; b c f h w&quot;, b=b, f=f, h=h2, w=w2)

def _apply_2d(self, x2d: torch.Tensor) -&gt; torch.Tensor:
    c = x2d.shape&#91;1&#93;
    weight = self.kernel.expand(c, 1, self.kernel_size, self.kernel_size)
    return nn.functional.conv2d(
        x2d,
        weight=weight,
        bias=None,
        stride=self.stride,
        padding=self.kernel_size // 2,
        groups=c,
    )</code></pre></details>
fastvideo.models.upsamplers.ltx2_upsampler.BlurDownsample.forward
forward(x: torch.Tensor) -> torch.Tensor
Source
def forward(self, x: torch.Tensor) -> torch.Tensor:
    if self.stride == 1:
        return x
    if self.dims == 2:
        return self._apply_2d(x)
    b, _, f, _, _ = x.shape
    x = rearrange(x, "b c f h w -> (b f) c h w")
    x = self._apply_2d(x)
    h2, w2 = x.shape[-2:]
    return rearrange(x, "(b f) c h w -> b c f h w", b=b, f=f, h=h2, w=w2)
fastvideo.models.upsamplers.ltx2_upsampler.SpatialRationalResampler
SpatialRationalResampler(mid_channels: int, scale: float)
Fully-learned rational spatial scaling: up by 'num' via PixelShuffle, then
anti-aliased downsample by 'den' using fixed blur + stride. Operates on H,W only.
For dims==3, work per-frame for spatial scaling (temporal axis untouched).
Source
class SpatialRationalResampler(nn.Module):
    """
    Fully-learned rational spatial scaling: up by 'num' via PixelShuffle, then
    anti-aliased downsample by 'den' using fixed blur + stride. Operates on H,W only.
    For dims==3, work per-frame for spatial scaling (temporal axis untouched).
    """
def __init__(self, mid_channels: int, scale: float) -&gt; None:
    super().__init__()
    self.scale = float(scale)
    self.num, self.den = _rational_for_scale(self.scale)
    self.conv = nn.Conv2d(mid_channels, (self.num**2) * mid_channels, kernel_size=3, padding=1)
    self.pixel_shuffle = PixelShuffleND(2, upscale_factors=(self.num, self.num))
    self.blur_down = BlurDownsample(dims=2, stride=self.den)

def forward(self, x: torch.Tensor) -&gt; torch.Tensor:
    b, _, f, _, _ = x.shape
    x = rearrange(x, &quot;b c f h w -&gt; (b f) c h w&quot;)
    x = self.conv(x)
    x = self.pixel_shuffle(x)
    x = self.blur_down(x)
    return rearrange(x, &quot;(b f) c h w -&gt; b c f h w&quot;, b=b, f=f)</code></pre></details>
fastvideo.models.upsamplers.ltx2_upsampler.SpatialRationalResampler.forward
forward(x: torch.Tensor) -> torch.Tensor
Source
def forward(self, x: torch.Tensor) -> torch.Tensor:
    b, _, f, _, _ = x.shape
    x = rearrange(x, "b c f h w -> (b f) c h w")
    x = self.conv(x)
    x = self.pixel_shuffle(x)
    x = self.blur_down(x)
    return rearrange(x, "(b f) c h w -> b c f h w", b=b, f=f)
fastvideo.models.upsamplers.ltx2_upsampler.ResBlock
ResBlock(channels: int, mid_channels: Optional[int] = None, dims: int = 3)
Residual block with two convolutional layers, group norm, and SiLU.
Source
class ResBlock(nn.Module):
    """Residual block with two convolutional layers, group norm, and SiLU."""
def __init__(self, channels: int, mid_channels: Optional&#91;int&#93; = None, dims: int = 3) -&gt; None:
    super().__init__()
    if mid_channels is None:
        mid_channels = channels

    conv = nn.Conv2d if dims == 2 else nn.Conv3d

    self.conv1 = conv(channels, mid_channels, kernel_size=3, padding=1)
    self.norm1 = nn.GroupNorm(32, mid_channels)
    self.conv2 = conv(mid_channels, channels, kernel_size=3, padding=1)
    self.norm2 = nn.GroupNorm(32, channels)
    self.activation = nn.SiLU()

def forward(self, x: torch.Tensor) -&gt; torch.Tensor:
    residual = x
    x = self.conv1(x)
    x = self.norm1(x)
    x = self.activation(x)
    x = self.conv2(x)
    x = self.norm2(x)
    x = self.activation(x + residual)
    return x</code></pre></details>
fastvideo.models.upsamplers.ltx2_upsampler.ResBlock.forward
forward(x: torch.Tensor) -> torch.Tensor
Source
def forward(self, x: torch.Tensor) -> torch.Tensor:
    residual = x
    x = self.conv1(x)
    x = self.norm1(x)
    x = self.activation(x)
    x = self.conv2(x)
    x = self.norm2(x)
    x = self.activation(x + residual)
    return x
fastvideo.models.upsamplers.ltx2_upsampler.LatentUpsampler
LatentUpsampler(in_channels: int = 128, mid_channels: int = 512, num_blocks_per_stage: int = 4, dims: int = 3, spatial_upsample: bool = True, temporal_upsample: bool = False, spatial_scale: float = 2.0, rational_resampler: bool = False)
Model to upsample VAE latents spatially and/or temporally.
Source
class LatentUpsampler(nn.Module):
    """
    Model to upsample VAE latents spatially and/or temporally.
    """
def __init__(
    self,
    in_channels: int = 128,
    mid_channels: int = 512,
    num_blocks_per_stage: int = 4,
    dims: int = 3,
    spatial_upsample: bool = True,
    temporal_upsample: bool = False,
    spatial_scale: float = 2.0,
    rational_resampler: bool = False,
) -&gt; None:
    super().__init__()

    self.in_channels = in_channels
    self.mid_channels = mid_channels
    self.num_blocks_per_stage = num_blocks_per_stage
    self.dims = dims
    self.spatial_upsample = spatial_upsample
    self.temporal_upsample = temporal_upsample
    self.spatial_scale = float(spatial_scale)
    self.rational_resampler = rational_resampler

    conv = nn.Conv2d if dims == 2 else nn.Conv3d

    self.initial_conv = conv(in_channels, mid_channels, kernel_size=3, padding=1)
    self.initial_norm = nn.GroupNorm(32, mid_channels)
    self.initial_activation = nn.SiLU()

    self.res_blocks = nn.ModuleList(&#91;ResBlock(mid_channels, dims=dims) for _ in range(num_blocks_per_stage)&#93;)

    if spatial_upsample and temporal_upsample:
        self.upsampler = nn.Sequential(
            nn.Conv3d(mid_channels, 8 * mid_channels, kernel_size=3, padding=1),
            PixelShuffleND(3),
        )
    elif spatial_upsample:
        if rational_resampler:
            self.upsampler = SpatialRationalResampler(mid_channels=mid_channels, scale=self.spatial_scale)
        else:
            self.upsampler = nn.Sequential(
                nn.Conv2d(mid_channels, 4 * mid_channels, kernel_size=3, padding=1),
                PixelShuffleND(2),
            )
    elif temporal_upsample:
        self.upsampler = nn.Sequential(
            nn.Conv3d(mid_channels, 2 * mid_channels, kernel_size=3, padding=1),
            PixelShuffleND(1),
        )
    else:
        raise ValueError(&quot;Either spatial_upsample or temporal_upsample must be True&quot;)

    self.post_upsample_res_blocks = nn.ModuleList(
        &#91;ResBlock(mid_channels, dims=dims) for _ in range(num_blocks_per_stage)&#93;)

    self.final_conv = conv(mid_channels, in_channels, kernel_size=3, padding=1)

def forward(self, latent: torch.Tensor) -&gt; torch.Tensor:
    b, _, f, _, _ = latent.shape

    if self.dims == 2:
        x = rearrange(latent, &quot;b c f h w -&gt; (b f) c h w&quot;)
        x = self.initial_conv(x)
        x = self.initial_norm(x)
        x = self.initial_activation(x)

        for block in self.res_blocks:
            x = block(x)

        x = self.upsampler(x)

        for block in self.post_upsample_res_blocks:
            x = block(x)

        x = self.final_conv(x)
        x = rearrange(x, &quot;(b f) c h w -&gt; b c f h w&quot;, b=b, f=f)
    else:
        x = self.initial_conv(latent)
        x = self.initial_norm(x)
        x = self.initial_activation(x)

        for block in self.res_blocks:
            x = block(x)

        if self.temporal_upsample:
            x = self.upsampler(x)
            x = x&#91;:, :, 1:, :, :&#93;
        elif isinstance(self.upsampler, SpatialRationalResampler):
            x = self.upsampler(x)
        else:
            x = rearrange(x, &quot;b c f h w -&gt; (b f) c h w&quot;)
            x = self.upsampler(x)
            x = rearrange(x, &quot;(b f) c h w -&gt; b c f h w&quot;, b=b, f=f)

        for block in self.post_upsample_res_blocks:
            x = block(x)

        x = self.final_conv(x)

    return x</code></pre></details>
fastvideo.models.upsamplers.ltx2_upsampler.LatentUpsampler.forward
forward(latent: torch.Tensor) -> torch.Tensor
Source
def forward(self, latent: torch.Tensor) -> torch.Tensor:
    b, _, f, _, _ = latent.shape
if self.dims == 2:
    x = rearrange(latent, &quot;b c f h w -&gt; (b f) c h w&quot;)
    x = self.initial_conv(x)
    x = self.initial_norm(x)
    x = self.initial_activation(x)

    for block in self.res_blocks:
        x = block(x)

    x = self.upsampler(x)

    for block in self.post_upsample_res_blocks:
        x = block(x)

    x = self.final_conv(x)
    x = rearrange(x, &quot;(b f) c h w -&gt; b c f h w&quot;, b=b, f=f)
else:
    x = self.initial_conv(latent)
    x = self.initial_norm(x)
    x = self.initial_activation(x)

    for block in self.res_blocks:
        x = block(x)

    if self.temporal_upsample:
        x = self.upsampler(x)
        x = x&#91;:, :, 1:, :, :&#93;
    elif isinstance(self.upsampler, SpatialRationalResampler):
        x = self.upsampler(x)
    else:
        x = rearrange(x, &quot;b c f h w -&gt; (b f) c h w&quot;)
        x = self.upsampler(x)
        x = rearrange(x, &quot;(b f) c h w -&gt; b c f h w&quot;, b=b, f=f)

    for block in self.post_upsample_res_blocks:
        x = block(x)

    x = self.final_conv(x)

return x</code></pre></details>
fastvideo.models.upsamplers.ltx2_upsampler.LatentUpsamplerConfigurator
LatentUpsamplerConfigurator()
Configurator for LatentUpsampler from a config dict.
Source
class LatentUpsamplerConfigurator:
    """Configurator for LatentUpsampler from a config dict."""
@classmethod
def from_config(cls, config: dict&#91;str, Any&#93;) -&gt; LatentUpsampler:
    cfg = dict(config)
    cfg.pop(&quot;_class_name&quot;, None)
    if &quot;upsampler&quot; in cfg and isinstance(cfg&#91;&quot;upsampler&quot;&#93;, dict):
        cfg = cfg&#91;&quot;upsampler&quot;&#93;

    return LatentUpsampler(
        in_channels=cfg.get(&quot;in_channels&quot;, 128),
        mid_channels=cfg.get(&quot;mid_channels&quot;, 512),
        num_blocks_per_stage=cfg.get(&quot;num_blocks_per_stage&quot;, 4),
        dims=cfg.get(&quot;dims&quot;, 3),
        spatial_upsample=cfg.get(&quot;spatial_upsample&quot;, True),
        temporal_upsample=cfg.get(&quot;temporal_upsample&quot;, False),
        spatial_scale=cfg.get(&quot;spatial_scale&quot;, 2.0),
        rational_resampler=cfg.get(&quot;rational_resampler&quot;, False),
    )</code></pre></details>
fastvideo.models.upsamplers.ltx2_upsampler.LatentUpsamplerConfigurator.from_config
from_config(config: dict[str, Any]) -> LatentUpsampler
Source
@classmethod
def from_config(cls, config: dict[str, Any]) -> LatentUpsampler:
    cfg = dict(config)
    cfg.pop("_class_name", None)
    if "upsampler" in cfg and isinstance(cfg["upsampler"], dict):
        cfg = cfg["upsampler"]
return LatentUpsampler(
    in_channels=cfg.get(&quot;in_channels&quot;, 128),
    mid_channels=cfg.get(&quot;mid_channels&quot;, 512),
    num_blocks_per_stage=cfg.get(&quot;num_blocks_per_stage&quot;, 4),
    dims=cfg.get(&quot;dims&quot;, 3),
    spatial_upsample=cfg.get(&quot;spatial_upsample&quot;, True),
    temporal_upsample=cfg.get(&quot;temporal_upsample&quot;, False),
    spatial_scale=cfg.get(&quot;spatial_scale&quot;, 2.0),
    rational_resampler=cfg.get(&quot;rational_resampler&quot;, False),
)</code></pre></details>
fastvideo.models.upsamplers.ltx2_upsampler.LTX2LatentUpsampler
LTX2LatentUpsampler(config: dict[str, Any])
Public wrapper for the LTX-2 latent upsampler.
Source
class LTX2LatentUpsampler(nn.Module):
    """Public wrapper for the LTX-2 latent upsampler."""
def __init__(self, config: dict&#91;str, Any&#93;):
    super().__init__()
    self.model: LatentUpsampler = LatentUpsamplerConfigurator.from_config(config)

def forward(self, latent: torch.Tensor) -&gt; torch.Tensor:
    return self.model(latent)</code></pre></details>
fastvideo.models.upsamplers.ltx2_upsampler.LTX2LatentUpsampler.forward
forward(latent: torch.Tensor) -> torch.Tensor
Source
def forward(self, latent: torch.Tensor) -> torch.Tensor:
    return self.model(latent)
fastvideo.models.upsamplers.ltx2_upsampler.upsample_video
upsample_video(latent: torch.Tensor, video_encoder: Any, upsampler: LatentUpsampler) -> torch.Tensor
Upsample a latent tensor with normalization based on the video encoder's per-channel statistics.
Source
def upsample_video(latent: torch.Tensor, video_encoder: Any, upsampler: LatentUpsampler) -> torch.Tensor:
    """
    Upsample a latent tensor with normalization based on the video encoder's per-channel statistics.
    """
    if not hasattr(video_encoder, "per_channel_statistics"):
        raise ValueError("video_encoder must expose per_channel_statistics for normalization")
    stats = video_encoder.per_channel_statistics
    latent = stats.un_normalize(latent)
    latent = upsampler(latent)
    latent = stats.normalize(latent)
    return latent