feat: add CPU-based toy training configs and fix net_g.module AttributeError for CPU training

models/CodeFormer/basicsr/models/codeformer_idx_model.py

@@ -96,7 +96,7 @@ class CodeFormerIdxModel(SRModel):
         
         if self.hq_feat_loss:
             # quant_feats
-            quant_feat_gt = self.net_g.module.quantize.get_codebook_feat(self.idx_gt, shape=[self.b,16,16,256])
+            quant_feat_gt = self.get_bare_model(self.net_g).quantize.get_codebook_feat(self.idx_gt, shape=[self.b,16,16,256])
 
         logits, lq_feat = self.net_g(self.input, w=0, code_only=True)
 

models/CodeFormer/basicsr/models/codeformer_joint_model.py

@@ -172,7 +172,7 @@ class CodeFormerJointModel(SRModel):
 
         if self.hq_feat_loss:
             # quant_feats
-            quant_feat_gt = self.net_g.module.quantize.get_codebook_feat(self.idx_gt, shape=[self.b,16,16,256])
+            quant_feat_gt = self.get_bare_model(self.net_g).quantize.get_codebook_feat(self.idx_gt, shape=[self.b,16,16,256])
 
         l_g_total = 0
         loss_dict = OrderedDict()
@@ -210,11 +210,11 @@ class CodeFormerJointModel(SRModel):
                     l_g_gan = self.cri_gan(fake_g_pred, True, is_disc=False)
                     recon_loss = l_g_pix + l_g_percep
                     if not self.fix_generator:
-                        last_layer = self.net_g.module.generator.blocks[-1].weight
+                        last_layer = self.get_bare_model(self.net_g).generator.blocks[-1].weight
                         d_weight = self.calculate_adaptive_weight(recon_loss, l_g_gan, last_layer, disc_weight_max=1.0)
                     else:
                         largest_fuse_size = self.opt['network_g']['connect_list'][-1]
-                        last_layer = self.net_g.module.fuse_convs_dict[largest_fuse_size].shift[-1].weight
+                        last_layer = self.get_bare_model(self.net_g).fuse_convs_dict[largest_fuse_size].shift[-1].weight
                         d_weight = self.calculate_adaptive_weight(recon_loss, l_g_gan, last_layer, disc_weight_max=1.0)
                     
                     d_weight *= self.scale_adaptive_gan_weight # 0.8

models/CodeFormer/basicsr/models/codeformer_model.py

@@ -155,7 +155,7 @@ class CodeFormerModel(SRModel):
 
         if self.hq_feat_loss:
             # quant_feats
-            quant_feat_gt = self.net_g.module.quantize.get_codebook_feat(self.idx_gt, shape=[self.b,16,16,256])
+            quant_feat_gt = self.get_bare_model(self.net_g).quantize.get_codebook_feat(self.idx_gt, shape=[self.b,16,16,256])
 
         l_g_total = 0
         loss_dict = OrderedDict()
@@ -192,11 +192,11 @@ class CodeFormerModel(SRModel):
                     l_g_gan = self.cri_gan(fake_g_pred, True, is_disc=False)
                     recon_loss = l_g_pix + l_g_percep
                     if not self.fix_generator:
-                        last_layer = self.net_g.module.generator.blocks[-1].weight
+                        last_layer = self.get_bare_model(self.net_g).generator.blocks[-1].weight
                         d_weight = self.calculate_adaptive_weight(recon_loss, l_g_gan, last_layer, disc_weight_max=1.0)
                     else:
                         largest_fuse_size = self.opt['network_g']['connect_list'][-1]
-                        last_layer = self.net_g.module.fuse_convs_dict[largest_fuse_size].shift[-1].weight
+                        last_layer = self.get_bare_model(self.net_g).fuse_convs_dict[largest_fuse_size].shift[-1].weight
                         d_weight = self.calculate_adaptive_weight(recon_loss, l_g_gan, last_layer, disc_weight_max=1.0)
                     
                     d_weight *= self.scale_adaptive_gan_weight # 0.8

models/CodeFormer/basicsr/models/vqgan_model.py

@@ -118,9 +118,9 @@ class VQGANModel(SRModel):
         logger = get_root_logger()
         loss_dict = OrderedDict()
         if self.opt['network_g']['quantizer'] == 'gumbel':
