backend/nn/base.py

import torch

from diffusers.configuration_utils import ConfigMixin, register_to_config
from torch import nn


class ModuleDict(torch.nn.Module):
    def __init__(self, module_dict):
        super(ModuleDict, self).__init__()
        for name, module in module_dict.items():
            self.add_module(name, module)


class ObjectDict:
    def __init__(self, module_dict):
        for name, module in module_dict.items():
            setattr(self, name, module)


class Dummy(nn.Module, ConfigMixin):
    config_name = 'config.json'

    @register_to_config
    def __init__(self):
        super().__init__()