import os import sys import cv2 import torch torch.set_num_threads(8) # Ryzen 7735HS has 8C/16T # oneDNN can be unstable on the Ryzen CPU during BasicSR training. torch.backends.mkldnn.enabled = False import yaml import subprocess import glob import argparse def main(): parser = argparse.ArgumentParser(description="Train CodeFormer with custom parameters.") parser.add_argument("--verify", action="store_true", help="Run 2 iterations for verification purposes.") args = parser.parse_args() 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) # Keep the deterministic benchmark holdout out of all training phases. # It is intentionally ignored by Git because it may identify private data; # create it with tools/prepare_benchmark.py before training. holdout_manifest = os.path.join(project_dir, "benchmarks", "holdout_paths.txt") if os.path.isfile(holdout_manifest): for dataset in config.get("datasets", {}).values(): dataset["exclude_manifest"] = holdout_manifest print(f"Excluding benchmark holdout from training: {holdout_manifest}") else: print("WARNING: benchmark holdout manifest not found; refusing quality claims for this run.") # Save original total_iter BEFORE any runtime overrides — we must restore # this after writing so --verify mode never permanently corrupts the file. original_total_iter = config.get("train", {}).get("total_iter", 20000) # Dynamically set GPU count config["num_gpu"] = num_gpus config["dist"] = False config["dist_params"] = None # Use multiple workers to saturate CPU cores when training on GPU; set to 0 on CPU to prevent Windows multiprocessing MemoryError/segfaults. if "datasets" in config: for phase in config["datasets"]: dataset = config["datasets"][phase] if device == "cpu": dataset["num_worker_per_gpu"] = 0 # MUST be null on CPU — 'cpu' prefetch spawns multiprocessing # workers which cause MemoryError/segfaults on Windows (Task 8). dataset["prefetch_mode"] = None else: dataset["num_worker_per_gpu"] = 4 # 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']}") # 4. Auto-detect resume state from the latest experiment checkpoint experiments_dir = os.path.join(codeformer_dir, "experiments") latest_state = None latest_state_iter = 0 # Search for existing experiment directories matching our config name pattern exp_pattern = os.path.join(experiments_dir, "*_CodeFormer_stage3_custom", "training_states", "*.state") state_files = glob.glob(exp_pattern) for state_file in state_files: basename = os.path.basename(state_file) try: iter_num = int(basename.replace(".state", "")) if iter_num > latest_state_iter: latest_state_iter = iter_num latest_state = state_file except ValueError: continue if latest_state: print(f"\n>>> RESUME MODE: Found checkpoint at iteration {latest_state_iter}") print(f" State file: {latest_state}") config["path"]["resume_state"] = latest_state config["path"]["pretrain_network_g"] = None if args.verify: config["train"]["total_iter"] = latest_state_iter + 2 print(f" Resuming training from iter {latest_state_iter} to {config['train']['total_iter']} (verification mode)...") else: if latest_state_iter >= config.get("train", {}).get("total_iter", 20000): print(f"\n>>> Training already completed ({latest_state_iter} >= {config.get('train', {}).get('total_iter', 20000)} total_iter)") sys.exit(0) print(f" Resuming training from iter {latest_state_iter} to {config['train']['total_iter']}...") else: print("\n>>> FRESH START: No previous checkpoint found. Starting from pretrained weights.") if args.verify: config["train"]["total_iter"] = 2 print(f" Training from scratch to {config['train']['total_iter']} (verification mode)...") else: print(f" Training from scratch to {config.get('train', {}).get('total_iter', 20000)}...") # 5. Write runtime config to a TEMP yml file — NEVER modify the original. # # Strategy: write all runtime overrides (num_gpu, prefetch_mode, resume_state, # total_iter) to a separate temp file. basicsr/train.py reads from that temp # file. The canonical yml with all comments is never touched again. # # This fixes two bugs permanently: # (a) --verify mode used to persist total_iter=checkpoint+2 to disc, causing # the next full run to stop after only 2 iterations. # (b) yaml.dump was stripping all comments from the original yml every run. import tempfile, shutil tmp_fd, tmp_config_path = tempfile.mkstemp(suffix=".yml", prefix="cf_runtime_", dir=os.path.join(codeformer_dir, "options")) try: with os.fdopen(tmp_fd, "w", encoding="utf-8") as f: yaml.dump(config, f, default_flow_style=False, sort_keys=False) print(f"Runtime config written to temp file: {os.path.basename(tmp_config_path)}") print(f" device={device.upper()}, num_gpu={num_gpus}, total_iter={config['train']['total_iter']}") # 6. Run the training process pointing at the temp config train_script = os.path.join("basicsr", "train.py") cmd = [ sys.executable, train_script, "-opt", os.path.relpath(tmp_config_path, codeformer_dir), "--launcher", "none" ] # Environment setup — CPU stability settings (Task 8 findings) env = os.environ.copy() for _k in ["OMP_NUM_THREADS", "MKL_NUM_THREADS", "OPENBLAS_NUM_THREADS", "NUMEXPR_NUM_THREADS", "VECLIB_MAXIMUM_THREADS"]: env[_k] = "8" # Match torch.set_num_threads — Ryzen 7735HS 8C/16T env["OPENCV_OPENCL_RUNTIME"] = "disabled" env["OPENCV_THREAD_LIMIT"] = "1" cv2.setNumThreads(0) 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: 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) finally: # Always clean up the temp file after the run (success or failure) if os.path.exists(tmp_config_path): os.remove(tmp_config_path) if __name__ == "__main__": main()