-            self.net_g.module.quantize.temperature = max(1/16, ((-1/160000) * current_iter) + 1)
+            self.get_bare_model(self.net_g).quantize.temperature = max(1/16, ((-1/160000) * current_iter) + 1)
             if current_iter%1000 == 0:
-                logger.info(f'temperature: {self.net_g.module.quantize.temperature}')
+                logger.info(f'temperature: {self.get_bare_model(self.net_g).quantize.temperature}')
 
         # optimize net_g
         for p in self.net_d.parameters():
@@ -150,7 +150,7 @@ class VQGANModel(SRModel):
                 fake_g_pred = self.net_d(self.output)
                 l_g_gan = self.cri_gan(fake_g_pred, True, is_disc=False)
                 recon_loss = l_g_total
-                last_layer = self.net_g.module.generator.blocks[-1].weight
+                last_layer = self.get_bare_model(self.net_g).generator.blocks[-1].weight
                 d_weight = self.calculate_adaptive_weight(recon_loss, l_g_gan, last_layer, disc_weight_max=1.0)
                 d_weight *= self.adopt_weight(1, current_iter, self.net_d_start_iter)
                 d_weight *= self.disc_weight # tamming setting 0.8

models/CodeFormer/options/CodeFormer_stage2_toy.yml

@@ -0,0 +1,128 @@
+name: CodeFormer_stage2
+model_type: CodeFormerIdxModel
+num_gpu: 0
+manual_seed: 0
+datasets:
+  train:
+    name: FFHQ
+    type: FFHQBlindDataset
+    dataroot_gt: datasets/ffhq/ffhq_512
+    filename_tmpl: '{}'
+    io_backend:
+      type: disk
+    in_size: 512
+    gt_size: 512
+    mean:
+    - 0.5
+    - 0.5
+    - 0.5
+    std:
+    - 0.5
+    - 0.5
+    - 0.5
+    use_hflip: true
+    use_corrupt: true
+    blur_kernel_size: 41
+    use_motion_kernel: false
+    motion_kernel_prob: 0.001
+    kernel_list:
+    - iso
+    - aniso
+    kernel_prob:
+    - 0.5
+    - 0.5
+    blur_sigma:
+    - 1
+    - 15
+    downsample_range:
+    - 4
+    - 30
+    noise_range:
+    - 0
+    - 20
+    jpeg_range:
+    - 30
+    - 80
+    latent_gt_path: null
+    num_worker_per_gpu: 0
+    batch_size_per_gpu: 1
+    dataset_enlarge_ratio: 1
+    prefetch_mode: cpu
+network_g:
+  type: CodeFormer
+  dim_embd: 512
+  n_head: 8
+  n_layers: 9
+  codebook_size: 1024
+  connect_list:
+  - '32'
+  - '64'
+  - '128'
+  - '256'
+  fix_modules:
+  - quantize
+  - generator
+  vqgan_path: ./experiments/pretrained_models/vqgan/vqgan_code1024.pth
+network_vqgan:
+  type: VQAutoEncoder
+  img_size: 512
+  nf: 64
+  ch_mult:
+  - 1
+  - 2
+  - 2
+  - 4
+  - 4
+  - 8
+  quantizer: nearest
+  codebook_size: 1024
+path:
+  pretrain_network_g: null
+  param_key_g: params_ema
+  strict_load_g: false
+  pretrain_network_d: null
+  strict_load_d: true
+  resume_state: null
+train:
+  use_hq_feat_loss: true
+  feat_loss_weight: 1.0
+  cross_entropy_loss: true
+  entropy_loss_weight: 0.5
+  fidelity_weight: 0
+  optim_g:
+    type: Adam
+    lr: 0.0001
+    weight_decay: 0
+    betas:
+    - 0.9
+    - 0.99
+  scheduler:
+    type: MultiStepLR
+    milestones:
+    - 400000
+    - 450000
+    gamma: 0.5
+  total_iter: 5
+  warmup_iter: -1
+  ema_decay: 0.995
+  use_adaptive_weight: true
+  net_g_start_iter: 0
+  net_d_iters: 1
+  net_d_start_iter: 0
+  manual_seed: 0
+val:
+  val_freq: 50000000000.0
+  save_img: true
+  metrics:
+    psnr:
+      type: calculate_psnr
+      crop_border: 4
+      test_y_channel: false
+logger:
+  print_freq: 1
+  save_checkpoint_freq: 5
+  use_tb_logger: false
+  wandb: null
+dist_params: null
+find_unused_parameters: true
+dist: false

