StyleSpace list now uses correct order.

model/sg2_model.py

@@ -527,38 +527,39 @@ class Generator(nn.Module):
             styles = [self.style(s) for s in styles]
 
         s_codes = [{# const block
-                   self.modulation_layers[0]:  self.modulation_layers[0](style[:, 0]),
-                   self.modulation_layers[1]:  self.modulation_layers[1](style[:, 1]),
+                   self.modulation_layers[0]:  self.modulation_layers[0](style[:, 0]),   #s0
+                   self.modulation_layers[1]:  self.modulation_layers[1](style[:, 1]),   #s1
                    # conv layers
-                   self.modulation_layers[2]:  self.modulation_layers[2](style[:, 2]),
-                   self.modulation_layers[3]:  self.modulation_layers[3](style[:, 3]),
-                   self.modulation_layers[5]:  self.modulation_layers[5](style[:, 4]),
-                   self.modulation_layers[6]:  self.modulation_layers[6](style[:, 5]),
-                   self.modulation_layers[8]:  self.modulation_layers[8](style[:, 6]),
-                   self.modulation_layers[9]:  self.modulation_layers[9](style[:, 7]),
-                   self.modulation_layers[11]: self.modulation_layers[11](style[:, 8]),
-                   self.modulation_layers[12]: self.modulation_layers[12](style[:, 9]),
-                   self.modulation_layers[14]: self.modulation_layers[14](style[:, 10]),
-                   self.modulation_layers[15]: self.modulation_layers[15](style[:, 11]),
-                   self.modulation_layers[17]: self.modulation_layers[17](style[:, 12]),
-                   self.modulation_layers[18]: self.modulation_layers[18](style[:, 13]),
-                   self.modulation_layers[20]: self.modulation_layers[20](style[:, 14]),
-                   self.modulation_layers[21]: self.modulation_layers[21](style[:, 15]),
-                   self.modulation_layers[23]: self.modulation_layers[23](style[:, 16]),
-                   self.modulation_layers[24]: self.modulation_layers[24](style[:, 17]),
+                   self.modulation_layers[2]:  self.modulation_layers[2](style[:, 1]),   #s2
+                   self.modulation_layers[3]:  self.modulation_layers[3](style[:, 2]),   #s3
+                   self.modulation_layers[4]:  self.modulation_layers[4](style[:, 3]),   #s5
+                   self.modulation_layers[5]:  self.modulation_layers[5](style[:, 4]),   #s6
+                   self.modulation_layers[6]:  self.modulation_layers[6](style[:, 5]),   #s8
+                   self.modulation_layers[7]:  self.modulation_layers[7](style[:, 6]),   #s9
+                   self.modulation_layers[8]:  self.modulation_layers[8](style[:, 7]),   #s11
+                   self.modulation_layers[9]:  self.modulation_layers[9](style[:, 8]),   #s12
+                   self.modulation_layers[10]: self.modulation_layers[10](style[:, 9]), #s14
+                   self.modulation_layers[11]: self.modulation_layers[11](style[:, 10]), #s15
+                   self.modulation_layers[12]: self.modulation_layers[12](style[:, 11]), #s17
+                   self.modulation_layers[13]: self.modulation_layers[13](style[:, 12]), #s18
+                   self.modulation_layers[14]: self.modulation_layers[14](style[:, 13]), #s20
+                   self.modulation_layers[15]: self.modulation_layers[15](style[:, 14]), #s21
+                   self.modulation_layers[16]: self.modulation_layers[16](style[:, 15]), #s23
+                   self.modulation_layers[17]: self.modulation_layers[17](style[:, 16]), #s24
                    # toRGB layers
-                   self.modulation_layers[4]: self.modulation_layers[4](style[:, 3]),
-                   self.modulation_layers[7]: self.modulation_layers[7](style[:, 5]),
-                   self.modulation_layers[10]: self.modulation_layers[10](style[:, 7]),
-                   self.modulation_layers[13]: self.modulation_layers[13](style[:, 9]),
-                   self.modulation_layers[16]: self.modulation_layers[16](style[:, 11]),
-                   self.modulation_layers[19]: self.modulation_layers[19](style[:, 13]),
-                   self.modulation_layers[22]: self.modulation_layers[22](style[:, 15]),
-                   self.modulation_layers[25]: self.modulation_layers[25](style[:, 17]),
+                   self.modulation_layers[18]: self.modulation_layers[18](style[:, 3]),  #s4
+                   self.modulation_layers[19]: self.modulation_layers[19](style[:, 5]),  #s7
+                   self.modulation_layers[20]: self.modulation_layers[20](style[:, 7]),  #s10
+                   self.modulation_layers[21]: self.modulation_layers[21](style[:, 9]),  #s13
+                   self.modulation_layers[22]: self.modulation_layers[22](style[:, 11]), #s16
+                   self.modulation_layers[23]: self.modulation_layers[23](style[:, 13]), #s19
+                   self.modulation_layers[24]: self.modulation_layers[24](style[:, 15]), #s22
+                   self.modulation_layers[25]: self.modulation_layers[25](style[:, 17]), #s25
         } for style in styles]
-
+            
         return s_codes    
 
+
     def forward(
         self,
         styles,

styleclip/styleclip_global.py

@@ -114,7 +114,7 @@ def expand_to_full_dim(partial_tensor):
     start_idx = 0
     for conv_start, conv_end in CONV_CODE_INDICES:
         length = conv_end - conv_start
-        full_dim_tensor[:, conv_start:conv_end] = partial_tensor[:, start_idx:start_idx + length]
+        full_dim_tensor[:, conv_start:conv_end] = partial_tensor[start_idx:start_idx + length]
         start_idx += length
 
     return full_dim_tensor