tmp/video-velocity-model/data_processing/flux2_ae/wrapper.py

from __future__ import annotations

from pathlib import Path

import torch
from safetensors.torch import load_file as load_safetensors
from torch import nn

from .autoencoder import AutoEncoder, AutoEncoderParams


class Flux2AutoEncoderWrapper(nn.Module):
    """Small adapter exposing FLUX.2 AE with the same encode/decode API as RAE."""

    def __init__(self, weight_path: str | Path):
        super().__init__()
        weight_path = Path(weight_path).expanduser()
        if not weight_path.exists():
            raise FileNotFoundError(f"FLUX.2 AE checkpoint not found: {weight_path}")
        with torch.device("meta"):
            self.ae = AutoEncoder(AutoEncoderParams())
        state_dict = load_safetensors(str(weight_path), device="cpu")
        self.ae.load_state_dict(state_dict, strict=True, assign=True)
        self.weight_path = str(weight_path)
        self.encoder_input_size = 256

    @torch.no_grad()
    def encode(self, x: torch.Tensor) -> torch.Tensor:
        # The surrounding pipeline provides [0, 1] RGB tensors; FLUX.2 AE expects [-1, 1].
        dtype = next(self.ae.parameters()).dtype
        z = self.ae.encode(x.to(dtype=dtype).mul(2.0).sub(1.0))
        return z.float()

    @torch.no_grad()
    def decode(self, z: torch.Tensor) -> torch.Tensor:
        dtype = next(self.ae.parameters()).dtype
        x = self.ae.decode(z.to(dtype=dtype)).float()
        return x.add(1.0).mul(0.5)