konfai 1.1.8__py3-none-any.whl → 1.2.0__py3-none-any.whl

konfai/models/generation/gan.py

@@ -1,9 +1,10 @@
 from functools import partial
+
 import torch
 
-from konfai.network import network, blocks
-from konfai.utils.config import config
 from konfai.data.patching import ModelPatch
+from konfai.network import blocks, network
+
 
 class Discriminator(network.Network):
 
@@ -11,121 +12,244 @@ class Discriminator(network.Network):
 
         def __init__(self, channels: list[int], strides: list[int], dim: int) -> None:
             super().__init__()
-            blockConfig = partial(blocks.BlockConfig, kernel_size=4, padding=1, bias=False, activation=partial(torch.nn.LeakyReLU, negative_slope = 0.2, inplace=True), normMode=blocks.NormMode.SYNCBATCH)
+            block_config = partial(
+                blocks.BlockConfig,
+                kernel_size=4,
+                padding=1,
+                bias=False,
+                activation=partial(torch.nn.LeakyReLU, negative_slope=0.2, inplace=True),
+                norm_mode=blocks.NormMode.SYNCBATCH,
+            )
             for i, (in_channels, out_channels, stride) in enumerate(zip(channels, channels[1:], strides)):
-                self.add_module("Layer_{}".format(i), blocks.ConvBlock(in_channels, out_channels, 1, blockConfig(stride=stride), dim))
-    
+                self.add_module(
+                    f"Layer_{i}",
+                    blocks.ConvBlock(in_channels, out_channels, [block_config(stride=stride)], dim),
+                )
+
     class DiscriminatorHead(network.ModuleArgsDict):
 
         def __init__(self, channels: int, dim: int) -> None:
             super().__init__()
-            self.add_module("Conv", blocks.getTorchModule("Conv", dim)(in_channels=channels, out_channels=1, kernel_size=4, stride=1, padding=1))
-            self.add_module("AdaptiveAvgPool", blocks.getTorchModule("AdaptiveAvgPool", dim)(tuple([1]*dim)))
+            self.add_module(
+                "Conv",
+                blocks.get_torch_module("Conv", dim)(
+                    in_channels=channels,
+                    out_channels=1,
+                    kernel_size=4,
+                    stride=1,
+                    padding=1,
+                ),
+            )
+            self.add_module(
+                "AdaptiveAvgPool",
+                blocks.get_torch_module("AdaptiveAvgPool", dim)(tuple([1] * dim)),
+            )
             self.add_module("Flatten", torch.nn.Flatten(1))
-    
-    def __init__(self,
-                    optimizer : network.OptimizerLoader = network.OptimizerLoader(),
-                    schedulers: dict[str, network.LRSchedulersLoader] = {"default:ReduceLROnPlateau": network.LRSchedulersLoader(0)},
-                    outputsCriterions: dict[str, network.TargetCriterionsLoader] = {"default" : network.TargetCriterionsLoader()},
-                    nb_batch_per_step: int = 64,
-                    dim : int = 3) -> None:
-        super().__init__(in_channels = 1, optimizer = optimizer, schedulers = schedulers, outputsCriterions = outputsCriterions, dim=dim, nb_batch_per_step=nb_batch_per_step)
+
+    def __init__(
+        self,
+        optimizer: network.OptimizerLoader = network.OptimizerLoader(),
+        schedulers: dict[str, network.LRSchedulersLoader] = {
+            "default:ReduceLROnPlateau": network.LRSchedulersLoader(0)
+        },
+        outputs_criterions: dict[str, network.TargetCriterionsLoader] = {"default": network.TargetCriterionsLoader()},
+        nb_batch_per_step: int = 64,
+        dim: int = 3,
+    ) -> None:
+        super().__init__(
+            in_channels=1,
+            optimizer=optimizer,
+            schedulers=schedulers,
+            outputs_criterions=outputs_criterions,
+            dim=dim,
+            nb_batch_per_step=nb_batch_per_step,
+        )
         channels = [1, 16, 32, 64, 64]
