Update toonmage/fluxpipeline.py

toonmage/fluxpipeline.py

@@ -27,7 +27,7 @@ class ToonMagePipeline(nn.Module):
         single_interval = 4
 
         # init encoder
-        self.toonmage_encoder = IDFormer().to(self.device, self.weight_dtype)
+        self.pulid_encoder = IDFormer().to(self.device, self.weight_dtype)
 
         num_ca = 19 // double_interval + 38 // single_interval
         if 19 % double_interval != 0:
@@ -174,7 +174,7 @@ class ToonMagePipeline(nn.Module):
 
         id_cond = torch.cat([id_ante_embedding, id_cond_vit], dim=-1)
 
-        id_embedding = self.toonmage_encoder(id_cond, id_vit_hidden)
+        id_embedding = self.pulid_encoder(id_cond, id_vit_hidden)
 
         if not cal_uncond:
             return id_embedding, None
@@ -183,6 +183,6 @@ class ToonMagePipeline(nn.Module):
         id_vit_hidden_uncond = []
         for layer_idx in range(0, len(id_vit_hidden)):
             id_vit_hidden_uncond.append(torch.zeros_like(id_vit_hidden[layer_idx]))
-        uncond_id_embedding = self.toonmage_encoder(id_uncond, id_vit_hidden_uncond)
+        uncond_id_embedding = self.pulid_encoder(id_uncond, id_vit_hidden_uncond)
 
         return id_embedding, uncond_id_embedding