fourm/utils/__init__.py

"""Minimal, torchvision-free fourm.utils for the MODUS HF Space.

The upstream fourm/utils/__init__.py is a heavy re-export hub (misc, timm, clip,
s3, logger ...) that imports torchvision/timm — which cannot be installed on the
ZeroGPU Space (custom torch 2.11, no matching torchvision). The fourm VQVAE
inference path only needs `to_2tuple` and `denormalize`, so we provide those two
here directly (torch-native), avoiding the heavy imports.
"""
import collections.abc
from itertools import repeat

import torch


# timm-style tuple helper (matches fourm.utils.misc: to_2tuple = _ntuple(2)).
def _ntuple(n):
    def parse(x):
        if isinstance(x, collections.abc.Iterable) and not isinstance(x, str):
            return tuple(x)
        return tuple(repeat(x, n))
    return parse


to_2tuple = _ntuple(2)

IMAGENET_DEFAULT_MEAN = (0.485, 0.456, 0.406)
IMAGENET_DEFAULT_STD = (0.229, 0.224, 0.225)


def denormalize(img, mean=IMAGENET_DEFAULT_MEAN, std=IMAGENET_DEFAULT_STD):
    """Inverse of torchvision Normalize: img * std + mean (per channel, CxHxW)."""
    m = torch.as_tensor(mean, device=img.device, dtype=img.dtype).view(-1, 1, 1)
    s = torch.as_tensor(std, device=img.device, dtype=img.dtype).view(-1, 1, 1)
    return img * s + m