-        strides = [2,2,2,1]
+        strides = [2, 2, 2, 1]
         self.add_module("Layers", Discriminator.DiscriminatorNLayers(channels, strides, dim))
         self.add_module("Head", Discriminator.DiscriminatorHead(channels[-1], dim))
 
+
 class Generator(network.Network):
 
     class GeneratorStem(network.ModuleArgsDict):
 
         def __init__(self, in_channels: int, out_channels: int, dim: int) -> None:
             super().__init__()
-            self.add_module("ConvBlock", blocks.ConvBlock(in_channels, out_channels, nb_conv=1, blockConfig=blocks.BlockConfig(bias=False, activation="ReLU", normMode="SYNCBATCH"), dim=dim))
+            self.add_module(
+                "ConvBlock",
+                blocks.ConvBlock(
+                    in_channels,
+                    out_channels,
+                    block_configs=[blocks.BlockConfig(bias=False, activation="ReLU", norm_mode="SYNCBATCH")],
+                    dim=dim,
+                ),
+            )
 
     class GeneratorHead(network.ModuleArgsDict):
 
         def __init__(self, in_channels: int, out_channels: int, dim: int) -> None:
             super().__init__()
-            self.add_module("ConvBlock", blocks.ConvBlock(in_channels, in_channels, nb_conv=1, blockConfig=blocks.BlockConfig(bias=False, activation="ReLU", normMode="SYNCBATCH"), dim=dim))
-            self.add_module("Conv", blocks.getTorchModule("Conv", dim)(in_channels, out_channels, kernel_size=1, bias=False))
+            self.add_module(
+                "ConvBlock",
+                blocks.ConvBlock(
+                    in_channels,
+                    in_channels,
+                    block_configs=[blocks.BlockConfig(bias=False, activation="ReLU", norm_mode="SYNCBATCH")],
+                    dim=dim,
+                ),
+            )
+            self.add_module(
+                "Conv",
+                blocks.get_torch_module("Conv", dim)(in_channels, out_channels, kernel_size=1, bias=False),
+            )
             self.add_module("Tanh", torch.nn.Tanh())
 
     class GeneratorDownSample(network.ModuleArgsDict):
 
         def __init__(self, in_channels: int, out_channels: int, dim: int) -> None:
             super().__init__()
-            self.add_module("ConvBlock", blocks.ConvBlock(in_channels, out_channels, nb_conv=1, blockConfig=blocks.BlockConfig(stride=2, bias=False, activation="ReLU", normMode="SYNCBATCH"), dim=dim))
-    
+            self.add_module(
+                "ConvBlock",
+                blocks.ConvBlock(
+                    in_channels,
+                    out_channels,
+                    block_configs=[blocks.BlockConfig(stride=2, bias=False, activation="ReLU", norm_mode="SYNCBATCH")],
+                    dim=dim,
+                ),
+            )
+
     class GeneratorUpSample(network.ModuleArgsDict):
 
         def __init__(self, in_channels: int, out_channels: int, dim: int) -> None:
             super().__init__()
-            self.add_module("ConvBlock", blocks.ConvBlock(in_channels, out_channels, nb_conv=1, blockConfig=blocks.BlockConfig(bias=False, activation="ReLU", normMode="SYNCBATCH"), dim=dim))
-            self.add_module("Upsample", torch.nn.Upsample(scale_factor=2, mode="bilinear" if dim < 3 else "trilinear"))
-    
+            self.add_module(
+                "ConvBlock",
+                blocks.ConvBlock(
+                    in_channels,
+                    out_channels,
+                    block_configs=[blocks.BlockConfig(bias=False, activation="ReLU", norm_mode="SYNCBATCH")],
+                    dim=dim,
+                ),
+            )
+            self.add_module(
+                "Upsample",
+                torch.nn.Upsample(scale_factor=2, mode="bilinear" if dim < 3 else "trilinear"),
+            )
+
     class GeneratorEncoder(network.ModuleArgsDict):
         def __init__(self, channels: list[int], dim: int) -> None:
             super().__init__()
             for i, (in_channels, out_channels) in enumerate(zip(channels, channels[1:])):
