feat: add multi-category dataset crawler + Real-ESRGAN training pipeline

- crawl_datasets.py: crawl face/landscape/anime images from HuggingFace
(FFHQ, mertcobanov/nature-dataset, amirali900/Anime-Face-Dataset-10k)
- train_realesrgan.py: setup + train Real-ESRGAN for general image enhancement
(landscape, anime, scenery) alongside CodeFormer face model
- train_sharpening_model.py: sharpening model training script
- scan_c_drive.py: disk scan utility
- CodeFormer_stage3_custom.yml: bump total_iter 350->600 to continue training

crawl_datasets.py

@@ -0,0 +1,193 @@
+import os
+import sys
+import io
+import pandas as pd
+from PIL import Image
+import requests
+import time
+
+# ─────────────────────────────────────────────
+# HuggingFace dataset sources
+# ─────────────────────────────────────────────
+SOURCES = {
+    "face": {
+        "parquet_url": "https://huggingface.co/datasets/Ryan-sjtu/ffhq512-caption/resolve/main/data/train-00000-of-00054-9b5f7c3e6bc03b3b.parquet",
+        "num_images": 300,
+        "subdir": "faces",
+        "prefix": "face",
+        "description": "FFHQ 512x512 real face photos",
+    },
+    "landscape": {
+        "parquet_url": "https://huggingface.co/datasets/mertcobanov/nature-dataset/resolve/refs%2Fconvert%2Fparquet/default/train/0000.parquet",
+        "num_images": 300,
+        "subdir": "landscapes",
+        "prefix": "landscape",
+        "description": "Nature & landscape scenery (50k images)",
+    },
+    "anime": {
+        "parquet_url": "https://huggingface.co/datasets/amirali900/Anime-Face-Dataset-10k/resolve/refs%2Fconvert%2Fparquet/default/train/0000.parquet",
+        "num_images": 300,
+        "subdir": "anime",
+        "prefix": "anime",
+        "description": "Anime Face Dataset 10k illustrations",
+    },
+}
+
+# Fallback: direct image download lists for anime / landscape
+ANIME_FALLBACK_URLS = [
+    "https://huggingface.co/datasets/huggan/anime-faces/resolve/main/data/train-00000-of-00001.parquet",
+]
+
+LANDSCAPE_FALLBACK_URLS = [
+    "https://huggingface.co/datasets/jlbaker361/flickr_humans/resolve/main/data/train-00000-of-00001.parquet",
+]
+
+
+def load_parquet_safe(url: str) -> pd.DataFrame | None:
+    """Try to load a parquet file from URL, return None on failure."""
+    print(f"  Loading: {url}")
+    try:
+        df = pd.read_parquet(url)
+        print(f"  [OK] Loaded {len(df)} rows.")
+        return df
+    except Exception as e:
+        print(f"  [FAIL] Failed: {e}")
+        return None
+
+
+def find_image_column(df: pd.DataFrame) -> str | None:
+    """Detect which column holds image data (bytes dict or PIL-compatible)."""
+    for col in df.columns:
+        sample = df[col].iloc[0]
+        if isinstance(sample, dict) and "bytes" in sample:
+            return col
+        if isinstance(sample, bytes):
+            return col
+    return None
+
+
+def extract_image_bytes(cell) -> bytes | None:
+    """Extract raw image bytes from a dataframe cell regardless of format."""
+    if isinstance(cell, dict):
+        return cell.get("bytes")
+    if isinstance(cell, bytes):
+        return cell
+    return None
+
+
+def crawl_source(name: str, cfg: dict, base_dataset_dir: str):
+    """Download images for a single source category."""
+    save_dir = os.path.join(base_dataset_dir, cfg["subdir"])
+    os.makedirs(save_dir, exist_ok=True)
+
+    print(f"\n{'='*55}")
+    print(f"  [{name.upper()}] {cfg['description']}")
+    print(f"  Save dir : {save_dir}")
+    print(f"  Target   : {cfg['num_images']} images")
+    print(f"{'='*55}")
+
+    df = load_parquet_safe(cfg["parquet_url"])
+
+    # Try fallbacks if primary failed
+    if df is None:
+        fallbacks = []
+        if name == "anime":
+            fallbacks = ANIME_FALLBACK_URLS
+        elif name == "landscape":
+            fallbacks = LANDSCAPE_FALLBACK_URLS
+        for fb_url in fallbacks:
+            print(f"  Trying fallback: {fb_url}")
+            df = load_parquet_safe(fb_url)
+            if df is not None:
+                break
+
+    if df is None:
+        print(f"  [FAIL] All sources failed for '{name}'. Skipping.")
+        return 0
+
+    img_col = find_image_column(df)
+    if img_col is None:
+        print(f"  [FAIL] No image column detected in dataset. Columns: {list(df.columns)}")
+        return 0
+
+    print(f"  Image column: '{img_col}'")
+
+    # Find existing highest index in this subdir to avoid overwrites
+    existing = [
+        f for f in os.listdir(save_dir)
+        if f.endswith(".png") and f.startswith(cfg["prefix"])
+    ]
+    highest_idx = -1
+    for f in existing:
+        try:
+            idx = int(f.replace(cfg["prefix"] + "_", "").replace(".png", ""))
+            highest_idx = max(highest_idx, idx)
+        except ValueError:
+            continue
+
+    start_idx = highest_idx + 1
+    saved = 0
+    target = cfg["num_images"]
+
+    for i in range(len(df)):
+        if saved >= target:
+            break
+        try:
+            cell = df[img_col].iloc[i]
+            img_bytes = extract_image_bytes(cell)
+            if not img_bytes:
+                continue
+
+            img = Image.open(io.BytesIO(img_bytes)).convert("RGB")
+            img = img.resize((512, 512), Image.Resampling.LANCZOS)
+
+            filename = f"{cfg['prefix']}_{start_idx + saved:05d}.png"
+            filepath = os.path.join(save_dir, filename)
+            img.save(filepath, "PNG")
+            saved += 1
+
+            if saved % 20 == 0 or saved == target:
+                print(f"  [{name}] {saved}/{target} images saved...")
+
+        except Exception as e:
+            print(f"  [warn] Row {i} error: {e}")
+            continue
+
+    print(f"\n  [OK] [{name.upper()}] Done: {saved} images saved to {save_dir}")
+    return saved
+
+
+def main():
+    project_dir = os.path.dirname(os.path.abspath(__file__))
+    codeformer_dir = os.path.join(project_dir, "models", "CodeFormer")
+    base_dataset_dir = os.path.join(codeformer_dir, "datasets", "ffhq", "ffhq_512")
+    os.makedirs(base_dataset_dir, exist_ok=True)
+
+    print("=" * 55)
+    print("  MULTI-CATEGORY DATASET CRAWLER")
+    print("  Categories: Face | Landscape | Anime Girl")
+    print("=" * 55)
+
+    total_saved = 0
+    for name, cfg in SOURCES.items():
+        saved = crawl_source(name, cfg, base_dataset_dir)
+        total_saved += saved
+        time.sleep(0.5)
+
+    # Count all images recursively
+    all_images = []
+    for root, _, files in os.walk(base_dataset_dir):
+        for f in files:
+            if f.lower().endswith(".png"):
+                all_images.append(f)
+
+    print("\n" + "=" * 55)
+    print(f"  CRAWL COMPLETE")
+    print(f"  New images this run : {total_saved}")
+    print(f"  Total dataset size  : {len(all_images)} images")
+    print(f"  Dataset location    : {base_dataset_dir}")
+    print("=" * 55)
+
+
+if __name__ == "__main__":
+    main()

