Modified s_code generation

app.py

@@ -194,7 +194,7 @@ class ImageEditor(object):
                     target_latents.append(project_code_by_edit_name(np_source_latent, attribute_name, strength))
 
         elif edit_choices["edit_type"] == "StyleCLIP":
-            source_s_dict = generators[0].get_s_code(source_latent, input_is_latent=True)
+            source_s_dict = generators[0].get_s_code(source_latent, input_is_latent=True)[0]
             target_latents.append(project_code_with_styleclip(source_s_dict, 
                                                               edit_choices["src_text"], 
                                                               edit_choices["tar_text"], 

model/sg2_model.py

@@ -526,36 +526,36 @@ class Generator(nn.Module):
         if not input_is_latent:
             styles = [self.style(s) for s in styles]
 
-        s_codes = {# const block
-                   self.modulation_layers[0]:  self.modulation_layers[0](styles[0]),
-                   self.modulation_layers[1]:  self.modulation_layers[1](styles[1]),
+        s_codes = [{# const block
+                   self.modulation_layers[0]:  self.modulation_layers[0](style[:, 0]),
+                   self.modulation_layers[1]:  self.modulation_layers[1](style[:, 1]),
                    # conv layers
-                   self.modulation_layers[2]:  self.modulation_layers[2](styles[2]),
-                   self.modulation_layers[3]:  self.modulation_layers[3](styles[3]),
-                   self.modulation_layers[5]:  self.modulation_layers[5](styles[4]),
-                   self.modulation_layers[6]:  self.modulation_layers[6](styles[5]),
-                   self.modulation_layers[8]:  self.modulation_layers[8](styles[6]),
-                   self.modulation_layers[9]:  self.modulation_layers[9](styles[7]),
-                   self.modulation_layers[11]: self.modulation_layers[11](styles[8]),
-                   self.modulation_layers[12]: self.modulation_layers[12](styles[9]),
-                   self.modulation_layers[14]: self.modulation_layers[14](styles[10]),
-                   self.modulation_layers[15]: self.modulation_layers[15](styles[11]),
-                   self.modulation_layers[17]: self.modulation_layers[17](styles[12]),
-                   self.modulation_layers[18]: self.modulation_layers[18](styles[13]),
-                   self.modulation_layers[20]: self.modulation_layers[20](styles[14]),
-                   self.modulation_layers[21]: self.modulation_layers[21](styles[15]),
-                   self.modulation_layers[23]: self.modulation_layers[23](styles[16]),
-                   self.modulation_layers[24]: self.modulation_layers[24](styles[17]),
+                   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]),
                    # toRGB layers
-                   self.modulation_layers[4]: self.modulation_layers[4](styles[3]),
-                   self.modulation_layers[7]: self.modulation_layers[7](styles[5]),
-                   self.modulation_layers[10]: self.modulation_layers[10](styles[7]),
-                   self.modulation_layers[13]: self.modulation_layers[13](styles[9]),
-                   self.modulation_layers[16]: self.modulation_layers[16](styles[11]),
-                   self.modulation_layers[19]: self.modulation_layers[19](styles[13]),
-                   self.modulation_layers[22]: self.modulation_layers[22](styles[15]),
-                   self.modulation_layers[25]: self.modulation_layers[25](styles[17]),
-        }
+                   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]),
+        } for style in styles]
 
         return s_codes