Update codeformer_arch.py

codeformer_arch.py

@@ -75,9 +75,9 @@ def _get_activation_fn(activation):
     raise RuntimeError(f"activation should be relu/gelu, not {activation}.")
 
 class TransformerSALayer(nn.Module):
-    def __init__(self, embed_dim, n_head=8, dim_mlp=2048, dropout=0.0, activation="gelu"):  # Changed nhead to n_head
+    def __init__(self, embed_dim, n_head=8, dim_mlp=2048, dropout=0.0, activation="gelu"):
         super().__init__()
-        self.self_attn = nn.MultiheadAttention(embed_dim, n_head, dropout=dropout)  # Changed nhead to n_head
+        self.self_attn = nn.MultiheadAttention(embed_dim, n_head, dropout=dropout)
         self.linear1 = nn.Linear(embed_dim, dim_mlp)
         self.dropout = nn.Dropout(dropout)
         self.linear2 = nn.Linear(dim_mlp, embed_dim)
@@ -140,7 +140,7 @@ class CodeFormer(VQAutoEncoder):
                  connect_list=['32', '64', '128', '256'],
                  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)
+        super().__init__(512, 64, down_factor, 'nearest', len(down_factor) - 1, 256, codebook_size)  # Changed z_channels from 16 to 256
 
         if vqgan_path is not None:
             self.load_state_dict(torch.load(vqgan_path, map_location='cpu')['params_ema'])