models/CodeFormer/options/CodeFormer_stage3_custom.yml

@@ -1,10 +1,7 @@
-# general settings
 name: CodeFormer_stage3_custom
 model_type: CodeFormerJointModel
 num_gpu: 0
 manual_seed: 0
-
-# dataset and data loader settings
 datasets:
   train:
     name: CustomDataset
@@ -13,142 +10,156 @@ datasets:
     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]
+    mean:
+    - 0.5
+    - 0.5
+    - 0.5
+    std:
+    - 0.5
+    - 0.5
+    - 0.5
     use_hflip: true
     use_corrupt: true
-
-    # Comprehensive degradation parameters (Motion blur, Gaussian blur, Downsampling, Noise, JPEG)
     blur_kernel_size: 41
     use_motion_kernel: true
     motion_kernel_prob: 0.05
-    kernel_list: ['iso', 'aniso']
-    kernel_prob: [0.5, 0.5]
-    
-    # Small degradation range
-    blur_sigma: [0.1, 10.0]
-    downsample_range: [1.0, 12.0]
-    noise_range: [0.0, 20.0]
-    jpeg_range: [50, 100]
-
-    # Large degradation range
-    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]
-
+    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
-
-    # Data loader configurations
     num_worker_per_gpu: 0
     batch_size_per_gpu: 1
     dataset_enlarge_ratio: 1
     prefetch_mode: cpu