models/CodeFormer/options/CodeFormer_stage3_toy.yml

@@ -0,0 +1,164 @@
+name: CodeFormer_stage3
+model_type: CodeFormerJointModel
+num_gpu: 0
+manual_seed: 0
+datasets:
+  train:
+    name: FFHQ
+    type: FFHQBlindJointDataset
+    dataroot_gt: datasets/ffhq/ffhq_512
+    filename_tmpl: '{}'
+    io_backend:
+      type: disk
+    in_size: 512
+    gt_size: 512
+    mean:
+    - 0.5
+    - 0.5
+    - 0.5
+    std:
+    - 0.5
+    - 0.5
+    - 0.5
+    use_hflip: true
+    use_corrupt: true
+    blur_kernel_size: 41
+    use_motion_kernel: false
+    motion_kernel_prob: 0.001
+    kernel_list:
+    - iso
+    - aniso
+    kernel_prob:
+    - 0.5
+    - 0.5
+    blur_sigma:
+    - 0.1
+    - 10
+    downsample_range:
+    - 1
+    - 12
+    noise_range:
+    - 0
+    - 15
+    jpeg_range:
+    - 60
+    - 100
+    blur_sigma_large:
+    - 1
+    - 15
+    downsample_range_large:
+    - 4
+    - 30
+    noise_range_large:
+    - 0
+    - 20
+    jpeg_range_large:
+    - 30
+    - 80
+    latent_gt_path: null
+    num_worker_per_gpu: 0
+    batch_size_per_gpu: 1
+    dataset_enlarge_ratio: 1
+    prefetch_mode: cpu
+network_g:
+  type: CodeFormer
+  dim_embd: 512
+  n_head: 8
+  n_layers: 9
+  codebook_size: 1024
+  connect_list:
+  - '32'
+  - '64'
+  - '128'
+  - '256'
+  fix_modules:
+  - quantize
+  - generator
+network_vqgan:
+  type: VQAutoEncoder
+  img_size: 512
+  nf: 64
+  ch_mult:
+  - 1
+  - 2
+  - 2
+  - 4
+  - 4
+  - 8
+  quantizer: nearest
+  codebook_size: 1024
+network_d:
+  type: VQGANDiscriminator
+  nc: 3
+  ndf: 64
+  n_layers: 4
+path:
+  pretrain_network_g: ./experiments/pretrained_models/CodeFormer_stage2/net_g_latest.pth
+  param_key_g: params_ema
+  strict_load_g: false
+  pretrain_network_d: null
+  resume_state: null
+train:
+  use_hq_feat_loss: true
+  feat_loss_weight: 1.0
+  cross_entropy_loss: true
+  entropy_loss_weight: 0.5
+  scale_adaptive_gan_weight: 0.1
+  optim_g:
+    type: Adam
+    lr: 5.0e-05
+    weight_decay: 0
+    betas:
+    - 0.9
+    - 0.99
+  optim_d:
+    type: Adam
+    lr: 5.0e-05
+    weight_decay: 0
+    betas:
+    - 0.9
+    - 0.99
+  scheduler:
+    type: CosineAnnealingRestartLR
+    periods:
+    - 150000
+    restart_weights:
+    - 1
+    eta_min: 2.0e-05
+  total_iter: 5
+  warmup_iter: -1
+  ema_decay: 0.997
+  pixel_opt:
+    type: L1Loss
+    loss_weight: 1.0
+    reduction: mean
+  perceptual_opt:
+    type: LPIPSLoss
+    loss_weight: 1.0
+    use_input_norm: true
+    range_norm: true
+  gan_opt:
+    type: GANLoss
+    gan_type: hinge
+    loss_weight: 1.0
+  use_adaptive_weight: true
+  net_g_start_iter: 0
+  net_d_iters: 1
+  net_d_start_iter: 5001
+  manual_seed: 0
+val:
+  val_freq: 50000000000.0
+  save_img: true
+  metrics:
+    psnr:
+      type: calculate_psnr
+      crop_border: 4
+      test_y_channel: false
+logger:
+  print_freq: 1
+  save_checkpoint_freq: 5
+  use_tb_logger: false
+  wandb: null
+dist_params: null
+find_unused_parameters: true
+dist: false

