feat: implement custom training configurations and runner script for model improvement

handover.md

@@ -92,3 +92,83 @@
 - **Files Staged:** `app.py`, `pipeline.py`, `download_weights.py`, `handover.md`
 - **Commit Message:** "feat: integrate Real-ESRGAN background upscaling and face detection threshold"
 - **Remote Push:** Scheduled for execution.
+
+---
+
+## Task 5: Peak End-to-End Model Improvement Plan
+
+### Overview
+This plan describes the comprehensive, peak end-to-end strategy to improve and fine-tune the CodeFormer face restoration model on custom target domain datasets, covering data preparation, degradation pipeline adjustment, advanced loss selection, distributed training, validation, and integration.
+
+---
+
+### Step 1: Data Preparation & Preprocessing Pipeline
+To fine-tune the model, you need a high-quality (HQ) training dataset. If you have low-quality (LQ) images, you also need to align them.
+1. **Acquire HQ Face Dataset:** Prepare 2,000 - 10,000 high-quality face images (e.g. from your target domain or high-res portraits).
+2. **Crop & Align Faces:**
+   Run the face detection and alignment helper to crop faces to $512 \times 512$ pixels:
+   ```bash
+   python models/CodeFormer/scripts/crop_align_face.py -i <input_raw_images_dir> -o <output_aligned_faces_dir>
+   ```
+3. **Data Splitting:** Divide aligned faces into training (90%), validation (5%), and test (5%) splits. Store them under `models/CodeFormer/datasets/custom_dataset/`.
+
+---
+
+### Step 2: Degradation Modeling Customization
+Modify the blind dataset configurations in your custom training option file (e.g. `CodeFormer_stage3_custom.yml`) to represent target real-world degradations:
+- **Motion Blur:** Set `motion_kernel_prob` and add motion blur kernels to model camera movement.
+- **Gaussian Blur:** Modify `blur_kernel_size` and `blur_sigma` to match degradation level.
+- **Noise:** Add Poisson and Gaussian noise with custom parameters (`noise_range` or `noise_range_large`).
+- **JPEG Compression:** Decrease the minimum of `jpeg_range` if dealing with high compression blockiness.
+
+---
+
+### Step 3: Architecture & Fine-Tuning Scenarios
+Depending on your project's goals, select one of the following training pathways:
+- **Scenario A: CFT Module Fine-Tuning (Stage III) - Recommended First Step**
+  - Keeps Stage 1 (VQGAN) and Stage 2 (Transformer) frozen. Fine-tunes the controllable feature transformation layers to balance likeness (fidelity) and quality.
+  - Very stable, relatively fast, and requires less GPU memory.
+- **Scenario B: Transformer & CFT Fine-Tuning (Stage II & III)**
+  - Fine-tunes the lookup transformer to map distorted inputs to the clean codebook indices.
+  - Useful if the degradations are highly non-linear or stylized (e.g. cartoons, oil paintings).
+- **Scenario C: Full VQGAN + Transformer Retraining (Stage I, II & III)**
+  - Re-trains the VQGAN codebook representation from scratch.
+  - Necessary only if restoring non-human faces (e.g., animal faces, fictional creatures).
+
+---
+
+### Step 4: Advanced Loss Function Adjustments
+To enhance qualitative results and identity preservation:
+1. **Identity Preservation (ArcFace Loss):** Integrate an ArcFace feature extractor to compute Cosine Similarity between restored and original faces:
+   $$\mathcal{L}_{id} = 1 - \cos(\text{ArcFace}(I_{rec}), \text{ArcFace}(I_{HQ}))$$
+2. **Structural & Detail Control:**
+   - **Perceptual (LPIPS) Loss:** Retain at weight `1.0` for natural textures.
+   - **GAN Loss:** Use Hinge GAN Loss (`loss_weight: 0.1`) to generate sharp details without artifacts.
+   - **Pixel (L1) Loss:** Retain at weight `1.0` to avoid drift in color/lighting.
+
+---
+
+### Step 5: Distributed GPU Training Setup
+For official training, use GPU(s) with CUDA:
+1. **Create Option File:** Save configuration to [CodeFormer_stage3_custom.yml](file:///c:/Users/admin/.gemini/antigravity-ide/scratch/custom-ai-enhancer/models/CodeFormer/options/CodeFormer_stage3_custom.yml). Set `num_gpu: 1` (or more).
+2. **Execute Training via torchrun (Distributed):**
+   ```bash
+   torchrun --nproc_per_node=gpu_num models/CodeFormer/basicsr/train.py -opt models/CodeFormer/options/CodeFormer_stage3_custom.yml --launcher pytorch
+   ```
+3. **Mixed Precision (AMP):** Enable AMP to save memory and speed up computation.
+
+---
+
+### Step 6: Evaluation & Metrics Validation
+Validate checkpoints quantitatively and qualitatively:
+- **PSNR / SSIM:** Measure reconstruction fidelity.
+- **LPIPS:** Measure perceptual closeness to human vision.
+- **FID:** Measure distribution quality of generated faces.
+- **ArcFace Cosine similarity:** Validate face identity preservation.
+
+---
+
+### Step 7: Streamlit Integration
+1. Export the best trained checkpoint (`params_ema` key) from `experiments/` to `weights/CodeFormer/codeformer_custom.pth`.
+2. Update [pipeline.py](file:///c:/Users/admin/.gemini/antigravity-ide/scratch/custom-ai-enhancer/pipeline.py) to point to the new model weights.
+3. Update [app.py](file:///c:/Users/admin/.gemini/antigravity-ide/scratch/custom-ai-enhancer/app.py) to add a model-selection dropdown or toggle, letting users compare the vanilla CodeFormer against your custom fine-tuned model.

models/CodeFormer/options/CodeFormer_stage3_custom.yml

@@ -0,0 +1,164 @@
+name: CodeFormer_stage3_custom
+model_type: CodeFormerJointModel
+num_gpu: 0
+manual_seed: 0
+datasets:
+  train:
+    name: CustomDataset
+    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: true
+    motion_kernel_prob: 0.05
+    kernel_list:
+    - iso
+    - aniso
+    kernel_prob:
+    - 0.5
+    - 0.5
+    blur_sigma:
+    - 0.1
+    - 10.0
+    downsample_range:
+    - 1.0
+    - 12.0
+    noise_range:
+    - 0.0
+    - 20.0
+    jpeg_range:
+    - 50
+    - 100
+    blur_sigma_large:
+    - 1.0
+    - 15.0
+    downsample_range_large:
+    - 4.0
+    - 30.0
+    noise_range_large:
+    - 0.0
+    - 30.0
+    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: ../../weights/CodeFormer/codeformer.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: 10
+  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:
+    project: null
+dist: false
+dist_params: null

train_custom.py

@@ -0,0 +1,99 @@
+import os
+import sys
+import torch
+import yaml
+import subprocess
+
+def main():
+    project_dir = os.path.dirname(os.path.abspath(__file__))
+    codeformer_dir = os.path.join(project_dir, "models", "CodeFormer")
+    
+    print("=== Custom CodeFormer Training Runner ===")
+    
+    # 1. Check GPU availability
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    num_gpus = torch.cuda.device_count() if device == "cuda" else 0
+    print(f"Device detected: {device.upper()}")
+    print(f"Number of GPUs available: {num_gpus}")
+    
+    # 2. Check and prepare dataset
+    dataset_dir = os.path.join(codeformer_dir, "datasets", "ffhq", "ffhq_512")
+    if not os.path.exists(dataset_dir) or len(os.listdir(dataset_dir)) == 0:
+        print("Dataset directory is empty. Preparing dataset images...")
+        try:
+            import prepare_toy_training
+            prepare_toy_training.main()
+            print("Dataset preparation completed.")
+        except Exception as e:
+            print(f"Error preparing dataset: {e}")
+            sys.exit(1)
+    else:
+        print(f"Dataset found at {dataset_dir} ({len(os.listdir(dataset_dir))} images).")
+        
+    # 3. Update configuration file
+    config_path = os.path.join(codeformer_dir, "options", "CodeFormer_stage3_custom.yml")
+    if not os.path.exists(config_path):
+        print(f"Error: Config file not found at {config_path}")
+        sys.exit(1)
+        
+    print(f"Reading configuration from {config_path}...")
+    with open(config_path, "r", encoding="utf-8") as f:
+        config = yaml.safe_load(f)
+        
+    # Dynamically set GPU count
+    config["num_gpu"] = num_gpus
+    config["dist"] = False
+    config["dist_params"] = None
+    
+    # Force single worker on Windows / CPU to prevent pickling issues
+    if "datasets" in config:
+        for phase in config["datasets"]:
+            dataset = config["datasets"][phase]
+            dataset["num_worker_per_gpu"] = 0
+            if device == "cpu":
+                dataset["prefetch_mode"] = "cpu"
+                
+    # Update weights path if they are in the project weights folder
+    project_weights_path = os.path.join(project_dir, "weights", "CodeFormer", "codeformer.pth")
+    if os.path.exists(project_weights_path):
+        # basicSR is relative to the running dir which is models/CodeFormer
+        config["path"]["pretrain_network_g"] = "../../weights/CodeFormer/codeformer.pth"
+        print(f"Configured pretrain generator path to: {config['path']['pretrain_network_g']}")
+    
+    # Write back the updated configuration
+    with open(config_path, "w", encoding="utf-8") as f:
+        yaml.dump(config, f, default_flow_style=False, sort_keys=False)
+    print(f"Updated configuration file for device={device.upper()}, num_gpu={num_gpus}")
+    
+    # 4. Run the training process
+    train_script = os.path.join("basicsr", "train.py")
+    cmd = [
+        sys.executable,
+        train_script,
+        "-opt",
+        os.path.join("options", "CodeFormer_stage3_custom.yml"),
+        "--launcher",
+        "none"
+    ]
+    
+    # Add models/CodeFormer to PYTHONPATH
+    env = os.environ.copy()
+    env["PYTHONPATH"] = os.path.pathsep.join([codeformer_dir, env.get("PYTHONPATH", "")])
+    
+    print("\nStarting training process. Command:")
+    print(" ".join(cmd))
+    print(f"Working directory: {codeformer_dir}")
+    print("------------------------------------------")
+    
+    try:
+        # Run subprocess under models/CodeFormer working directory
+        result = subprocess.run(cmd, cwd=codeformer_dir, env=env, check=True)
+        print("------------------------------------------")
+        print("Training execution completed successfully!")
+    except subprocess.CalledProcessError as e:
+        print("------------------------------------------")
+        print(f"Training failed with exit code: {e.returncode}")
+        sys.exit(e.returncode)
+
+if __name__ == "__main__":
+    main()