Update codeformer_arch.py

codeformer_arch.py

@@ -138,9 +138,8 @@ 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', 'generator'], vqgan_path=None):
-        # Adjust down_factor to ensure it works with channel scaling
-        down_factor = [1, 2, 2, 4, 4, 8]  # Ensure this matches the number of steps
+                 fix_modules=['quantize', 'decoder'], vqgan_path=None):  # Changed 'generator' to 'decoder'
+        down_factor = [1, 2, 2, 4, 4, 8]
         super().__init__(512, 64, down_factor, 'nearest', len(down_factor) - 1, 16, codebook_size)
 
         if vqgan_path is not None:
@@ -159,7 +158,7 @@ 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=n_head, dim_mlp=self.dim_mlp, dropout=0.0) 
+        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.idx_pred_layer = nn.Sequential(
@@ -226,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):
+        for i, block in enumerate(self.decoder):  # Changed 'generator' to 'decoder'
             x = block(x)
             if i in fuse_list:
                 f_size = str(x.shape[-1])