-
-# network structures
 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']
-
+  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'
+  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 settings
 path:
-  pretrain_network_g: ../../weights/CodeFormer/codeformer.pth
+  pretrain_network_g: null
   param_key_g: params_ema
   strict_load_g: false
   pretrain_network_d: null
-  resume_state: null
-
-# training loss and optimizer settings
+  resume_state: D:\.gemini-scratch\custom-ai-enhancer\models\CodeFormer\experiments\20260708_201102_CodeFormer_stage3_custom\training_states\350.state
 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: !!float 5e-5
+    lr: 5.0e-05
     weight_decay: 0
-    betas: [0.9, 0.99]
+    betas:
+    - 0.9
+    - 0.99
   optim_d:
     type: Adam
-    lr: !!float 5e-5
+    lr: 5.0e-05
     weight_decay: 0
-    betas: [0.9, 0.99]
-
+    betas:
+    - 0.9
+    - 0.99
   scheduler:
     type: CosineAnnealingRestartLR
-    periods: [150000]
-    restart_weights: [1]
-    eta_min: !!float 2e-5
-
-  total_iter: 100
+    periods:
+    - 150000
+    restart_weights:
+    - 1
+    eta_min: 2.0e-05
+  total_iter: 600
   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: !!float 1.0
-
+    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
-
-# validation settings
 val:
-  val_freq: !!float 5e10
+  val_freq: 50000000000.0
   save_img: true
   metrics:
     psnr:
       type: calculate_psnr
       crop_border: 4
       test_y_channel: false
-
-# logging settings
 logger:
   print_freq: 1
   save_checkpoint_freq: 5
   use_tb_logger: false
   wandb:
     project: null
-
 find_unused_parameters: true
 dist_params: null
 dist: false

scan_c_drive.py