models/CodeFormer/options/VQGAN_toy.yml

@@ -0,0 +1,111 @@
+name: VQGAN-512-ds32-nearest-stage1
+model_type: VQGANModel
+num_gpu: 0
+manual_seed: 0
+datasets:
+  train:
+    name: FFHQ
+    type: FFHQBlindDataset
+    dataroot_gt: datasets/ffhq/ffhq_512
+    filename_tmpl: '{}'
+    io_backend:
+      type: disk
+    in_size: 512
+    gt_size: 512
+    mean:
+    - 0.5
+    - 0.5
+    - 0.5
+    std:
+    - 0.5
+    - 0.5
+    - 0.5
+    use_hflip: true
+    use_corrupt: false
+    num_worker_per_gpu: 0
+    batch_size_per_gpu: 1
+    dataset_enlarge_ratio: 1
+    prefetch_mode: cpu
+    num_prefetch_queue: 4
+network_g:
+  type: VQAutoEncoder
+  img_size: 512
+  nf: 64
+  ch_mult:
+  - 1
+  - 2
+  - 2
+  - 4
+  - 4
+  - 8
+  quantizer: nearest
+  codebook_size: 1024
+network_d:
+  type: VQGANDiscriminator
+  nc: 3
+  ndf: 64
+path:
+  pretrain_network_g: null
+  param_key_g: params_ema
+  strict_load_g: true
+  pretrain_network_d: null
+  strict_load_d: true
+  resume_state: null
+train:
+  optim_g:
+    type: Adam
+    lr: 7.0e-05
+    weight_decay: 0
+    betas:
+    - 0.9
+    - 0.99
+  optim_d:
+    type: Adam
+    lr: 7.0e-05
+    weight_decay: 0
+    betas:
+    - 0.9
+    - 0.99
+  scheduler:
+    type: CosineAnnealingRestartLR
+    periods:
+    - 1600000
+    restart_weights:
+    - 1
+    eta_min: 6.0e-05
+  total_iter: 5
+  warmup_iter: -1
+  ema_decay: 0.995
+  pixel_opt:
+    type: L1Loss
+    loss_weight: 1.0
+    reduction: mean
+  perceptual_opt:
+    type: LPIPSLoss
+    loss_weight: 1.0
+    use_input_norm: true
+    range_norm: true
+  gan_opt:
+    type: GANLoss
+    gan_type: hinge
+    loss_weight: 1.0
+  net_g_start_iter: 0
+  net_d_iters: 1
+  net_d_start_iter: 30001
+  manual_seed: 0
+val:
+  val_freq: 50000000000.0
+  save_img: true
+  metrics:
+    psnr:
+      type: calculate_psnr
+      crop_border: 4
+      test_y_channel: false
+logger:
+  print_freq: 1
+  save_checkpoint_freq: 5
+  use_tb_logger: false
+  wandb: null
+dist_params: null
+find_unused_parameters: true
+dist: false

prepare_toy_training.py

