konfai 1.0.0__py3-none-any.whl

konfai/models/classification/convNeXt.py

@@ -0,0 +1,175 @@
+import torch
+import torch.nn.functional as F
+from KonfAI.konfai.network import network, blocks
+from KonfAI.konfai.utils.config import config
+from KonfAI.konfai.data.HDF5 import ModelPatch
+
+"""
+"convnext_tiny_1k": "https://dl.fbaipublicfiles.com/convnext/convnext_tiny_1k_224_ema.pth", depths=[3, 3, 9, 3], dims=[96, 192, 384, 768]
+"convnext_small_1k": "https://dl.fbaipublicfiles.com/convnext/convnext_small_1k_224_ema.pth", depths=[3, 3, 27, 3], dims=[96, 192, 384, 768]
+"convnext_base_1k": "https://dl.fbaipublicfiles.com/convnext/convnext_base_1k_224_ema.pth", depths=[3, 3, 27, 3], dims=[128, 256, 512, 1024]
+"convnext_large_1k": "https://dl.fbaipublicfiles.com/convnext/convnext_large_1k_224_ema.pth",[3, 3, 27, 3], dims=[192, 384, 768, 1536]
+"convnext_tiny_22k": "https://dl.fbaipublicfiles.com/convnext/convnext_tiny_22k_224.pth",
+"convnext_small_22k": "https://dl.fbaipublicfiles.com/convnext/convnext_small_22k_224.pth",
+"convnext_base_22k": "https://dl.fbaipublicfiles.com/convnext/convnext_base_22k_224.pth",
+"convnext_large_22k": "https://dl.fbaipublicfiles.com/convnext/convnext_large_22k_224.pth",
+"convnext_xlarge_22k": "https://dl.fbaipublicfiles.com/convnext/convnext_xlarge_22k_224.pth", depths=[3, 3, 27, 3], dims=[256, 512, 1024, 2048]
+"""
+
+class LayerNorm(torch.nn.Module):
+
+    def __init__(self, normalized_shape : int, eps : float = 1e-6, data_format : str="channels_last"):
+        super().__init__()
+        self.weight = torch.nn.parameter.Parameter(torch.ones(normalized_shape))
+        self.bias = torch.nn.parameter.Parameter(torch.zeros(normalized_shape))
+        self.eps = eps
+        self.data_format = data_format
+        if self.data_format not in ["channels_last", "channels_first"]:
+            raise NotImplementedError 
+        self.normalized_shape = (normalized_shape, )
+    
+    def forward(self, x):
+        if self.data_format == "channels_last":
+            return F.layer_norm(x, self.normalized_shape, self.weight, self.bias, self.eps)
+        elif self.data_format == "channels_first":
+            u = x.mean(1, keepdim=True)
+            s = (x - u).pow(2).mean(1, keepdim=True)
+            x = (x - u) / torch.sqrt(s + self.eps)
+            if len(x.shape) == 3:
+                x = self.weight[:, None] * x + self.bias[:, None]
+            elif len(x.shape) == 4:
+                x = self.weight[:, None, None] * x + self.bias[:, None, None]
+            else:
+                x = self.weight[:, None, None, None] * x + self.bias[:, None, None, None]
+            return x
+
+    def extra_repr(self):
+        return "normalized_shape={}, eps={}, data_format={})".format(self.normalized_shape, self.eps, self.data_format)
+
+class DropPath(torch.nn.Module):
+
+    def __init__(self, drop_prob: float = 0., scale_by_keep: bool = True):
+        super(DropPath, self).__init__()
+        self.drop_prob = drop_prob
+        self.scale_by_keep = scale_by_keep
+
+    def forward(self, x):
+        if self.drop_prob == 0. or not self.training:
+            return x
+        keep_prob = 1 - self.drop_prob
+        shape = (x.shape[0],) + (1,) * (x.ndim - 1)
+        random_tensor = x.new_empty(shape).bernoulli_(keep_prob)
+        if keep_prob > 0.0 and self.scale_by_keep:
+            random_tensor.div_(keep_prob)
+        return x * random_tensor
+
+    def extra_repr(self):
+        return "drop_prob={}".format(round(self.drop_prob,3))
+
+class LayerScaler(torch.nn.Module):
+    
+    def __init__(self, init_value : float, dimensions : int):
+        super().__init__()
+        self.init_value = init_value
+        self.gamma = torch.nn.Parameter(torch.ones(dimensions, 1, 1) * init_value)
+
+    def forward(self, input : torch.Tensor) -> torch.Tensor:
+        return self.gamma * input
+
+    def extra_repr(self):
+        return "init_value={}".format(self.init_value)
+
+class BottleNeckBlock(network.ModuleArgsDict):
+    def __init__(   self,
+                    features: int,
+                    drop_p: float,
+                    layer_scaler_init_value: float,
+                    dim: int):
+        super().__init__()
+        self.add_module("Conv_0", blocks.getTorchModule("Conv", dim)(features, features, kernel_size=7, padding=3, groups=features), alias=["dwconv"])
+        self.add_module("ToFeatures", blocks.ToFeatures(dim))
+        self.add_module("LayerNorm", LayerNorm(features, eps=1e-6), alias=["norm"])
+        self.add_module("Linear_1",  torch.nn.Linear(features, features * 4), alias=["pwconv1"])
+        self.add_module("GELU", torch.nn.GELU())
+        self.add_module("Linear_2", torch.nn.Linear(features * 4, features), alias=["pwconv2"])
+        self.add_module("ToChannels", blocks.ToChannels(dim))
+        self.add_module("LayerScaler", LayerScaler(init_value=layer_scaler_init_value, dimensions=features), alias=[""])
+        self.add_module("StochasticDepth", DropPath(drop_p))
+        self.add_module("Residual", blocks.Add(), in_branch=[0,1])
+        
+class DownSample(network.ModuleArgsDict):
+    
+    def __init__(   self, 
+                    in_features: int,
+                    out_features: int,
+                    dim : int):
+        super().__init__()
+        self.add_module("LayerNorm", LayerNorm(in_features, eps=1e-6, data_format="channels_first"), alias=["0"])
+        self.add_module("Conv", blocks.getTorchModule("Conv", dim)(in_features, out_features, kernel_size=2, stride=2), alias=["1"])
+
+class ConvNexStage(network.ModuleArgsDict):
+    
+    def __init__(   self, 
+                    features: int,
+                    depth: int, 
+                    drop_p: list[float],
+                    dim : int):
+        super().__init__()
+        for i in range(depth):
+            self.add_module("BottleNeckBlock_{}".format(i), BottleNeckBlock(features=features, drop_p=drop_p[i], layer_scaler_init_value=1e-6, dim=dim), alias=["{}".format(i)])
+
+class ConvNextStem(network.ModuleArgsDict):
+
+    def __init__(self, in_features: int, out_features: int, dim: int):
+        super().__init__()
+        self.add_module("Conv", blocks.getTorchModule("Conv", dim)(in_features, out_features, kernel_size=4, stride=4), alias=["0"])
+        self.add_module("LayerNorm", LayerNorm(out_features, eps=1e-6, data_format="channels_first"), alias=["1"])
+
+class ConvNextEncoder(network.ModuleArgsDict):
+
+    def __init__(   self,
+                    in_channels: int,
+                    depths: list[int],
+                    widths: list[int],
+                    drop_p: float,
+                    dim : int):
+        super().__init__()
+        self.add_module("ConvNextStem", ConvNextStem(in_channels, widths[0], dim=dim), alias=["downsample_layers.0"])
+        
+        drop_probs = [x.item() for x in torch.linspace(0, drop_p, sum(depths))]
+        self.add_module("ConvNexStage_0", ConvNexStage(features=widths[0], depth=depths[0], drop_p=drop_probs[:depths[0]], dim=dim), alias=["stages.0"])
+        
+        for i, (in_features, out_features) in enumerate(list(zip(widths[:], widths[1:]))):
+            self.add_module("DownSample_{}".format(i+1), DownSample(in_features=in_features, out_features=out_features, dim=dim), alias=["downsample_layers.{}".format(i+1)])
+            self.add_module("ConvNexStage_{}".format(i+1), ConvNexStage(features=out_features, depth=depths[i+1], drop_p=drop_probs[sum(depths[:i+1]):sum(depths[:i+2])], dim=dim), alias=["stages.{}".format(i+1)])
+
+class Head(network.ModuleArgsDict):
+
+    def __init__(self, in_features : int, num_classes : list[int], dim : int) -> None:
+        super().__init__()
+        self.add_module("AdaptiveAvgPool", blocks.getTorchModule("AdaptiveAvgPool", dim)(tuple([1]*dim)))
+        self.add_module("Flatten", torch.nn.Flatten(1))
+        self.add_module("LayerNorm", torch.nn.LayerNorm(in_features, eps=1e-6), alias=["norm"])
+        
+        for i, nb_classe in enumerate(num_classes):
+            self.add_module("Linear_{}".format(i), torch.nn.Linear(in_features, nb_classe), pretrained=False, alias=["head"], out_branch=[i+1])
+            self.add_module("Unsqueeze_{}".format(i), blocks.Unsqueeze(2), in_branch=[i+1], out_branch=[-1])
+
+class ConvNeXt(network.Network):
+    
+    @config("ConvNeXt")
+    def __init__(   self,
+                    optimizer : network.OptimizerLoader = network.OptimizerLoader(),
+                    schedulers : network.LRSchedulersLoader = network.LRSchedulersLoader(),
+                    outputsCriterions: dict[str, network.TargetCriterionsLoader] = {"default" : network.TargetCriterionsLoader()},
+                    patch : ModelPatch = ModelPatch(),
+                    dim : int = 3,
+                    in_channels: int = 1,
+                    depths: list[int] = [3,3,27,3],
+                    widths: list[int] = [128, 256, 512, 1024],
+                    drop_p: float = 0.1,
+                    num_classes: list[int] = [4, 7]):
+
+        super().__init__(in_channels = in_channels, optimizer = optimizer, schedulers = schedulers, outputsCriterions = outputsCriterions, dim = dim, patch=patch, init_type = "trunc_normal", init_gain=0.02)
+        self.add_module("ConvNextEncoder", ConvNextEncoder(in_channels=in_channels, depths=depths, widths=widths, drop_p=drop_p, dim=dim))        
+        self.add_module("Head", Head(in_features=widths[-1], num_classes=num_classes, dim=dim))