-                self.add_module("DownSample_{}".format(i), Generator.GeneratorDownSample(in_channels=in_channels, out_channels=out_channels, dim=dim))
-    
+                self.add_module(
+                    f"DownSample_{i}",
+                    Generator.GeneratorDownSample(in_channels=in_channels, out_channels=out_channels, dim=dim),
+                )
+
     class GeneratorResnetBlock(network.ModuleArgsDict):
 
-        def __init__(self, channels : int, dim : int):
+        def __init__(self, channels: int, dim: int):
             super().__init__()
-            self.add_module("Conv_0", blocks.getTorchModule("Conv", dim)(channels, channels, kernel_size=3, padding=1, bias=False))
+            self.add_module(
+                "Conv_0",
+                blocks.get_torch_module("Conv", dim)(channels, channels, kernel_size=3, padding=1, bias=False),
+            )
             self.add_module("Norm", torch.nn.LeakyReLU(0.2, inplace=True))
-            self.add_module("Conv_1", blocks.getTorchModule("Conv", dim)(channels, channels, kernel_size=3, padding=1, bias=False))
-            self.add_module("Residual", blocks.Add(), in_branch=[0,1])
+            self.add_module(
+                "Conv_1",
+                blocks.get_torch_module("Conv", dim)(channels, channels, kernel_size=3, padding=1, bias=False),
+            )
+            self.add_module("Residual", blocks.Add(), in_branch=[0, 1])
 
     class GeneratorNResnetBlock(network.ModuleArgsDict):
 
         def __init__(self, channels: int, nb_conv: int, dim: int) -> None:
             super().__init__()
             for i in range(nb_conv):
-                self.add_module("ResnetBlock_{}".format(i), Generator.GeneratorResnetBlock(channels=channels, dim=dim))
+                self.add_module(
+                    f"ResnetBlock_{i}",
+                    Generator.GeneratorResnetBlock(channels=channels, dim=dim),
+                )
 
     class GeneratorDecoder(network.ModuleArgsDict):
         def __init__(self, channels: list[int], dim: int) -> None:
             super().__init__()
             for i, (in_channels, out_channels) in enumerate(zip(reversed(channels), reversed(channels[:-1]))):
-                self.add_module("UpSample_{}".format(i), Generator.GeneratorUpSample(in_channels=in_channels, out_channels=out_channels, dim=dim))
-    
+                self.add_module(
+                    f"UpSample_{i}",
+                    Generator.GeneratorUpSample(in_channels=in_channels, out_channels=out_channels, dim=dim),
+                )
+
     class GeneratorAutoEncoder(network.ModuleArgsDict):
 
         def __init__(self, ngf: int, dim: int) -> None:
             super().__init__()
-            channels = [ngf, ngf*2]
+            channels = [ngf, ngf * 2]
             self.add_module("Encoder", Generator.GeneratorEncoder(channels, dim))
-            self.add_module("NResBlock", Generator.GeneratorNResnetBlock(channels=channels[-1], nb_conv=6, dim=dim))
+            self.add_module(
+                "NResBlock",
+                Generator.GeneratorNResnetBlock(channels=channels[-1], nb_conv=6, dim=dim),
+            )
             self.add_module("Decoder", Generator.GeneratorDecoder(channels, dim))