@@ -0,0 +1,100 @@
+import os
+import requests
+import yaml
+
+def main():
+    project_dir = os.path.dirname(os.path.abspath(__file__))
+    codeformer_dir = os.path.join(project_dir, "models", "CodeFormer")
+    
+    # 1. Create target dataset directory
+    dataset_dir = os.path.join(codeformer_dir, "datasets", "ffhq", "ffhq_512")
+    os.makedirs(dataset_dir, exist_ok=True)
+    print(f"Created dataset directory: {dataset_dir}")
+    
+    # 2. Download toy dataset (30 face images from Unsplash)
+    unsplash_ids = [
+        "1544005313-94ddf0286df2", "1506794778202-cad84cf45f1d", "1534528741775-53994a69daeb",
+        "1507003211169-0a1dd7228f2d", "1522075469751-3a6694fb2f61", "1544717305-2782549b5136",
+        "1554151228-14d9def656e4", "1531746020798-e6953c6e8e04", "1500648767791-00dcc994a43e",
+        "1508214751196-bcfd4ca60f91", "1494790108377-be9c29b29330", "1517841905240-472988babdf9",
+        "1539571696357-5a69c17a67c6", "1438761681033-6461ffad8d80", "1524504388940-b1c1722653e1",
+        "1519085360753-af0119f7cbe7", "1491528920044-4531310b6531", "1503023344727-8982f00d2947",
+        "1534308983496-4fabb1a015ee", "1542206395-9feb3edaa68d", "1501196354995-cbb51c65aaea",
+        "1506863530036-1775a06bfa37", "1508214751196-bcfd4ca60f91", "1513956589380-bad6acb9b9d4",
+        "1519345182560-3f2917c472ef", "1520155707334-757655122b38", "1530577197743-7adf14294584",
+        "1531123897727-8f129e1688ce", "1539571696357-5a69c17a67c6", "1548142813-c348350df52b"
+    ]
+    
+    headers = {
+        "User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/120.0.0.0 Safari/537.36"
+    }
+    
+    print("Downloading 30 face images from Unsplash (512x512 cropped)...")
+    for i, img_id in enumerate(unsplash_ids):
+        filename = f"{i:05d}.png"
+        filepath = os.path.join(dataset_dir, filename)
+        if os.path.exists(filepath):
+            print(f"Skipping {filename} (already exists)")
+            continue
+            
+        url = f"https://images.unsplash.com/photo-{img_id}?w=512&h=512&fit=crop&q=80"
+        try:
+            r = requests.get(url, headers=headers, timeout=15)
+            if r.status_code == 200:
+                with open(filepath, "wb") as f:
+                    f.write(r.content)
+                print(f"Downloaded {filename}")
+            else:
+                print(f"Failed to download {filename} (Status: {r.status_code})")
+        except Exception as e:
+            print(f"Error downloading {filename}: {e}")
+            
+    # 3. Create local toy configurations for Stages I, II, and III
+    configs_to_make = [
+        ("VQGAN_512_ds32_nearest_stage1.yml", "VQGAN_toy.yml"),
+        ("CodeFormer_stage2.yml", "CodeFormer_stage2_toy.yml"),
+        ("CodeFormer_stage3.yml", "CodeFormer_stage3_toy.yml")
+    ]
+    
+    for orig_name, toy_name in configs_to_make:
+        orig_path = os.path.join(codeformer_dir, "options", orig_name)
+        toy_path = os.path.join(codeformer_dir, "options", toy_name)
+        
+        if not os.path.exists(orig_path):
+            print(f"Error: Original config {orig_name} not found!")
+            continue
+            
+        with open(orig_path, "r", encoding="utf-8") as f:
+            config = yaml.safe_load(f)
+            
+        # Modify for local CPU toy training
+        config["num_gpu"] = 0
+        config["dist_params"] = None
+        config["dist"] = False
+        
+        if "datasets" in config:
+            for phase in config["datasets"]:
+                dataset = config["datasets"][phase]
+                dataset["num_worker_per_gpu"] = 0  # No multiprocessing on CPU (avoid pickling/Windows issues)
+                dataset["batch_size_per_gpu"] = 1
+                dataset["dataset_enlarge_ratio"] = 1  # Don't enlarge dataset, keep it small
+                if "prefetch_mode" in dataset:
+                    dataset["prefetch_mode"] = "cpu"
+                    
+        if "train" in config:
+            config["train"]["total_iter"] = 5  # Train for only 5 iterations to test
+            
+        if "logger" in config:
+            config["logger"]["print_freq"] = 1
+            config["logger"]["save_checkpoint_freq"] = 5
+            config["logger"]["use_tb_logger"] = False
+            config["logger"]["wandb"] = None
+            
+        with open(toy_path, "w", encoding="utf-8") as f:
+            yaml.dump(config, f, default_flow_style=False, sort_keys=False)
+        print(f"Created local toy config: {toy_path}")
+        
+    print("\nSUCCESS: Toy training setup complete!")
+
+if __name__ == "__main__":
+    main()