Update app.py

app.py

@@ -36,8 +36,11 @@ def setup_environment():
 # Load CodeFormer model
 def load_codeformer():
     setup_environment()
-    model = CodeFormer(dim_embd=512, codebook_size=1024, n_head=8, n_layers=9, connect_list=['32', '64', '128', '256'])  # Changed n_layer to n_layers
-    model.load_state_dict(torch.load("weights/CodeFormer/codeformer.pth", map_location='cpu'))
+    model = CodeFormer(dim_embd=512, codebook_size=1024, n_head=8, n_layers=9, connect_list=['32', '64', '128', '256'])
+    # Load the state dict, extracting the 'params_ema' key
+    checkpoint = torch.load("weights/CodeFormer/codeformer.pth", map_location='cpu')
+    state_dict = checkpoint['params_ema'] if 'params_ema' in checkpoint else checkpoint
+    model.load_state_dict(state_dict, strict=False)  # Use strict=False to ignore missing keys
     model.eval()
     model = model.to('cpu')  # Force CPU
     return model