-        
-    def __init__(self, 
-                    optimizer : network.OptimizerLoader = network.OptimizerLoader(),
-                    schedulers: dict[str, network.LRSchedulersLoader] = {"default:ReduceLROnPlateau": network.LRSchedulersLoader(0)},
-                    patch : ModelPatch = ModelPatch(),
-                    outputsCriterions: dict[str, network.TargetCriterionsLoader] = {"default" : network.TargetCriterionsLoader()},
-                    nb_batch_per_step: int = 64,
-                    dim : int = 3) -> None:
-        super().__init__(optimizer=optimizer, in_channels=1, schedulers=schedulers, patch=patch, outputsCriterions=outputsCriterions, dim=dim, nb_batch_per_step=nb_batch_per_step)
-        ngf=32
+
+    def __init__(
+        self,
+        optimizer: network.OptimizerLoader = network.OptimizerLoader(),
+        schedulers: dict[str, network.LRSchedulersLoader] = {
+            "default:ReduceLROnPlateau": network.LRSchedulersLoader(0)
+        },
+        patch: ModelPatch = ModelPatch(),
+        outputs_criterions: dict[str, network.TargetCriterionsLoader] = {"default": network.TargetCriterionsLoader()},
+        nb_batch_per_step: int = 64,
+        dim: int = 3,
+    ) -> None:
+        super().__init__(
+            optimizer=optimizer,
+            in_channels=1,
+            schedulers=schedulers,
+            patch=patch,
+            outputs_criterions=outputs_criterions,
+            dim=dim,
+            nb_batch_per_step=nb_batch_per_step,
+        )
+        ngf = 32
         self.add_module("Stem", Generator.GeneratorStem(1, ngf, dim))
         self.add_module("AutoEncoder", Generator.GeneratorAutoEncoder(ngf, dim))
         self.add_module("Head", Generator.GeneratorHead(in_channels=ngf, out_channels=1, dim=dim))
 
-    def getName(self):
+    def get_name(self):
         return "Generator"
 
+
 class Gan(network.Network):
 
-    def __init__(self, generator : Generator = Generator(), discriminator : Discriminator = Discriminator()) -> None:
+    def __init__(
+        self,
+        generator: Generator = Generator(),
+        discriminator: Discriminator = Discriminator(),
+    ) -> None:
         super().__init__()
-        self.add_module("Discriminator_B", discriminator, in_branch=[1], out_branch=[-1], requires_grad=True)
+        self.add_module(
+            "Discriminator_B",
+            discriminator,
+            in_branch=[1],
+            out_branch=[-1],
+            requires_grad=True,
+        )
         self.add_module("Generator_A_to_B", generator, in_branch=[0], out_branch=["pB"])
-        
-        self.add_module("detach", blocks.Detach(), in_branch=["pB"], out_branch=["pB_detach"])
-        self.add_module("Discriminator_pB_detach", discriminator, in_branch=["pB_detach"], out_branch=[-1])
-          
-        self.add_module("Discriminator_pB", discriminator, in_branch=["pB"], out_branch=[-1], requires_grad=False)
 
+        self.add_module("detach", blocks.Detach(), in_branch=["pB"], out_branch=["pB_detach"])
+        self.add_module(
+            "Discriminator_pB_detach",
+            discriminator,
+            in_branch=["pB_detach"],
+            out_branch=[-1],
+        )
+
+        self.add_module(
+            "Discriminator_pB",
+            discriminator,
+            in_branch=["pB"],
+            out_branch=[-1],
+            requires_grad=False,
+        )

konfai/models/generation/vae.py

@@ -1,70 +1,170 @@
 import torch
 
-from konfai.network import network, blocks
-from konfai.utils.config import config
+from konfai.network import blocks, network
+
 
 class VAE(network.Network):
 
     class AutoEncoderBlock(network.ModuleArgsDict):
 