@@ -0,0 +1,74 @@
+import os
+import sys
+
+def get_dir_size(path, max_depth=4, current_depth=0):
+    total = 0
+    if current_depth > max_depth:
+        return 0
+    try:
+        for entry in os.scandir(path):
+            try:
+                if entry.is_file(follow_symlinks=False):
+                    total += entry.stat().st_size
+                elif entry.is_dir(follow_symlinks=False):
+                    total += get_dir_size(entry.path, max_depth, current_depth + 1)
+            except Exception:
+                continue
+    except Exception:
+        pass
+    return total
+
+def main():
+    user_home = os.path.expanduser("~")
+    print(f"Scanning large directories in: {user_home}...")
+    
+    large_dirs = []
+    
+    # List top level directories in user home
+    try:
+        entries = list(os.scandir(user_home))
+    except Exception as e:
+        print(f"Error reading home directory: {e}")
+        return
+        
+    for entry in entries:
+        if entry.is_dir(follow_symlinks=False):
+            # Skip some standard tiny folders or system symbolic links to speed up
+            if entry.name.startswith('.') and entry.name not in ['.cache', '.conda', '.docker', '.npm', '.nuget', '.vscode', '.gradle']:
+                continue
+            
+            print(f"Checking size of {entry.name}...")
+            size = get_dir_size(entry.path, max_depth=5)
+            size_gb = size / (1024 * 1024 * 1024)
+            if size_gb > 0.5:  # larger than 500 MB
+                large_dirs.append((entry.path, size_gb))
+                print(f"-> Found large dir: {entry.name} ({size_gb:.2f} GB)")
+                
+    print("\n=== Scan Results (Directories > 0.5 GB) ===")
+    large_dirs.sort(key=lambda x: x[1], reverse=True)
+    for path, size_gb in large_dirs:
+        print(f"- {path}: {size_gb:.2f} GB")
+        
+        # If it is AppData, scan one level deeper to pinpoint the culprit
+        if "AppData" in path:
+            print("  Pinpointing AppData contents...")
+            for sub_name in ["Local", "Roaming", "LocalLow"]:
+                sub_path = os.path.join(path, sub_name)
+                if os.path.exists(sub_path):
+                    sub_size = get_dir_size(sub_path, max_depth=4)
+                    sub_size_gb = sub_size / (1024 * 1024 * 1024)
+                    if sub_size_gb > 0.5:
+                        print(f"    - {sub_path}: {sub_size_gb:.2f} GB")
+                        # Pinpoint Local/Roaming further
+                        try:
+                            for local_entry in os.scandir(sub_path):
+                                if local_entry.is_dir(follow_symlinks=False):
+                                    l_size = get_dir_size(local_entry.path, max_depth=3)
+                                    l_size_gb = l_size / (1024 * 1024 * 1024)
+                                    if l_size_gb > 0.5:
+                                        print(f"      - {local_entry.path}: {l_size_gb:.2f} GB")
+                        except Exception:
+                            pass
+
+if __name__ == "__main__":
+    main()

train_realesrgan.py

