Fixed minor bug in hflow, updated README

Author: myscience

README.md

@@ -8,12 +8,19 @@ The authors release their official implementation which can be found [here](http
 
 ```python
 import torch
-from src.hflow import HierarchicalFlow
+from src.hflow import HierarchyFlow
 
 flow = HierarchyFlow(
     inp_channels=3,
-    out_channels=[30, 120],
-    pad_size=10,
+    flow_channel_mult=[10, 4, 4], # Set channel mult of the hierarchy convs
+    feat_channel_mult=[3, 3, 3],  # Set the channel mult factors of the style convs 
+    pad_size = 10, # Input pad size
+    pad_mode = 'reflect',
+    style_out_dim = 8, # Number of channel of final style features
+    style_conv_kw = [
+        {'kernel_size' : 7, 'stride' : 1, 'padding' : 3},
+        *[{'kernel_size' : 4, 'stride' : 2, 'padding' : 1}] * 2
+    ] # Parameter for style convolutional layers, should match length of feat_channel_mult
 )
 
 x = torch.randn(1, 3, 256, 256)

src/hflow.py

@@ -5,7 +5,7 @@
 
 from typing import List
 from torch import Tensor
-from itertools import pairwise
+from itertools import pairwise, repeat
 
 class HierarchyFlow(nn.Module):
     '''
@@ -20,12 +20,12 @@ def __init__(
         feat_channel_mult : List[int] = [3, 3],
         pad_size : int = 10,
         pad_mode : str = 'reflect',
-        style_dim : int = 8,
-        style_kw : dict | None = None,
+        style_out_dim : int = 8,
+        style_conv_kw : dict | None = None,
     ):
         super(HierarchyFlow, self).__init__()
 
-        style_kw = default(style_kw, {'padding' : 0})
+        style_conv_kw = default(style_conv_kw, repeat({'kernel_size' : 3}))
 
         self.inp_channels = inp_channels
         self.flow_channel_mult = flow_channel_mult
@@ -39,18 +39,18 @@ def __init__(
                 ReversiblePad2d(pad_size, pad_mode=pad_mode),
                 *[HierarchyBlock(
                     inp_chn, out_chn,
-                    mlp_inp_dim=style_dim)
+                    mlp_inp_dim=style_out_dim)
                 for inp_chn, out_chn in pairwise(flow_channels)]
             ]
         )
 
         self.style_block = nn.Sequential([
             nn.Sequential(
-                nn.Conv2d(inp_chn, out_chn, **style_kw),
+                nn.Conv2d(inp_chn, out_chn, **style),
                 nn.ReLU(),
-            ) for inp_chn, out_chn in zip(feat_channels)],
+            ) for (inp_chn, out_chn), style in zip(feat_channels, style_conv_kw)],
             nn.AdaptiveAvgPool2d(1), # global average pooling
-            nn.Conv2d(feat_channels[-1], style_dim, 1, 1, 0)
+            nn.Conv2d(feat_channels[-1], style_out_dim, 1, 1, 0)
         )
 
     def forward(

src/losses.py

@@ -50,9 +50,10 @@ class StyleLoss(nn.Module):
         the `content_weight` parameter.
 
         Args:
-            content_weight (float): The weight of the content loss.
-            style_weight (float): The weight of the style loss.
-            vgg_layers (list): The layers of the VGG-19 model to use. 
+        - enc_depth (List[int]): List of integers representing the
+          backbone layer depth to use as feature encoders.
+        - backbone (str): Torchvision model to use as feature extractor
+        - content_weight (float): The weight of the content loss.
     '''
 
     def __init__(