konfai/models/classification/resnet.py

@@ -0,0 +1,116 @@
+from abc import ABC
+from typing import Type
+import torch
+from KonfAI.konfai.network import network, blocks
+from KonfAI.konfai.utils.config import config
+from KonfAI.konfai.data.HDF5 import ModelPatch
+
+"""
+'resnet18': 'https://download.pytorch.org/models/resnet18-5c106cde.pth', dim = 2, in_channels = 3, depths=[2, 2, 2, 2], widths = [64, 64, 128, 256, 512],  num_classes=1000, useBottleneck=False
+
+'resnet34': 'https://download.pytorch.org/models/resnet34-333f7ec4.pth', dim = 2, in_channels = 3, depths=[3, 4, 6, 3], widths = [64, 64, 128, 256, 512],  num_classes=1000, useBottleneck=False
+
+'resnet50': 'https://download.pytorch.org/models/resnet50-19c8e357.pth', dim = 2, in_channels = 3, depths=[3, 4, 6, 3], widths = [64, 256, 512, 1024, 2048],  num_classes=1000, useBottleneck=True
+
+'resnet101': 'https://download.pytorch.org/models/resnet101-5d3b4d8f.pth', dim = 2, in_channels = 3, depths=[3, 4, 23, 3], widths = [64, 256, 512, 1024, 2048],  num_classes=1000, useBottleneck=True
+
+'resnet152': 'https://download.pytorch.org/models/resnet152-b121ed2d.pth', dim = 2, in_channels = 3, depths=[3, 8, 36, 3], widths = [64, 256, 512, 1024, 2048],  num_classes=1000, useBottleneck=True
+
+'resnext50_32x4d': 'https://download.pytorch.org/models/resnext50_32x4d-7cdf4587.pth',
+'resnext101_32x8d': 'https://download.pytorch.org/models/resnext101_32x8d-8ba56ff5.pth',
+'wide_resnet50_2': 'https://download.pytorch.org/models/wide_resnet50_2-95faca4d.pth',
+'wide_resnet101_2': 'https://download.pytorch.org/models/wide_resnet101_2-32ee1156.pth',
+"""
+class AbstractResBlock(network.ModuleArgsDict, ABC):
+
+    def __init__(self, in_channels: int, out_channels: int, downsample: bool, dim: int):
+        super().__init__()
+
+class ResBlock(AbstractResBlock):
+
+    def __init__(self, in_channels: int, out_channels: int, downsample: bool, dim: int):
+        super().__init__(in_channels, out_channels, downsample, dim)
+        if downsample:
+            self.add_module("Shortcut", blocks.ConvBlock(in_channels, out_channels, 1, blocks.BlockConfig(kernel_size=1, stride=2, padding=0, bias=False, activation="None", normMode=blocks.NormMode.BATCH.name), dim=dim, alias=[["0"], ["1"], []]), in_branch=[1], out_branch=[1], alias=["downsample"])     
+                   
+        self.add_module("ConvBlock_0", blocks.ConvBlock(in_channels, out_channels, 1, blocks.BlockConfig(kernel_size=3, stride=2 if downsample else 1, padding=1, bias=False, activation="ReLU", normMode=blocks.NormMode.BATCH.name), dim=dim, alias=[["conv1"], ["bn1"], []]))
+        self.add_module("ConvBlock_1", blocks.ConvBlock(out_channels, out_channels, 1, blocks.BlockConfig(kernel_size=3, stride=1, padding=1, bias=False, activation="None", normMode=blocks.NormMode.BATCH.name), dim=dim, alias=[["conv2"], ["bn2"], []]))
+        self.add_module("Residual", blocks.Add(), in_branch=[0,1])
+        self.add_module("ReLU", torch.nn.ReLU())
+
+class ResBottleneckBlock(AbstractResBlock):
+    def __init__(self, in_channels: int, out_channels: int, downsample: bool, dim: int):
+        super().__init__(in_channels, out_channels, downsample, dim)
+        self.add_module("ConvBlock_0", blocks.ConvBlock(in_channels, out_channels//4, 1, blocks.BlockConfig(kernel_size=1, stride=1, padding=0, bias=False, activation="ReLU", normMode=blocks.NormMode.BATCH.name), dim=dim, alias=[["conv1"], ["bn1"], []]))
+        self.add_module("ConvBlock_1", blocks.ConvBlock(out_channels//4, out_channels//4, 1, blocks.BlockConfig(kernel_size=3, stride=2 if downsample else 1, padding=1, bias=False, activation="ReLU", normMode=blocks.NormMode.BATCH.name), dim=dim, alias=[["conv2"], ["bn2"], []]))
+        self.add_module("ConvBlock_2", blocks.ConvBlock(out_channels//4, out_channels, 1, blocks.BlockConfig(kernel_size=1, stride=1, padding=0, bias=False, activation="ReLU", normMode=blocks.NormMode.BATCH.name), dim=dim, alias=[["conv3"], ["bn3"], []]))
+        
+        if downsample or in_channels != out_channels:
+            self.add_module("Shortcut", blocks.ConvBlock(in_channels, out_channels, 1, blocks.BlockConfig(kernel_size=1, stride=2 if downsample else 1, padding=0, bias=False, activation="None", normMode=blocks.NormMode.BATCH.name), dim=dim, alias=[["0"], ["1"], []]), in_branch=[1], out_branch=[1], alias=["downsample"])
+        
+        self.add_module("Residual", blocks.Add(), in_branch=[0,1])
+        self.add_module("ReLU", torch.nn.ReLU())
+
+class ResNetStage(network.ModuleArgsDict):
+
+    def __init__(self, 
+                    in_channels: int,
+                    out_channels: int,
+                    depth: int,
+                    block : Type[AbstractResBlock],
+                    downsample : bool,
+                    dim : int):
+        super().__init__()
+        self.add_module("BottleNeckBlock_0", block(in_channels=in_channels, out_channels=out_channels, downsample=downsample, dim=dim), alias=["0"])
+        for i in range(1, depth):
+            self.add_module("BottleNeckBlock_{}".format(i), block(in_channels=out_channels, out_channels=out_channels, downsample=False, dim=dim), alias=["{}".format(i)])
+
+class ResNetStem(network.ModuleArgsDict):
+
+    def __init__(self, in_channels: int, out_features: int, dim: int):
+        super().__init__()
+        self.add_module("ConvBlock", blocks.ConvBlock(in_channels, out_features, 1, blocks.BlockConfig(kernel_size=7, stride=2, padding=3, bias=False, activation="ReLU", normMode=blocks.NormMode.BATCH.name), dim=dim, alias=[["conv1"], ["bn1"], []]))
+        self.add_module("MaxPool", blocks.getTorchModule("MaxPool", dim)(kernel_size=3, stride=2, padding=1)) 
+
+class ResNetEncoder(network.ModuleArgsDict):
+    
+    def __init__(self,
+                    in_channels: int,
+                    depths: list[int],
+                    widths: list[int],
+                    useBottleneck : bool,
+                    dim : int):
+        super().__init__()
+        self.add_module("ResNetStem", ResNetStem(in_channels, widths[0], dim=dim))
+
+        for i, (in_channels, out_channels, depth) in enumerate(list(zip(widths[:], widths[1:], depths))):
+            self.add_module("ResNetStage_{}".format(i), ResNetStage(in_channels=in_channels, out_channels=out_channels, depth=depth, block=ResBottleneckBlock if useBottleneck else ResBlock, downsample=i!=0, dim=dim), alias=["layer{}".format(i+1)])
+
+class Head(network.ModuleArgsDict):
+    
+    def __init__(self, in_features : int, num_classes : int, dim : int):
+        super().__init__()
+        self.add_module("AdaptiveAvgPool", blocks.getTorchModule("AdaptiveAvgPool", dim)(1))
+        self.add_module("Flatten", torch.nn.Flatten(1))
+        self.add_module("Linear", torch.nn.Linear(in_features, num_classes), pretrained=False, alias=["fc"])
+        self.add_module("Unsqueeze", blocks.Unsqueeze(2))
+
+        
+class ResNet(network.Network):
+
+    @config("ResNet")
+    def __init__(   self,
+                    optimizer : network.OptimizerLoader = network.OptimizerLoader(),
+                    schedulers : network.LRSchedulersLoader = network.LRSchedulersLoader(),
+                    outputsCriterions: dict[str, network.TargetCriterionsLoader] = {"default" : network.TargetCriterionsLoader()},
+                    patch : ModelPatch = ModelPatch(),
+                    dim : int = 3,
+                    in_channels: int = 1,
+                    depths: list[int] = [2, 2, 2, 2],
+                    widths: list[int] = [64, 64, 128, 256, 512],
+                    num_classes: int = 10, 
+                    useBottleneck=False):
+        super().__init__(in_channels = in_channels, optimizer = optimizer, schedulers = schedulers, outputsCriterions = outputsCriterions, dim = dim, patch=patch, init_type = "trunc_normal", init_gain=0.02)
+        self.add_module("ResNetEncoder", ResNetEncoder(in_channels=in_channels, depths=depths, widths=widths, useBottleneck=useBottleneck, dim=dim))        
+        self.add_module("Head", Head(in_features=widths[-1], num_classes=num_classes, dim=dim))
+        

konfai/models/generation/cStyleGan.py

@@ -0,0 +1,137 @@
+import importlib
+import torch
+from KonfAI.konfai.network import network, blocks
+from KonfAI.konfai.utils.config import config
+from KonfAI.konfai.data.HDF5 import ModelPatch    
+class MappingNetwork(network.ModuleArgsDict):
+    def __init__(self, z_dim: int, c_dim: int, w_dim: int, num_layers: int, embed_features: int, layer_features: int):
+        super().__init__()
+        
+        self.add_module("Concat_1", blocks.Concat(), in_branch=[0,1])
+        
+        features = [z_dim + embed_features if c_dim > 0 else 0] + [layer_features] * (num_layers - 1) + [w_dim]    
+        if c_dim > 0:
+            self.add_module("Linear", torch.nn.Linear(c_dim, embed_features), out_branch=["Embed"])
+
+        self.add_module("Noise", blocks.NormalNoise(z_dim), in_branch=["Embed"])
+        if c_dim > 0:
+            self.add_module("Concat", blocks.Concat(), in_branch=[0,"Embed"])
+        
+        for i, (in_features, out_features) in enumerate(zip(features, features[1:])):
+            self.add_module("Linear_{}".format(i), torch.nn.Linear(in_features, out_features))
+    
+class ModulatedConv(torch.nn.Module):
+
+    class _ModulatedConv(torch.nn.Module):
+
+        def __init__(self, w_dim: int, conv: torch.nn.modules.conv._ConvNd, dim: int) -> None:
+            super().__init__()
+            self.affine = torch.nn.Linear(w_dim, conv.in_channels)
+            self.isConv = True
+            self.in_channels = conv.in_channels
+            self.out_channels = conv.out_channels
+            self.padding = conv.padding
+            self.stride = conv.stride
+            if isinstance(conv, torch.nn.modules.conv._ConvTransposeNd):
+                self.weight = torch.nn.parameter.Parameter(torch.randn((conv.in_channels, conv.out_channels, *conv.kernel_size)))
+                self.isConv = False
+            else:
+                self.weight = torch.nn.parameter.Parameter(torch.randn((conv.out_channels, conv.in_channels, *conv.kernel_size)))
+            conv.forward = self.forward
+            self.styles = None
+            self.dim = dim
+
+        def setStyle(self, styles: torch.Tensor) -> None:
+            self.styles = styles
+
+        def forward(self, input: torch.Tensor) -> torch.Tensor:
+            b = input.shape[0]
+            self.affine.to(input.device)
+            styles = self.affine(self.styles)
+            w1 = styles.reshape(b, -1, 1, *[1 for _ in range(self.dim)]) if not self.isConv else styles.reshape(b, 1, -1, *[1 for _ in range(self.dim)])
+            w2 = self.weight.unsqueeze(0).to(input.device)
+            weights = w2 * (w1 + 1)
+
+            d = torch.rsqrt((weights ** 2).sum(dim=tuple([i+2 for i in range(len(weights.shape)-2)]), keepdim=True) + 1e-8)
+            weights = weights * d
+            
+            input = input.reshape(1, -1, *input.shape[2:])
+
+            _, _, *ws = weights.shape
+            if not self.isConv:
+                out = getattr(importlib.import_module("torch.nn.functional"), "conv_transpose{}d".format(self.dim))(input, weights.reshape(b * self.in_channels, *ws), stride=self.stride, padding=self.padding, groups=b)
+            else:
+                out = getattr(importlib.import_module("torch.nn.functional"), "conv{}d".format(self.dim))(input, weights.reshape(b * self.out_channels, *ws), padding=self.padding, groups=b, stride=self.stride)
+            
+            out = out.reshape(-1, self.out_channels, *out.shape[2:])
+            return out
+    
+    def __init__(self, w_dim: int, module: torch.nn.Module) -> None:
+        super().__init__()
+        self.w_dim = w_dim
+        self.module = module
+        self.convs = torch.nn.ModuleList() 
+        self.module.apply(self.apply)
+        
+    def forward(self, input: torch.Tensor, styles: torch.Tensor) -> torch.Tensor:
+        for conv in self.convs:
+            conv.setStyle(styles.clone())
+        return self.module(input)
+
+    def apply(self, module: torch.nn.Module):
+        if isinstance(module, torch.nn.modules.conv._ConvNd):
+            delattr(module, "weight")
+            module.bias = None
+            
+            str_dim = module.__class__.__name__[-2:]
+            dim = 1
+            if str_dim == "2d":
+                dim = 2
+            elif str_dim == "3d":
+                dim = 3
+            self.convs.append(ModulatedConv._ModulatedConv(self.w_dim, module, dim=dim))
+        
+class UNetBlock(network.ModuleArgsDict):
+
+    def __init__(self, w_dim: int, channels: list[int], nb_conv_per_stage: int, blockConfig: blocks.BlockConfig, downSampleMode: blocks.DownSampleMode, upSampleMode: blocks.UpSampleMode, attention : bool, dim: int, i : int = 0) -> None:
+        super().__init__()
+        if i > 0:
+            self.add_module(downSampleMode.name, blocks.downSample(in_channels=channels[0], out_channels=channels[1], downSampleMode=downSampleMode, dim=dim))
+        self.add_module("DownConvBlock", blocks.ConvBlock(in_channels=channels[1 if downSampleMode == blocks.DownSampleMode.CONV_STRIDE and i > 0 else 0], out_channels=channels[1], nb_conv=nb_conv_per_stage, blockConfig=blockConfig, dim=dim))
+        if len(channels) > 2:
+            self.add_module("UNetBlock_{}".format(i+1), UNetBlock(w_dim, channels[1:], nb_conv_per_stage, blockConfig, downSampleMode, upSampleMode, attention, dim, i+1), in_branch=[0,1])
+            self.add_module("UpConvBlock", ModulatedConv(w_dim, blocks.ConvBlock((channels[1]+channels[2]) if upSampleMode != blocks.UpSampleMode.CONV_TRANSPOSE else channels[1]*2, out_channels=channels[1], nb_conv=nb_conv_per_stage, blockConfig=blockConfig, dim=dim)), in_branch=[0,1])
+        if i > 0:
+            if attention:
+                self.add_module("Attention", blocks.Attention(F_g=channels[1], F_l=channels[0], F_int=channels[0], dim=dim), in_branch=["Skip", 0], out_branch=["Skip"])
+            self.add_module(upSampleMode.name, ModulatedConv(w_dim, blocks.upSample(in_channels=channels[1], out_channels=channels[0], upSampleMode=upSampleMode, dim=dim)), in_branch=[0, 1])
+            self.add_module("SkipConnection", blocks.Concat(), in_branch=[0, "Skip"])
+
+class Generator(network.Network):
+
+    class GeneratorHead(network.ModuleArgsDict):
+
+        def __init__(self, in_channels: int, out_channels: int, dim: int) -> None:
+            super().__init__()
+            self.add_module("Conv", blocks.getTorchModule("Conv", dim)(in_channels = in_channels, out_channels = out_channels, kernel_size = 3, stride = 1, padding = 1))
+            self.add_module("Tanh", torch.nn.Tanh())
+
+    @config("Generator")
+    def __init__(self, 
+                    optimizer : network.OptimizerLoader = network.OptimizerLoader(),
+                    schedulers : network.LRSchedulersLoader = network.LRSchedulersLoader(),
+                    patch : ModelPatch = ModelPatch(),
+                    outputsCriterions: dict[str, network.TargetCriterionsLoader] = {"default" : network.TargetCriterionsLoader()},
+                    channels: list[int]=[1, 64, 128, 256, 512, 1024],
+                    nb_batch_per_step: int = 64,
+                    z_dim: int = 512,
+                    c_dim: int = 1,
+                    w_dim: int = 512,
+                    dim : int = 3) -> None:
+        super().__init__(optimizer=optimizer, in_channels=channels[0], schedulers=schedulers, patch=patch, outputsCriterions=outputsCriterions, dim=dim, nb_batch_per_step=nb_batch_per_step)
+
+        self.add_module("MappingNetwork", MappingNetwork(z_dim=z_dim, c_dim=c_dim, w_dim=w_dim, num_layers=8, embed_features=w_dim, layer_features=w_dim), in_branch=[1,2], out_branch=["Style"])
+        nb_conv_per_stage = 2
+        blockConfig = blocks.BlockConfig(kernel_size=3, stride=1, padding=1, bias=True, activation="ReLU", normMode="INSTANCE") 
+        self.add_module("UNetBlock_0", UNetBlock(w_dim, channels, nb_conv_per_stage, blockConfig, downSampleMode=blocks.DownSampleMode.MAXPOOL, upSampleMode=blocks.UpSampleMode.CONV_TRANSPOSE, attention=False, dim=dim), in_branch=[0, "Style"])
+        self.add_module("Head", Generator.GeneratorHead(in_channels=channels[1], out_channels=1, dim=dim))

konfai/models/generation/ddpm.py

@@ -0,0 +1,218 @@
+
+
+
+from typing import Union, Callable
+import torch
+import tqdm
+import numpy as np
+
+from KonfAI.konfai.network import network, blocks
+from KonfAI.konfai.utils.config import config
+from KonfAI.konfai.data.HDF5 import ModelPatch
+from KonfAI.konfai.utils.utils import gpuInfo
+from KonfAI.konfai.metric.measure import Criterion
+
+def cosine_beta_schedule(timesteps, s=0.008):
+    steps = timesteps + 1
+    x = torch.linspace(0, timesteps, steps)
+    alphas_cumprod = torch.cos(((x / timesteps) + s) / (1 + s) * torch.pi * 0.5) ** 2
+    alphas_cumprod = alphas_cumprod / alphas_cumprod[0]
+    betas = 1 - (alphas_cumprod[1:] / alphas_cumprod[:-1])
+    return torch.clip(betas, 0.0001, 0.9999)
+
+class DDPM(network.Network):
+    
+    class DDPM_TE(torch.nn.Module):
+
+        def __init__(self, in_channels: int, out_channels: int) -> None:
+            super().__init__()
+            self.linear_0 = torch.nn.Linear(in_channels, out_channels)
+            self.siLU = torch.nn.SiLU()
+            self.linear_1 = torch.nn.Linear(out_channels, out_channels)
+        
+        def forward(self, input: torch.Tensor, t: torch.Tensor) -> torch.Tensor:
+            return input + self.linear_1(self.siLU(self.linear_0(t))).reshape(input.shape[0], -1, *[1 for _ in range(len(input.shape)-2)])
+            
+    class DDPM_UNetBlock(network.ModuleArgsDict):
+        
+        def __init__(self, channels: list[int], nb_conv_per_stage: int, blockConfig: blocks.BlockConfig, downSampleMode: blocks.DownSampleMode, upSampleMode: blocks.UpSampleMode, attention : bool, time_embedding_dim: int, dim: int, i : int = 0) -> None:
+            super().__init__()
+            if i > 0:
+                self.add_module(downSampleMode.name, blocks.downSample(in_channels=channels[0], out_channels=channels[1], downSampleMode=downSampleMode, dim=dim))
+            self.add_module("Te_down", DDPM.DDPM_TE(time_embedding_dim, channels[1 if downSampleMode == blocks.DownSampleMode.CONV_STRIDE and i > 0 else 0]), in_branch=[0, 1])
+            self.add_module("DownConvBlock", blocks.ResBlock(in_channels=channels[1 if downSampleMode == blocks.DownSampleMode.CONV_STRIDE and i > 0 else 0], out_channels=channels[1], nb_conv=nb_conv_per_stage, blockConfig=blockConfig, dim=dim))
+            if len(channels) > 2:
+                self.add_module("UNetBlock_{}".format(i+1), DDPM.DDPM_UNetBlock(channels[1:], nb_conv_per_stage, blockConfig, downSampleMode, upSampleMode, attention, time_embedding_dim, dim, i+1), in_branch=[0,1])
+                self.add_module("Te_up", DDPM.DDPM_TE(time_embedding_dim, (channels[1]+channels[2]) if upSampleMode != blocks.UpSampleMode.CONV_TRANSPOSE else channels[1]*2), in_branch=[0, 1])
+                self.add_module("UpConvBlock", blocks.ResBlock(in_channels=(channels[1]+channels[2]) if upSampleMode != blocks.UpSampleMode.CONV_TRANSPOSE else channels[1]*2, out_channels=channels[1], nb_conv=nb_conv_per_stage, blockConfig=blockConfig, dim=dim))
+            if i > 0:
+                if attention:
+                    self.add_module("Attention", blocks.Attention(F_g=channels[1], F_l=channels[0], F_int=channels[0], dim=dim), in_branch=[2, 0], out_branch=[2])
+                self.add_module(upSampleMode.name, blocks.upSample(in_channels=channels[1], out_channels=channels[0], upSampleMode=upSampleMode, dim=dim))
+                self.add_module("SkipConnection", blocks.Concat(), in_branch=[0, 2])
+    
+    class DDPM_UNetHead(network.ModuleArgsDict):
+
+        def __init__(self, in_channels: int, out_channels: int, dim: int) -> None:
+            super().__init__()
+            self.add_module("Conv", blocks.getTorchModule("Conv", dim)(in_channels = in_channels, out_channels = out_channels, kernel_size = 3, stride = 1, padding = 1))
+
+    class DDPM_ForwardProcess(torch.nn.Module):
+
+        def __init__(self, noise_step: int=1000, beta_start: float = 1e-4, beta_end: float = 0.02) -> None:
+            super().__init__()
+            self.betas = torch.linspace(beta_start, beta_end, noise_step)
+            self.betas = DDPM.DDPM_ForwardProcess.enforce_zero_terminal_snr(self.betas)
+            self.alphas = 1 - self.betas
+            self.alpha_hat = torch.concat((torch.ones(1), torch.cumprod(self.alphas, dim=0)))
+
+        def enforce_zero_terminal_snr(betas: torch.Tensor):
+            alphas = 1 - betas
+            alphas_bar = alphas.cumprod(0)
+            alphas_bar_sqrt = alphas_bar.sqrt()
+            alphas_bar_sqrt_0 = alphas_bar_sqrt[0].clone()
+            alphas_bar_sqrt_T = alphas_bar_sqrt[-1].clone()
+            alphas_bar_sqrt -= alphas_bar_sqrt_T
+            alphas_bar_sqrt *= alphas_bar_sqrt_0 / (
+            alphas_bar_sqrt_0 - alphas_bar_sqrt_T)
+            alphas_bar = alphas_bar_sqrt ** 2
+            alphas = alphas_bar[1:] / alphas_bar[:-1]
+            alphas = torch.cat([alphas_bar[0:1], alphas])
+            betas = 1 - alphas
+            return betas
+
+        def forward(self, input: torch.Tensor, t: torch.Tensor, eta: torch.Tensor) -> torch.Tensor:
+            alpha_hat_t = self.alpha_hat[t.cpu()].to(input.device).reshape(input.shape[0], *[1 for _ in range(len(input.shape)-1)])
+            result = alpha_hat_t.sqrt() * input + (1 - alpha_hat_t).sqrt() * eta
+            return result
+    
+    class DDPM_SampleT(torch.nn.Module):
+
+        def __init__(self, noise_step: int) -> None:
+            super().__init__()
+            self.noise_step = noise_step
+    
+        def forward(self, input: torch.Tensor) -> torch.Tensor:
+            return torch.randint(0, self.noise_step, (input.shape[0],)).to(input.device)
+    
+    class DDPM_TimeEmbedding(torch.nn.Module):
+        
+        def sinusoidal_embedding(noise_step: int, time_embedding_dim: int):
+            noise_step += 1
+            embedding = torch.zeros(noise_step, time_embedding_dim)
+            wk = torch.tensor([1 / 10_000 ** (2 * j / time_embedding_dim) for j in range(time_embedding_dim)])
+            wk = wk.reshape((1, time_embedding_dim))
+            t = torch.arange(noise_step).reshape((noise_step, 1))
+            embedding[:,::2] = torch.sin(t * wk[:,::2])
+            embedding[:,1::2] = torch.cos(t * wk[:,::2])
+            return embedding
+
+        def __init__(self, noise_step: int=1000, time_embedding_dim: int=100) -> None:
+            super().__init__()
+            self.time_embed = torch.nn.Embedding(noise_step, time_embedding_dim)
+            self.time_embed.weight.data = DDPM.DDPM_TimeEmbedding.sinusoidal_embedding(noise_step, time_embedding_dim)
+            self.time_embed.requires_grad_(False)
+            self.noise_step = noise_step
+        
+        def forward(self, input: torch.Tensor, p: torch.Tensor) -> torch.Tensor:
+            return self.time_embed((p*self.noise_step).long().repeat((input.shape[0])))
+        
+    class DDPM_UNet(network.ModuleArgsDict):
+
+        def __init__(self, noise_step: int, channels: list[int], blockConfig: blocks.BlockConfig, nb_conv_per_stage: int, downSampleMode: str, upSampleMode: str, attention : bool, time_embedding_dim: int, dim : int) -> None:
+            super().__init__()
+            self.add_module("t", DDPM.DDPM_TimeEmbedding(noise_step, time_embedding_dim), in_branch=[1], out_branch=["te"])
+            self.add_module("UNetBlock_0", DDPM.DDPM_UNetBlock(channels, nb_conv_per_stage, blockConfig, downSampleMode=blocks.DownSampleMode._member_map_[downSampleMode], upSampleMode=blocks.UpSampleMode._member_map_[upSampleMode], attention=attention, time_embedding_dim=time_embedding_dim, dim=dim), in_branch=[0,"te"])
+            self.add_module("Head", DDPM.DDPM_UNetHead(in_channels=channels[1], out_channels=1, dim=dim))
+
+    class DDPM_Inference(torch.nn.Module):
+
+        def __init__(self, model: Callable[[torch.Tensor, torch.Tensor], torch.Tensor], train_noise_step: int, inference_noise_step: int, beta_start: float, beta_end: float) -> None:
+            super().__init__()
+            self.model = model
+            self.train_noise_step = train_noise_step
+            self.inference_noise_step = inference_noise_step
+            self.forwardProcess = DDPM.DDPM_ForwardProcess(train_noise_step, beta_start, beta_end)
+        
+        def forward(self, input: torch.Tensor) -> torch.Tensor:
+            x = torch.randn_like(input).to(input.device)
+            result = []
+            result.append(x.unsqueeze(1))
+            description = lambda : "Inference : "+gpuInfo(input.device)
+            offset = self.train_noise_step // self.inference_noise_step
+            t_list = np.round(np.arange(self.train_noise_step-1, 0, -offset)).astype(int).tolist()
+            with tqdm.tqdm(iterable = enumerate(t_list), desc = description(), total=len(t_list), leave=False, disable=True) as batch_iter:
+                for i, t in batch_iter:
+                    alpha_t_hat = self.forwardProcess.alpha_hat[t]
+                    alpha_t_hat_prev = self.forwardProcess.alpha_hat[t - offset] if t - offset >= 0 else self.forwardProcess.alpha_hat[0]
+                    eta_theta = self.model(torch.concat((x, input), dim=1), (torch.ones(input.shape[0], 1) * t).to(input.device).long())
+
+                    predicted = (x - (1 - alpha_t_hat).sqrt() * eta_theta)/alpha_t_hat.sqrt()*alpha_t_hat_prev.sqrt()
+
+                    variance = ((1 - alpha_t_hat_prev) / (1 - alpha_t_hat)) * (1 - alpha_t_hat / alpha_t_hat_prev)
+                    
+                    direction = (1-alpha_t_hat_prev-variance).sqrt()*eta_theta
+                    
+                    x = predicted+direction
+                    x += variance.sqrt()*torch.randn_like(input).to(input.device)
+                    
+                    result.append(x.unsqueeze(1))
+                    
+                    batch_iter.set_description(description())
+            result[-1] = torch.clip(result[-1], -1, 1)
+            return torch.concat(result, dim=1)
+
+    class DDPM_VLoss(torch.nn.Module):
+
+        def __init__(self, alpha_hat: torch.Tensor) -> None:
+            super().__init__()
+            self.alpha_hat = alpha_hat
+
+        def v(self, input: torch.Tensor, noise: torch.Tensor, t: torch.Tensor):
+            alpha_hat_t = self.alpha_hat[t.cpu()].to(input.device).reshape(input.shape[0], *[1 for _ in range(len(input.shape)-1)])
+            return alpha_hat_t.sqrt()*input-(1-alpha_hat_t).sqrt()*noise
+            
+        def forward(self, input: torch.Tensor, eta: torch.Tensor, eta_hat: torch.Tensor, t: int) -> torch.Tensor:
+            return torch.concat((self.v(input, eta, t), self.v(input, eta_hat, t)), dim=1)
+
+    @config("DDPM")
+    def __init__(   self,
+                    optimizer : network.OptimizerLoader = network.OptimizerLoader(),
+                    schedulers : network.LRSchedulersLoader = network.LRSchedulersLoader(),
+                    outputsCriterions: dict[str, network.TargetCriterionsLoader] = {"default" : network.TargetCriterionsLoader()},
+                    patch : Union[ModelPatch, None] = None,
+                    train_noise_step: int = 1000,
+                    inference_noise_step: int = 100,
+                    beta_start: float = 1e-4, 
+                    beta_end: float = 0.02,
+                    time_embedding_dim: int = 100,
+                    channels: list[int]=[1, 64, 128, 256, 512, 1024],
+                    blockConfig: blocks.BlockConfig = blocks.BlockConfig(),
+                    nb_conv_per_stage: int = 2,
+                    downSampleMode: str = "MAXPOOL",
+                    upSampleMode: str = "CONV_TRANSPOSE",
+                    attention : bool = False,
+                    dim : int = 3) -> None:
+        super().__init__(in_channels=1, optimizer=optimizer, schedulers=schedulers, outputsCriterions=outputsCriterions, patch=patch, dim=dim)
+        self.add_module("Noise", blocks.NormalNoise(), out_branch=["eta"], training=True)
+        self.add_module("Sample", DDPM.DDPM_SampleT(train_noise_step), out_branch=["t"], training=True)
+        self.add_module("Forward", DDPM.DDPM_ForwardProcess(train_noise_step, beta_start, beta_end), in_branch=[0, "t", "eta"], out_branch=["x_t"], training=True)
+        self.add_module("Concat", blocks.Concat(), in_branch=["x_t",1], out_branch=["xy_t"], training=True)
+        self.add_module("UNet", DDPM.DDPM_UNet(train_noise_step, channels, blockConfig, nb_conv_per_stage, downSampleMode, upSampleMode, attention, time_embedding_dim, dim), in_branch=["xy_t", "t"], out_branch=["eta_hat"], training=True)
+
+        self.add_module("Noise_optim", DDPM.DDPM_VLoss(self["Forward"].alpha_hat), in_branch=[0, "eta", "eta_hat", "t"], out_branch=["noise"], training=True)
+        
+        self.add_module("Inference", DDPM.DDPM_Inference(self.inference, train_noise_step, inference_noise_step, beta_start, beta_end), in_branch=[1], training=False)
+        self.add_module("LastImage", blocks.Select([slice(None, None), slice(-1, None)]), training=False)
+    
+    def inference(self, input: torch.Tensor, t: torch.Tensor):
+        return self["UNet"](input, t)
+
+class MSE(Criterion):
+
+    def __init__(self):
+        super().__init__()
+        self.loss = torch.nn.MSELoss()
+    
+    def forward(self, input: torch.Tensor) -> torch.Tensor:
+        return self.loss(input[:, 0, ...], input[:, 1, ...])