-        def __init__(self, channels: list[int], nb_conv_per_stage: int, blockConfig: blocks.BlockConfig, downSampleMode: blocks.DownSampleMode, upSampleMode: blocks.UpSampleMode, dim: int, block: type, i : int = 0) -> None:
+        def __init__(
+            self,
+            channels: list[int],
+            nb_conv_per_stage: int,
+            block_config: blocks.BlockConfig,
+            downsample_mode: blocks.DownsampleMode,
+            upsample_mode: blocks.UpsampleMode,
+            dim: int,
+            block: type,
+            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("DownBlock", block(in_channels=channels[1 if downSampleMode == blocks.DownSampleMode.CONV_STRIDE and i > 0 else 0], out_channels=channels[1], blockConfigs=[blockConfig]*nb_conv_per_stage, dim=dim))
+                self.add_module(
+                    downsample_mode.name,
+                    blocks.downsample(
+                        in_channels=channels[0],
+                        out_channels=channels[1],
+                        downsample_mode=downsample_mode,
+                        dim=dim,
+                    ),
+                )
+            self.add_module(
+                "DownBlock",
+                block(
+                    in_channels=channels[(1 if downsample_mode == blocks.DownsampleMode.CONV_STRIDE and i > 0 else 0)],
+                    out_channels=channels[1],
+                    block_configs=[block_config] * nb_conv_per_stage,
+                    dim=dim,
+                ),
+            )
             if len(channels) > 2:
-                self.add_module("AutoEncoder_{}".format(i+1), VAE.AutoEncoderBlock(channels[1:], nb_conv_per_stage, blockConfig, downSampleMode, upSampleMode, dim, block, i+1))
-                self.add_module("UpBlock", block(in_channels=channels[2] if upSampleMode != blocks.UpSampleMode.CONV_TRANSPOSE else channels[1], out_channels=channels[1], blockConfigs=[blockConfig]*nb_conv_per_stage, dim=dim))
+                self.add_module(
+                    f"AutoEncoder_{i + 1}",
+                    VAE.AutoEncoderBlock(
+                        channels[1:],
+                        nb_conv_per_stage,
+                        block_config,
+                        downsample_mode,
+                        upsample_mode,
+                        dim,
+                        block,
+                        i + 1,
+                    ),
+                )
+                self.add_module(
+                    "UpBlock",
+                    block(
+                        in_channels=(
+                            channels[2] if upsample_mode != blocks.UpsampleMode.CONV_TRANSPOSE else channels[1]
+                        ),
+                        out_channels=channels[1],
+                        block_configs=[block_config] * nb_conv_per_stage,
+                        dim=dim,
+                    ),
+                )
             if i > 0:
-               self.add_module(upSampleMode.name, blocks.upSample(in_channels=channels[1], out_channels=channels[0], upSampleMode=upSampleMode, dim=dim))
-    
-    class VAE_Head(network.ModuleArgsDict):
+                self.add_module(
+                    upsample_mode.name,
+                    blocks.upsample(
+                        in_channels=channels[1],
+                        out_channels=channels[0],
+                        upsample_mode=upsample_mode,
+                        dim=dim,
+                    ),
+                )
+
+    class VAEHead(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()) 
-
-    def __init__(self,
-                    optimizer: network.OptimizerLoader = network.OptimizerLoader(),
-                    schedulers: network.LRSchedulersLoader = network.LRSchedulersLoader(),
-                    outputsCriterions: dict[str, network.TargetCriterionsLoader] = {"default" : network.TargetCriterionsLoader()},
-                    dim : int = 3,
-                    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",
-                    blockType: str = "Conv") -> None:
-        
-        super().__init__(in_channels = channels[0], init_type="normal", optimizer = optimizer, schedulers = schedulers, outputsCriterions = outputsCriterions, dim=dim, nb_batch_per_step=1)
-        self.add_module("AutoEncoder_0", VAE.AutoEncoderBlock(channels, nb_conv_per_stage, blockConfig, downSampleMode=blocks.DownSampleMode._member_map_[downSampleMode], upSampleMode=blocks.UpSampleMode._member_map_[upSampleMode], dim=dim, block = blocks.ConvBlock if blockType == "Conv" else blocks.ResBlock))
-        self.add_module("Head", VAE.VAE_Head(channels[1], channels[0], dim))
+            self.add_module(
+                "Conv",
+                blocks.get_torch_module("Conv", dim)(
+                    in_channels=in_channels,
+                    out_channels=out_channels,
+                    kernel_size=3,
+                    stride=1,
+                    padding=1,
+                ),
+            )
+            self.add_module("Tanh", torch.nn.Tanh())
+
+    def __init__(
+        self,
+        optimizer: network.OptimizerLoader = network.OptimizerLoader(),
+        schedulers: dict[str, network.LRSchedulersLoader] = {
+            "default:ReduceLROnPlateau": network.LRSchedulersLoader(0)
+        },
+        outputs_criterions: dict[str, network.TargetCriterionsLoader] = {"default": network.TargetCriterionsLoader()},
+        dim: int = 3,
+        channels: list[int] = [1, 64, 128, 256, 512, 1024],
+        block_config: blocks.BlockConfig = blocks.BlockConfig(),
+        nb_conv_per_stage: int = 2,
+        downsample_mode: str = "MAXPOOL",
+        upsample_mode: str = "CONV_TRANSPOSE",
+        block_type: str = "Conv",
+    ) -> None:
+
+        super().__init__(
+            in_channels=channels[0],
+            init_type="normal",
+            optimizer=optimizer,
+            schedulers=schedulers,
+            outputs_criterions=outputs_criterions,
+            dim=dim,
+            nb_batch_per_step=1,
+        )
+        self.add_module(
+            "AutoEncoder_0",
+            VAE.AutoEncoderBlock(
+                channels,
+                nb_conv_per_stage,
+                block_config,
+                downsample_mode=blocks.DownsampleMode[downsample_mode],
+                upsample_mode=blocks.UpsampleMode[upsample_mode],
+                dim=dim,
+                block=blocks.ConvBlock if block_type == "Conv" else blocks.ResBlock,
+            ),
+        )
+        self.add_module("Head", VAE.VAEHead(channels[1], channels[0], dim))
 
 
 class LinearVAE(network.Network):
 
-    class LinearVAE_DenseLayer(network.ModuleArgsDict):
+    class LinearVAEDenseLayer(network.ModuleArgsDict):
 
         def __init__(self, in_features: int, out_features: int) -> None:
             super().__init__()
             self.add_module("Linear", torch.nn.Linear(in_features, out_features))
-            #self.add_module("Norm", torch.nn.BatchNorm1d(out_features))
+            # self.add_module("Norm", torch.nn.BatchNorm1d(out_features))
             self.add_module("Activation", torch.nn.LeakyReLU())
 
-    class LinearVAE_Head(network.ModuleArgsDict):
+    class LinearVAEHead(network.ModuleArgsDict):
 
         def __init__(self, in_features: int, out_features: int) -> None:
             super().__init__()
             self.add_module("Linear", torch.nn.Linear(in_features, out_features))
             self.add_module("Tanh", torch.nn.Tanh())
 
-    def __init__(self,
-                    optimizer: network.OptimizerLoader = network.OptimizerLoader(),
-                    schedulers: network.LRSchedulersLoader = network.LRSchedulersLoader(),
-                    outputsCriterions: dict[str, network.TargetCriterionsLoader] = {"default" : network.TargetCriterionsLoader()},) -> None:
-        super().__init__(in_channels = 1, init_type="normal", optimizer = optimizer, schedulers = schedulers, outputsCriterions = outputsCriterions, dim=1, nb_batch_per_step=1)
-        self.add_module("DenseLayer_0", LinearVAE.LinearVAE_DenseLayer(23343, 5))
-        #self.add_module("Head", LinearVAE.DenseLayer(100, 28590))
-        self.add_module("Head", LinearVAE.LinearVAE_Head(5, 23343))
-        #self.add_module("DenseLayer_5", LinearVAE.DenseLayer(5000, 28590))
+    def __init__(
+        self,
+        optimizer: network.OptimizerLoader = network.OptimizerLoader(),
+        schedulers: dict[str, network.LRSchedulersLoader] = {
+            "default:ReduceLROnPlateau": network.LRSchedulersLoader(0)
+        },
+        outputs_criterions: dict[str, network.TargetCriterionsLoader] = {"default": network.TargetCriterionsLoader()},
+    ) -> None:
+        super().__init__(
+            in_channels=1,
+            init_type="normal",
+            optimizer=optimizer,
+            schedulers=schedulers,
+            outputs_criterions=outputs_criterions,
+            dim=1,
+            nb_batch_per_step=1,
+        )
+        self.add_module("DenseLayer_0", LinearVAE.LinearVAEDenseLayer(23343, 5))
+        # self.add_module("Head", LinearVAE.DenseLayer(100, 28590))
+        self.add_module("Head", LinearVAE.LinearVAEHead(5, 23343))
+        # self.add_module("DenseLayer_5", LinearVAE.DenseLayer(5000, 28590))