jumbo merge

modules/codeformer/codeformer_arch.py

@@ -7,7 +7,7 @@ from torch import nn, Tensor
 import torch.nn.functional as F
 from typing import Optional, List
 
-from modules.codeformer.vqgan_arch import *
+from modules.codeformer.vqgan_arch import VQAutoEncoder, ResBlock
 from basicsr.utils import get_root_logger
 from basicsr.utils.registry import ARCH_REGISTRY
 
@@ -160,12 +160,12 @@ class Fuse_sft_block(nn.Module):
 
 @ARCH_REGISTRY.register()
 class CodeFormer(VQAutoEncoder):
-    def __init__(self, dim_embd=512, n_head=8, n_layers=9, 
+    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']):
+                connect_list=('32', '64', '128', '256'),
+                fix_modules=('quantize', 'generator')):
         super(CodeFormer, self).__init__(512, 64, [1, 2, 2, 4, 4, 8], 'nearest',2, [16], codebook_size)
-
+      
         if fix_modules is not None:
             for module in fix_modules:
                 for param in getattr(self, module).parameters():
@@ -180,14 +180,14 @@ class CodeFormer(VQAutoEncoder):
         self.feat_emb = nn.Linear(256, self.dim_embd)
 
         # transformer
-        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=n_head, dim_mlp=self.dim_mlp, dropout=0.0)
                                     for _ in range(self.n_layers)])
 
         # logits_predict head
         self.idx_pred_layer = nn.Sequential(
             nn.LayerNorm(dim_embd),
             nn.Linear(dim_embd, codebook_size, bias=False))
-        
+
         self.channels = {
             '16': 512,
             '32': 256,