Update flux/model.py

flux/model.py

@@ -79,9 +79,9 @@ class Flux(nn.Module):
 
         self.final_layer = LastLayer(self.hidden_size, 1, self.out_channels)
 
-        self.toonmage_ca = None
-        self.toonmage_double_interval = 2
-        self.toonmage_single_interval = 4
+        self.pulid_ca = None
+        self.pulid_double_interval = 2
+        self.pulid_single_interval = 4
 
     def forward(
         self,
@@ -115,8 +115,8 @@ class Flux(nn.Module):
         for i, block in enumerate(self.double_blocks):
             img, txt = block(img=img, txt=txt, vec=vec, pe=pe)
 
-            if i % self.toonmage_double_interval == 0 and id is not None:
-                img = img + id_weight * self.toonmage_ca[ca_idx](id, img)
+            if i % self.pulid_double_interval == 0 and id is not None:
+                img = img + id_weight * self.pulid_ca[ca_idx](id, img)
                 ca_idx += 1
 
         img = torch.cat((txt, img), 1)
@@ -124,8 +124,8 @@ class Flux(nn.Module):
             x = block(img, vec=vec, pe=pe)
             real_img, txt = x[:, txt.shape[1]:, ...], x[:, :txt.shape[1], ...]
 
-            if i % self.toonmage_single_interval == 0 and id is not None:
-                real_img = real_img + id_weight * self.toonmage_ca[ca_idx](id, real_img)
+            if i % self.pulid_single_interval == 0 and id is not None:
+                real_img = real_img + id_weight * self.pulid_ca[ca_idx](id, real_img)
                 ca_idx += 1
 
             img = torch.cat((txt, real_img), 1)
@@ -133,3 +133,4 @@ class Flux(nn.Module):
 
         img = self.final_layer(img, vec)  # (N, T, patch_size ** 2 * out_channels)
         return img
+