@@ -0,0 +1,335 @@
+"""
+Train Real-ESRGAN for general image enhancement (landscape, anime, scenery).
+This script:
+  1. Clones Real-ESRGAN repo if not present
+  2. Prepares a training config pointing to our crawled dataset
+  3. Launches the training process
+"""
+import os
+import sys
+import subprocess
+import shutil
+import yaml
+import glob
+
+PROJECT_DIR = os.path.dirname(os.path.abspath(__file__))
+REALESRGAN_DIR = os.path.join(PROJECT_DIR, "models", "Real-ESRGAN")
+
+# Dataset paths (populated by crawl_datasets.py)
+DATASET_BASE = os.path.join(
+    PROJECT_DIR, "models", "CodeFormer", "datasets", "ffhq", "ffhq_512"
+)
+LANDSCAPE_DIR = os.path.join(DATASET_BASE, "landscapes")
+ANIME_DIR = os.path.join(DATASET_BASE, "anime")
+FACE_DIR = os.path.join(DATASET_BASE, "faces")
+
+# Combined GT (high-quality) folder for Real-ESRGAN training
+REALESRGAN_GT_DIR = os.path.join(PROJECT_DIR, "datasets", "realesrgan_gt")
+
+
+def clone_realesrgan():
+    """Clone Real-ESRGAN repository if not already present."""
+    if os.path.exists(REALESRGAN_DIR):
+        print("[OK] Real-ESRGAN repo already present.")
+        return True
+    print("Cloning Real-ESRGAN repository...")
+    result = subprocess.run(
+        ["git", "clone", "https://github.com/xinntao/Real-ESRGAN.git", REALESRGAN_DIR],
+        capture_output=False,
+    )
+    if result.returncode != 0:
+        print("[FAIL] Failed to clone Real-ESRGAN.")
+        return False
+    print("[OK] Real-ESRGAN cloned successfully.")
+    return True
+
+
+def install_realesrgan_deps():
+    """Install Real-ESRGAN dependencies."""
+    print("Installing Real-ESRGAN dependencies...")
+    requirements_file = os.path.join(REALESRGAN_DIR, "requirements.txt")
+    if os.path.exists(requirements_file):
+        subprocess.run(
+            [sys.executable, "-m", "pip", "install", "-r", requirements_file],
+            check=True,
+        )
+    # Also install basicsr for Real-ESRGAN
+    subprocess.run(
+        [sys.executable, "-m", "pip", "install", "basicsr", "facexlib", "gfpgan"],
+        check=True,
+    )
+    print("[OK] Dependencies installed.")
+
+
+def prepare_gt_dataset():
+    """Combine landscape + anime + face images into a single GT folder."""
+    os.makedirs(REALESRGAN_GT_DIR, exist_ok=True)
+
+    total_copied = 0
+    sources = {
+        "landscape": LANDSCAPE_DIR,
+        "anime": ANIME_DIR,
+        "face": FACE_DIR,
+    }
+
+    for category, src_dir in sources.items():
+        if not os.path.exists(src_dir):
+            print(f"  [skip] {category} dir not found: {src_dir}")
+            continue
+        images = glob.glob(os.path.join(src_dir, "*.png"))
+        print(f"  Copying {len(images)} {category} images...")
+        for img_path in images:
+            dest = os.path.join(REALESRGAN_GT_DIR, os.path.basename(img_path))
+            if not os.path.exists(dest):
+                shutil.copy2(img_path, dest)
+                total_copied += 1
+
+    # Also copy root-level face images (original FFHQ)
+    root_images = glob.glob(os.path.join(DATASET_BASE, "*.png"))
+    print(f"  Copying {len(root_images)} root FFHQ face images...")
+    for img_path in root_images:
+        dest = os.path.join(REALESRGAN_GT_DIR, os.path.basename(img_path))
+        if not os.path.exists(dest):
+            shutil.copy2(img_path, dest)
+            total_copied += 1
+
+    total = len([f for f in os.listdir(REALESRGAN_GT_DIR) if f.endswith(".png")])
+    print(f"[OK] GT dataset ready: {total} images (copied {total_copied} new)")
+    return total
+
+
+def create_training_config(num_images: int) -> str:
+    """Create a Real-ESRGAN training YAML config."""
+    config_dir = os.path.join(REALESRGAN_DIR, "options")
+    os.makedirs(config_dir, exist_ok=True)
+    config_path = os.path.join(config_dir, "train_realesrgan_custom.yml")
+
+    # Relative path from Real-ESRGAN dir to GT dataset
+    gt_path_rel = os.path.relpath(REALESRGAN_GT_DIR, REALESRGAN_DIR).replace("\\", "/")
+
+    config = {
+        "name": "train_RealESRGAN_custom",
+        "model_type": "RealESRGANModel",
+        "scale": 4,
+        "num_gpu": 0,
+        "manual_seed": 0,
+        "datasets": {
+            "train": {
+                "name": "CustomMixedDataset",
+                "type": "RealESRGANDataset",
+                "dataroot_gt": gt_path_rel,
+                "meta_info": None,
+                "io_backend": {"type": "disk"},
+                "gt_size": 256,
+                "use_hflip": True,
+                "use_rot": False,
+                "blur_kernel_size": 21,
+                "kernel_list": ["iso", "aniso", "generalized_iso", "generalized_aniso",
+                                "plateau_iso", "plateau_aniso"],
+                "kernel_prob": [0.45, 0.25, 0.12, 0.03, 0.12, 0.03],
+                "sinc_prob": 0.1,
+                "blur_sigma": [0.2, 3.0],
+                "betag_range": [0.5, 4.0],
+                "betap_range": [1, 2.0],
+                "blur_kernel_size2": 21,
+                "kernel_list2": ["iso", "aniso", "generalized_iso", "generalized_aniso",
+                                 "plateau_iso", "plateau_aniso"],
+                "kernel_prob2": [0.45, 0.25, 0.12, 0.03, 0.12, 0.03],
+                "sinc_prob2": 0.1,
+                "blur_sigma2": [0.2, 1.5],
+                "betag_range2": [0.5, 4.0],
+                "betap_range2": [1, 2.0],
+                "final_sinc_prob": 0.8,
+                "resize_prob": [0.2, 0.7, 0.1],
+                "resize_range": [0.15, 1.5],
+                "gaussian_noise_prob": 0.5,
+                "noise_range": [1, 30],
+                "poisson_scale_range": [0.05, 3.0],
+                "gray_noise_prob": 0.4,
+                "jpeg_range": [30, 95],
+                "resize_prob2": [0.3, 0.4, 0.3],
+                "resize_range2": [0.3, 1.2],
+                "gaussian_noise_prob2": 0.5,
+                "noise_range2": [1, 25],
+                "poisson_scale_range2": [0.05, 2.5],
+                "gray_noise_prob2": 0.4,
+                "jpeg_range2": [30, 95],
+                "gt_size_h": 256,
+                "gt_size_w": 256,
+                "batch_size_per_gpu": 4,
+                "num_worker_per_gpu": 0,
+                "dataset_enlarge_ratio": 1,
+                "prefetch_mode": "cpu",
+                "num_prefetch_queue": 1,
+            }
+        },
+        "network_g": {
+            "type": "RRDBNet",
+            "num_in_ch": 3,
+            "num_out_ch": 3,
+            "num_feat": 64,
+            "num_block": 23,
+            "num_grow_ch": 32,
+            "scale": 4,
+        },
+        "network_d": {
+            "type": "UNetDiscriminatorSN",
+            "num_in_ch": 3,
+            "num_feat": 64,
+            "skip_connection": True,
+        },
+        "path": {
+            "pretrain_network_g": None,
+            "param_key_g": "params_ema",
+            "strict_load_g": False,
+            "resume_state": None,
+        },
+        "train": {
+            "ema_decay": 0.999,
+            "optim_g": {
+                "type": "Adam",
+                "lr": 1.0e-4,
+                "weight_decay": 0,
+                "betas": [0.9, 0.99],
+            },
+            "optim_d": {
+                "type": "Adam",
+                "lr": 1.0e-4,
+                "weight_decay": 0,
+                "betas": [0.9, 0.99],
+            },
+            "scheduler": {
+                "type": "MultiStepLR",
+                "milestones": [400000],
+                "gamma": 0.5,
+            },
+            "total_iter": 500,
+            "warmup_iter": -1,
+            "pixel_opt": {
+                "type": "L1Loss",
+                "loss_weight": 1.0,
+                "reduction": "mean",
+            },
+            "perceptual_opt": {
+                "type": "PerceptualLoss",
+                "layer_weights": {"conv1_2": 0.1, "conv2_2": 0.1, "conv3_4": 1.0,
+                                  "conv4_4": 1.0, "conv5_4": 1.0},
+                "vgg_type": "vgg19",
+                "use_input_norm": True,
+                "range_norm": False,
+                "perceptual_weight": 1.0,
+                "style_weight": 0.0,
+                "criterion": "l1",
+            },
+            "gan_opt": {
+                "type": "GANLoss",
+                "gan_type": "vanilla",
+                "real_label_val": 1.0,
+                "fake_label_val": 0.0,
+                "loss_weight": 1.0e-1,
+            },
+            "net_d_iters": 1,
+            "net_d_start_iter": 0,
+        },
+        "val": {
+            "val_freq": 500,
+            "save_img": True,
+        },
+        "logger": {
+            "print_freq": 1,
+            "save_checkpoint_freq": 50,
+            "use_tb_logger": False,
+            "wandb": {"project": None},
+        },
+        "dist_params": None,
+        "dist": False,
+        "find_unused_parameters": False,
+    }
+
+    with open(config_path, "w", encoding="utf-8") as f:
+        yaml.dump(config, f, default_flow_style=False, sort_keys=False)
+
+    print(f"[OK] Config written: {config_path}")
+    return config_path
+
+
+def find_latest_state() -> str | None:
+    """Find the latest training state checkpoint."""
+    pattern = os.path.join(
+        REALESRGAN_DIR, "experiments", "*train_RealESRGAN_custom*",
+        "training_states", "*.state"
+    )
+    states = glob.glob(pattern)
+    if not states:
+        return None
+    latest = max(states, key=lambda p: int(os.path.basename(p).replace(".state", "")))
+    return latest
+
+
+def main():
+    print("=" * 55)
+    print("  Real-ESRGAN Custom Training Runner")
+    print("  Target: Landscape + Anime + Face enhancement")
+    print("=" * 55)
+
+    # 1. Clone repo
+    if not clone_realesrgan():
+        sys.exit(1)
+
+    # 2. Install deps
+    install_realesrgan_deps()
+
+    # 3. Prepare combined GT dataset
+    print("\n[Step 3] Preparing GT dataset...")
+    num_images = prepare_gt_dataset()
+    if num_images == 0:
+        print("[FAIL] No images found in dataset directories. Run crawl_datasets.py first.")
+        sys.exit(1)
+
+    # 4. Create training config
+    print("\n[Step 4] Creating training config...")
+    config_path = create_training_config(num_images)
+
+    # 5. Check for resume state
+    latest_state = find_latest_state()
+    if latest_state:
+        iter_num = int(os.path.basename(latest_state).replace(".state", ""))
+        print(f"\n>>> RESUME MODE: Found checkpoint at iteration {iter_num}")
+        # Patch config to resume
+        with open(config_path, "r", encoding="utf-8") as f:
+            cfg = yaml.safe_load(f)
+        cfg["path"]["resume_state"] = latest_state
+        cfg["path"]["pretrain_network_g"] = None
+        with open(config_path, "w", encoding="utf-8") as f:
+            yaml.dump(cfg, f, default_flow_style=False, sort_keys=False)
+    else:
+        print("\n>>> FRESH START: No previous checkpoint. Starting from scratch.")
+
+    # 6. Run training
+    train_script = os.path.join("basicsr", "train.py")
+    cmd = [
+        sys.executable,
+        train_script,
+        "-opt", os.path.join("options", "train_realesrgan_custom.yml"),
+        "--launcher", "none",
+    ]
+
+    env = os.environ.copy()
+    env["PYTHONPATH"] = os.path.pathsep.join([REALESRGAN_DIR, env.get("PYTHONPATH", "")])
+
+    print(f"\nStarting Real-ESRGAN training...")
+    print(f"Working directory: {REALESRGAN_DIR}")
+    print(f"Command: {' '.join(cmd)}")
+    print("-" * 55)
+
+    try:
+        subprocess.run(cmd, cwd=REALESRGAN_DIR, env=env, check=True)
+        print("-" * 55)
+        print("[OK] Real-ESRGAN training completed!")
+    except subprocess.CalledProcessError as e:
+        print(f"[FAIL] Training failed with exit code: {e.returncode}")
+        sys.exit(e.returncode)
+
+
+if __name__ == "__main__":
+    main()

