Update codeformer_arch.py

codeformer_arch.py

@@ -138,7 +138,7 @@ class CodeFormer(VQAutoEncoder):
     def __init__(self, dim_embd=512, n_head=8, n_layers=9, 
                  codebook_size=1024, latent_size=256,
                  connect_list=['32', '64', '128', '256'],
-                 fix_modules=['quantize', 'decoder'], vqgan_path=None):  # Changed 'generator' to 'decoder'
+                 fix_modules=['quantize', 'decoder'], vqgan_path=None):
         down_factor = [1, 2, 2, 4, 4, 8]
         super().__init__(512, 64, down_factor, 'nearest', len(down_factor) - 1, 16, codebook_size)
 
@@ -158,8 +158,8 @@ class CodeFormer(VQAutoEncoder):
         self.position_emb = nn.Parameter(torch.zeros(latent_size, self.dim_embd))
         self.feat_emb = nn.Linear(256, self.dim_embd)
 
-        self.ft_layers = nn.Sequential(*[TransformerSALayer(embed_dim=dim_embd, nhead=nhead, dim_mlp=self.dim_mlp, dropout=0.0) 
-                                        for _ in range(self.n_layers)])
+        self.ft_layers = nn.Sequential(*[TransformerSALayer(embed_dim=dim_embd, n_head=n_head, dim_mlp=self.dim_mlp, dropout=0.0) 
+                                        for _ in range(self.n_layers)])  # Changed nhead to n_head
 
         self.idx_pred_layer = nn.Sequential(
             nn.LayerNorm(dim_embd),
@@ -225,7 +225,7 @@ class CodeFormer(VQAutoEncoder):
 
         x = quant_feat
         fuse_list = [self.fuse_generator_block[f_size] for f_size in self.connect_list]
-        for i, block in enumerate(self.decoder):  # Changed 'generator' to 'decoder'
+        for i, block in enumerate(self.decoder):
             x = block(x)
             if i in fuse_list:
                 f_size = str(x.shape[-1])