foleycrafter/models/specvqgan/modules/misc/raw_feats.py

import torch

class RawFeatsStage(object):
    def __init__(self):
        pass

    def eval(self):
        return self

    def encode(self, c):
        """fake vqmodel interface because self.cond_stage_model should have something
        similar to coord.py but even more `dummy`"""
        # assert 0.0 <= c.min() and c.max() <= 1.0
        info = None, None, c
        return c, None, info

    def decode(self, c):
        return c

    def get_input(self, batch, k):
        x = batch[k]
        x = x.permute(0, 2, 1).to(memory_format=torch.contiguous_format)
        return x.float()