train_sharpening_model.py

@@ -0,0 +1,90 @@
+import os
+import argparse
+import glob
+import time
+
+import torch
+import torch.nn as nn
+import torch.optim as optim
+from torch.utils.data import DataLoader, Dataset
+from torchvision import transforms
+from PIL import Image
+
+class ImageDataset(Dataset):
+    def __init__(self, root_dir, transform=None):
+        self.paths = sorted(glob.glob(os.path.join(root_dir, "*.png")))
+        self.transform = transform
+
+    def __len__(self):
+        return len(self.paths)
+
+    def __getitem__(self, idx):
+        img = Image.open(self.paths[idx]).convert("RGB")
+        if self.transform:
+            img = self.transform(img)
+        return img
+
+class SimpleUNet(nn.Module):
+    """Lightweight UNet‑style model for image sharpening (placeholder)."""
+    def __init__(self):
+        super().__init__()
+        self.encoder = nn.Sequential(
+            nn.Conv2d(3, 16, kernel_size=3, padding=1), nn.ReLU(inplace=True),
+            nn.Conv2d(16, 32, kernel_size=3, padding=1), nn.ReLU(inplace=True)
+        )
+        self.decoder = nn.Sequential(
+            nn.Conv2d(32, 16, kernel_size=3, padding=1), nn.ReLU(inplace=True),
+            nn.Conv2d(16, 3, kernel_size=3, padding=1)
+        )
+
+    def forward(self, x):
+        x = self.encoder(x)
+        x = self.decoder(x)
+        return x
+
+def train(dataset_dir, checkpoint_dir, epochs, batch_size, lr, use_gpu):
+    device = torch.device("cuda" if torch.cuda.is_available() and use_gpu else "cpu")
+    transform = transforms.Compose([
+        transforms.Resize((512, 512)),
+        transforms.ToTensor()
+    ])
+    dataset = ImageDataset(dataset_dir, transform=transform)
+    dataloader = DataLoader(dataset, batch_size=batch_size, shuffle=True, num_workers=2)
+
+    model = SimpleUNet().to(device)
+    criterion = nn.L1Loss()
+    optimizer = optim.Adam(model.parameters(), lr=lr)
+
+    print(f"Start training on {len(dataset)} images (device: {device})")
+    for epoch in range(1, epochs + 1):
+        epoch_loss = 0.0
+        for batch in dataloader:
+            batch = batch.to(device)
+            optimizer.zero_grad()
+            outputs = model(batch)
+            loss = criterion(outputs, batch)
+            loss.backward()
+            optimizer.step()
+            epoch_loss += loss.item() * batch.size(0)
+        epoch_loss /= len(dataset)
+        print(f"Epoch {epoch}/{epochs} - Loss: {epoch_loss:.4f}")
+        ckpt_path = os.path.join(checkpoint_dir, f"epoch_{epoch}.pth")
+        torch.save(model.state_dict(), ckpt_path)
+    final_path = os.path.join(checkpoint_dir, "sharpening_model_final.pth")
+    torch.save(model.state_dict(), final_path)
+    print("Training completed. Model saved to", final_path)
+
+def main():
+    parser = argparse.ArgumentParser(description="Train sharpening model on 512x512 PNG images.")
+    parser.add_argument("--dataset_dir", type=str, required=True, help="Folder containing PNG images.")
+    parser.add_argument("--checkpoint_dir", type=str, default="checkpoints", help="Directory to store checkpoints.")
+    parser.add_argument("--epochs", type=int, default=10, help="Number of epochs.")
+    parser.add_argument("--batch_size", type=int, default=8, help="Batch size.")
+    parser.add_argument("--lr", type=float, default=1e-4, help="Learning rate.")
+    parser.add_argument("--gpu", action="store_true", help="Use GPU if available.")
+    args = parser.parse_args()
+    os.makedirs(args.checkpoint_dir, exist_ok=True)
+    train(args.dataset_dir, args.checkpoint_dir, args.epochs, args.batch_size, args.lr, args.gpu)
+
+if __name__ == "__main__":
+    main()