feat: complete custom training pipeline bugfixes, dataset crawling, and validation

app.py

@@ -440,6 +440,16 @@ with st.sidebar:
             except Exception as e:
                 st.error(f"Failed to import settings: {e}")
 
+    if st.button("🔄 Reset UI State"):
+        st.session_state.processing = False
+        st.session_state.enhanced_img = None
+        st.session_state.processing_error = None
+        st.session_state.last_run_params = None
+        if pipeline:
+            pipeline.cancel_flag = False
+        st.success("UI State reset successfully!")
+        st.rerun()
+
     # ── Training Dashboard ──────────────────────────────────────────────────────
     st.markdown("<div class='sidebar-section'>📊 Training Dashboard</div>", unsafe_allow_html=True)
     status = get_training_status()
@@ -605,35 +615,37 @@ with tab_single:
 
         st.markdown("<hr>", unsafe_allow_html=True)
 
-        start_time = time.time()
-        # Async processing using a background thread
-        if not st.session_state.get('processing'):
-            # Initialize state
-            st.session_state.processing = True
-            st.session_state.progress_state = {'stage': None, 'progress': 0.0, 'message': '', 'active': False}
-            
-            # Progress bar and status
-            progress_bar = st.progress(0)
-            status_text = st.empty()
-            stage_text = st.empty()
-            
-            # Cancel button
-            cancel_col, _ = st.columns([1, 3])
-            with cancel_col:
-                cancel_button = st.button("❌ Cancel", type="secondary")
-            
-            if cancel_button:
-                st.session_state.progress_state['cancelled'] = True
-                st.warning("⚠️ Processing cancelled by user")
-                st.stop()
-            
-            result_container = {}
-            def _run():
-                max_retries = 2
-                retry_count = 0
-                last_error = None
+        # Track parameters to auto-rerun if they change
+        current_params = {
+            'img_name': img_name,
+            'w': fidelity_weight,
+            'detection_model': face_detector,
+            'upscale': upscale_factor,
+            'blend_softness': blend_softness,
+            'bg_upsampler': 'realesrgan' if bg_upscale_toggle else None,
+            'det_threshold': det_threshold,
+            'sharpen_amount': sharpen_amount,
+            'face_upsample': face_upscale_toggle,
+            'face_restore': enable_face_restoration
+        }
+        
+        if st.session_state.get('last_run_params') != current_params:
+            st.session_state.enhanced_img = None
+            st.session_state.processing_error = None
+            st.session_state.processing = False
+            st.session_state.process_duration = None
+            if pipeline:
+                pipeline.cancel_flag = False
+
+        if st.session_state.enhanced_img is None and st.session_state.get('processing_error') is None:
+            if not st.session_state.get('processing'):
+                st.session_state.processing = True
+                st.session_state.progress_state = {'stage': 'initialization', 'progress': 0.0, 'message': 'Starting...', 'active': True, 'cancelled': False}
+                st.session_state.start_time = time.time()
+                if pipeline:
+                    pipeline.cancel_flag = False
                 
-                while retry_count <= max_retries:
+                def _run():
                     try:
                         result = pipeline.process_image(
                             img,
@@ -649,68 +661,78 @@ with tab_single:
                             batch_size=4,
                             face_restore=enable_face_restoration
                         )
-                        result_container['enhanced_img'] = result
-                        break
+                        if pipeline.cancel_flag:
+                            return
+                        st.session_state.enhanced_img = result
+                        st.session_state.process_duration = time.time() - st.session_state.start_time
+                        st.session_state.last_run_params = current_params
                     except Exception as e:
-                        last_error = e
-                        retry_count += 1
-                        if retry_count <= max_retries:
-                            st.session_state.progress_state['message'] = f"Retry {retry_count}/{max_retries}: {str(e)[:50]}..."
-                            time.sleep(1)
-                        else:
-                            result_container['error'] = str(e)
-                finally:
-                    st.session_state.processing = False
-                    st.session_state.progress_state['active'] = False
-            
-            threading.Thread(target=_run, daemon=True).start()
-            
-            # Wait for result (polling with progress updates)
-            stage_names = {
-                'initialization': '🔧 Initializing',
-                'detection': '👁️ Detecting Faces',
-                'background': '🖼️ Upscaling Background',
-                'restoration': '✨ Restoring Faces',
-                'blending': '🎨 Blending Faces',
-                'complete': '✅ Complete'
-            }
-            
-            while st.session_state.processing:
-                time.sleep(0.1)
-                progress_data = st.session_state.progress_state
+                        import traceback
+                        traceback.print_exc()
+                        st.session_state.processing_error = str(e)
+                    finally:
+                        st.session_state.processing = False
+                        st.session_state.progress_state['active'] = False
+
+                threading.Thread(target=_run, daemon=True).start()
+                st.rerun()
+
+            else:
+                # We are actively processing
+                progress_bar = st.progress(0)
+                status_text = st.empty()
+                stage_text = st.empty()
                 
-                # Check for cancellation
-                if progress_data.get('cancelled', False):
-                    st.session_state.processing = False
-                    st.warning("⚠️ Processing cancelled by user")
-                    progress_bar.empty()
-                    status_text.empty()
-                    stage_text.empty()
-                    st.stop()
+                stage_names = {
+                    'initialization': '🔧 Initializing',
+                    'detection': '👁️ Detecting Faces',
+                    'background': '🖼️ Upscaling Background',
+                    'restoration': '✨ Restoring Faces',
+                    'blending': '🎨 Blending Faces',
+                    'complete': '✅ Complete'
+                }
                 
-                if progress_data['active']:
+                progress_data = st.session_state.progress_state
+                if progress_data.get('active'):
                     stage_name = stage_names.get(progress_data['stage'], progress_data['stage'])
                     stage_text.text(f"**{stage_name}**")
                     status_text.text(progress_data['message'])
                     progress_bar.progress(progress_data['progress'])
                 else:
                     status_text.text('⏳ Starting...')
-            
-            enhanced_img = result_container.get('enhanced_img')
-            process_duration = time.time() - start_time
-            
-            # Clear progress UI
-            progress_bar.empty()
-            status_text.empty()
-            stage_text.empty()
-            
-            # Check for errors
-            if 'error' in result_container:
-                st.error(f"❌ Processing failed after retries: {result_container['error']}")
-                st.info("💡 Try reducing the upscale factor or disabling some features.")
+                
+                # Render Cancel button
+                cancel_col, _ = st.columns([1, 3])
+                with cancel_col:
+                    cancel_button = st.button("❌ Cancel", type="secondary")
+                
+                if cancel_button:
+                    if pipeline:
+                        pipeline.cancel_flag = True
+                    st.session_state.processing = False
+                    st.session_state.progress_state['cancelled'] = True
+                    st.warning("⚠️ Processing cancelled by user")
+                    st.rerun()
+                
+                time.sleep(0.1)
+                st.rerun()
                 st.stop()
-            
-            # Save to history
+
+        if st.session_state.get('processing_error') is not None:
+            st.error(f"❌ Processing failed: {st.session_state.processing_error}")
+            st.info("💡 Try reducing the upscale factor or disabling some features.")
+            if st.button("🔄 Try Again"):
+                st.session_state.processing_error = None
+                st.session_state.processing = False
+                st.rerun()
+            st.stop()
+
+        # If we reach here, we have the enhanced image
+        enhanced_img = st.session_state.enhanced_img
+        process_duration = st.session_state.process_duration
+
+        # Save to history
+        if st.session_state.get('history_added_for') != current_params:
             history_item = {
                 'name': img_name,
                 'timestamp': time.time(),
@@ -725,12 +747,9 @@ with tab_single:
                 'enhanced_shape': enhanced_img.shape[:2]
             }
             st.session_state.history.insert(0, history_item)
-            # Keep only last 10 items
             if len(st.session_state.history) > 10:
                 st.session_state.history = st.session_state.history[:10]
-        else:
-            st.warning('Processing is already running. Please wait.')
-            st.stop()
+            st.session_state.history_added_for = current_params
 
         h_orig, w_orig = img.shape[:2]
         h_enh,  w_enh  = enhanced_img.shape[:2]
@@ -853,6 +872,9 @@ with tab_batch:
     )
     if uploaded_files:
         if st.button("🚀 Process Batch", key="batch_button"):
+            if pipeline is None:
+                st.error("Cannot process batch: The AI Enhancer pipeline is offline or failed to initialize. Please check the logs.")
+                st.stop()
             import zipfile
             from io import BytesIO
             

datasets/crop_and_align_facexlib.py

@@ -16,9 +16,11 @@ from facelib.utils.face_restoration_helper import FaceRestoreHelper
 # Configuration
 # ----------------------------------------------------------------------
 SRC_ROOT = os.path.join(project_dir, "datasets", "game_characters")
-DST_ROOT = os.path.join(project_dir, "datasets", "ffhq", "ffhq_512")
+DST_ROOT = os.path.join(project_dir, "models", "CodeFormer", "datasets", "ffhq", "ffhq_512")
 DETECTOR = "retinaface_mobile0.25"  # Fast CPU detector
 BLUR_THRESHOLD = 80.0  # Discard crops with Laplacian variance below this
+MIN_RESOLUTION = 128  # Discard raw images smaller than this
+MIN_FACE_COVERAGE = 0.15  # Discard face crops that cover less than 15% of the original image area
 
 def is_blurry(img, threshold=BLUR_THRESHOLD):
     """Check if an image is blurry using the variance of Laplacian method."""
@@ -33,9 +35,7 @@ def main():
     # Configure face detector path
     os.environ['FACE_DETECTOR_PATH'] = os.path.join(project_dir, "weights", "facelib")
     
-    # Clear and recreate destination folder
-    if os.path.exists(DST_ROOT):
-        shutil.rmtree(DST_ROOT)
+    # Ensure destination folder exists
     os.makedirs(DST_ROOT, exist_ok=True)
     
     # Initialize face helper
@@ -61,6 +61,8 @@ def main():
     count = 0
     failures = 0
     skipped_blur = 0
+    skipped_res = 0
+    skipped_coverage = 0
     
     for i, img_path in enumerate(raw_images):
         try:
@@ -69,6 +71,12 @@ def main():
                 failures += 1
                 continue
                 
+            # Resolution check
+            h_img, w_img, _ = img.shape
+            if h_img < MIN_RESOLUTION or w_img < MIN_RESOLUTION:
+                skipped_res += 1
+                continue
+                
             face_helper.clean_all()
             face_helper.read_image(img)
             
@@ -87,18 +95,29 @@ def main():
             face_helper.align_warp_face()
             
             # Save cropped faces
-            for cropped_face in face_helper.cropped_faces:
+            for idx, cropped_face in enumerate(face_helper.cropped_faces):
+                # Face coverage check
+                if hasattr(face_helper, 'face_det_list') and idx < len(face_helper.face_det_list):
+                    det_box = face_helper.face_det_list[idx]
+                    x1, y1, x2, y2 = det_box[:4]
+                    face_area = (x2 - x1) * (y2 - y1)
+                    img_area = w_img * h_img
+                    ratio = face_area / img_area
+                    if ratio < MIN_FACE_COVERAGE:
+                        skipped_coverage += 1
+                        continue
+
                 # Blur filter check
                 if is_blurry(cropped_face):
                     skipped_blur += 1
                     continue
                     
                 count += 1
-                dest_path = os.path.join(DST_ROOT, f"face_{count:06d}.png")
+                dest_path = os.path.join(DST_ROOT, f"face_align_{count:06d}.png")
                 cv2.imwrite(dest_path, cropped_face)
                 
             if (i + 1) % 100 == 0:
-                print(f"  Processed {i + 1}/{len(raw_images)} images. Extracted {count} faces (skipped {skipped_blur} blurry)...")
+                print(f"  Processed {i + 1}/{len(raw_images)} images. Extracted {count} faces (skipped {skipped_blur} blurry, {skipped_coverage} small, {skipped_res} low-res)...")
                 
         except Exception as e:
             # print(f"Error processing {img_path}: {e}")
@@ -108,6 +127,8 @@ def main():
     print(f"  Total raw images: {len(raw_images)}")
     print(f"  Extracted face crops: {count}")
     print(f"  Skipped blurry crops: {skipped_blur}")
+    print(f"  Skipped low-res images: {skipped_res}")
+    print(f"  Skipped small face crops: {skipped_coverage}")
     print(f"  Failed loads: {failures}")
     print(f"  Destination: {DST_ROOT}\n")
 

datasets/resize_and_clean.py

@@ -10,7 +10,7 @@ import os, cv2, glob
 # Source folder containing the raw character folders
 SRC_ROOT = os.path.abspath(os.path.join(os.path.dirname(__file__), "game_characters"))
 # Destination folder expected by CodeFormer
-DST_ROOT = os.path.abspath(os.path.join(os.path.dirname(__file__), "..", "ffhq", "ffhq_512"))
+DST_ROOT = os.path.abspath(os.path.join(os.path.dirname(__file__), "..", "models", "CodeFormer", "datasets", "ffhq", "ffhq_512"))
 
 os.makedirs(DST_ROOT, exist_ok=True)
 

download_starrail.py

@@ -355,7 +355,22 @@ def main():
     print("#" * 60)
     # A. Official Skins (113 characters)
     download_official_skins("bluearchive", bluearchive_dir)
-    # B. Safebooru Mix
+    # B. Safebooru Characters
+    ba_chars = {
+        "shiroko": "shiroko_(blue_archive)",
+        "aru": "aru_(blue_archive)",
+        "hoshino": "hoshino_(blue_archive)",
+        "nonomi": "nonomi_(blue_archive)",
+        "serika": "serika_(blue_archive)",
+        "momoi": "momoi_(blue_archive)",
+        "midori": "midori_(blue_archive)",
+        "yuzu": "yuzu_(blue_archive)",
+        "mutsuki": "mutsuki_(blue_archive)",
+        "iori": "iori_(blue_archive)"
+    }
+    for key, tag in ba_chars.items():
+        download_from_safebooru(tag, bluearchive_dir, f"ba_{key}", 30)
+    # C. Safebooru Mix
     download_from_safebooru("blue_archive", bluearchive_dir, "ba_general", 300)
 
     # ─────────────────────────────────────────────────────────────

pipeline.py

@@ -38,6 +38,7 @@ class LocalAIEnhancerPipeline:
             self.device = torch.device(device)
         
         self.progress_callback = progress_callback
+        self.cancel_flag = False
             
         print(f"[Pipeline] Initializing pipeline on device: {self.device}")
         
@@ -104,12 +105,7 @@ class LocalAIEnhancerPipeline:
     
     def _check_cancelled(self):
         """Check if processing was cancelled by user."""
-        if self.progress_callback:
-            # Check session state through callback
-            import streamlit as st
-            if hasattr(st, 'session_state') and 'progress_state' in st.session_state:
-                return st.session_state.progress_state.get('cancelled', False)
-        return False
+        return self.cancel_flag
     
     def _get_onnx_session(self, path, providers=None):
         """Get or create cached ONNX session."""
@@ -154,7 +150,7 @@ class LocalAIEnhancerPipeline:
 
     def run_onnx_batch(self, faces_np, w_val):
         """Helper to run ONNX batch inference."""
-        w_np = np.full((faces_np.shape[0], 1), w_val, dtype=np.float32)
+        w_np = np.full((faces_np.shape[0],), w_val, dtype=np.float32)
         ort_inputs = {
             self.ort_session_cf.get_inputs()[0].name: faces_np,
             self.ort_session_cf.get_inputs()[1].name: w_np

requirements.txt

@@ -116,3 +116,4 @@ websockets==16.0
 Werkzeug==3.1.8
 wheel==0.47.0
 yapf==0.43.0
+streamlit-image-comparison

tools/download_ffhq_extended.py

@@ -0,0 +1,123 @@
+import os
+import sys
+import io
+import pandas as pd
+from PIL import Image
+import requests
+import time
+
+SHARDS = [
+    "train-00001-of-00054-c76023902046cea3.parquet",
+    "train-00002-of-00054-b43f6d454561b047.parquet",
+    "train-00003-of-00054-0637d6a1c5f28946.parquet",
+    "train-00004-of-00054-17a50ec36be02fa9.parquet",
+    "train-00005-of-00054-b370d8ca7905b127.parquet",
+    "train-00006-of-00054-e7f70526908ad428.parquet",
+    "train-00007-of-00054-36469d3079484e03.parquet",
+    "train-00008-of-00054-564af51540d658a8.parquet",
+    "train-00009-of-00054-57224ef56d82da23.parquet"
+]
+
+def load_parquet_safe(url):
+    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):
+    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):
+    if isinstance(cell, dict):
+        return cell.get("bytes")
+    if isinstance(cell, bytes):
+        return cell
+    return None
+
+def main():
+    project_dir = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
+    codeformer_dir = os.path.join(project_dir, "models", "CodeFormer")
+    save_dir = os.path.join(codeformer_dir, "datasets", "ffhq", "ffhq_512", "faces")
+    os.makedirs(save_dir, exist_ok=True)
+    
+    print("=" * 55)
+    print("  EXTENDED FFHQ REAL FACE DOWNLOADER")
+    print(f"  Target directory: {save_dir}")
+    print("=" * 55)
+    
+    # Check current count in faces subdirectory to determine start_idx
+    existing = [f for f in os.listdir(save_dir) if f.endswith('.png') and f.startswith('face')]
+    highest_idx = -1
+    for f in existing:
+        try:
+            idx = int(f.replace('face_', '').replace('.png', ''))
+            highest_idx = max(highest_idx, idx)
+        except ValueError:
+            continue
+    start_idx = highest_idx + 1
+    print(f"Starting index for new files: face_{start_idx:05d}.png")
+    
+    images_per_shard = 300
+    total_saved = 0
+    
+    for shard_idx, shard_name in enumerate(SHARDS):
+        print(f"\n[{shard_idx + 1}/{len(SHARDS)}] Processing {shard_name}...")
+        url = f"https://huggingface.co/datasets/Ryan-sjtu/ffhq512-caption/resolve/main/data/{shard_name}"
+        
+        df = load_parquet_safe(url)
+        if df is None:
+            continue
+            
+        img_col = find_image_column(df)
+        if img_col is None:
+            print("  No image column found in dataframe.")
+            continue
+            
+        saved_in_shard = 0
+        for i in range(len(df)):
+            if saved_in_shard >= images_per_shard:
+                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"face_{start_idx + total_saved:05d}.png"
+                filepath = os.path.join(save_dir, filename)
+                img.save(filepath, "PNG")
+                
+                saved_in_shard += 1
+                total_saved += 1
+                
+                if saved_in_shard % 50 == 0 or saved_in_shard == images_per_shard:
+                    print(f"  Saved {saved_in_shard}/{images_per_shard} from this shard...")
+                    
+            except Exception as e:
+                print(f"  [warn] Row {i} error: {e}")
+                continue
+                
+        print(f"  Shard completed. Saved {saved_in_shard} images.")
+        time.sleep(0.5)
+        
+    print("\n" + "=" * 55)
+    print("  DOWNLOAD COMPLETED")
+    print(f"  Total downloaded this run: {total_saved}")
+    print(f"  Total images in directory: {len(os.listdir(save_dir))}")
+    print("=" * 55)
+
+if __name__ == "__main__":
+    main()

tools/download_highres_datasets.py

@@ -0,0 +1,167 @@
+import os
+import sys
+import argparse
+import urllib.request
+import zipfile
+import shutil
+from PIL import Image
+
+# Setup sys.path to bypass local "datasets" folder conflict
+PROJECT_DIR = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
+DATASET_BASE = os.path.join(PROJECT_DIR, "models", "CodeFormer", "datasets", "ffhq", "ffhq_512")
+
+# Setup clean path for HuggingFace datasets
+sys.path = [p for p in sys.path if p not in ('', '.')]
+from datasets import load_dataset
+
+def download_div2k(test_run=False):
+    print("\n" + "="*50)
+    print("  DOWNLOADING DIV2K DATASET (2K RESOLUTION)")
+    print("="*50)
+    
+    target_dir = os.path.join(DATASET_BASE, "div2k")
+    os.makedirs(target_dir, exist_ok=True)
+    
+    # Check if we already have images
+    existing = len([f for f in os.listdir(target_dir) if f.endswith(".png")])
+    if existing >= (5 if test_run else 800):
+        print(f"[OK] DIV2K already contains {existing} images. Skipping download.")
+        return
+        
+    zip_url = "http://data.vision.ee.ethz.ch/cvl/DIV2K/DIV2K_train_HR.zip"
+    zip_path = os.path.join(PROJECT_DIR, "DIV2K_train_HR.zip")
+    extract_dir = os.path.join(PROJECT_DIR, "DIV2K_temp")
+    
+    # In test mode we don't download 3.5GB zip, we just download 5 images directly if possible
+    if test_run:
+        print("[Test Mode] Downloading 5 sample images directly from DIV2K repo...")
+        for i in range(1, 6):
+            img_url = f"https://raw.githubusercontent.com/eugenesiow/super-image-data/master/div2k/DIV2K_train_HR/{i:04d}.png"
+            dest = os.path.join(target_dir, f"div2k_{i:04d}.png")
+            try:
+                urllib.request.urlretrieve(img_url, dest)
+                print(f"  Saved sample {i}/5 -> {dest}")
+            except Exception as e:
+                print(f"  Failed sample {i}: {e}")
+        return
+
+    print(f"Downloading {zip_url} (approx. 3.5 GB)...")
+    def report_hook(block_num, block_size, total_size):
+        read_so_far = block_num * block_size
+        if total_size > 0:
+            percent = min(100, read_so_far * 100 / total_size)
+            sys.stdout.write(f"\rProgress: {percent:.2f}% ({read_so_far / 1024 / 1024:.2f} MB of {total_size / 1024 / 1024:.2f} MB)")
+        else:
+            sys.stdout.write(f"\rProgress: {read_so_far / 1024 / 1024:.2f} MB")
+        sys.stdout.flush()
+        
+    try:
+        urllib.request.urlretrieve(zip_url, zip_path, reporthook=report_hook)
+        print("\n[OK] DIV2K Zip downloaded. Extracting...")
+        
+        os.makedirs(extract_dir, exist_ok=True)
+        with zipfile.ZipFile(zip_path, 'r') as zip_ref:
+            zip_ref.extractall(extract_dir)
+            
+        print("[*] Copying images to target directory...")
+        copied = 0
+        for root, _, files in os.walk(extract_dir):
+            for file in files:
+                if file.lower().endswith(".png"):
+                    src = os.path.join(root, file)
+                    dest = os.path.join(target_dir, f"div2k_{file}")
+                    shutil.copy2(src, dest)
+                    copied += 1
+                    
+        print(f"[OK] Successfully extracted and copied {copied} DIV2K images.")
+        
+    except Exception as e:
+        print(f"\n[FAIL] Error occurred during DIV2K setup: {e}")
+    finally:
+        # Cleanup
+        try:
+            if os.path.exists(zip_path):
+                os.remove(zip_path)
+            if os.path.exists(extract_dir):
+                shutil.rmtree(extract_dir)
+        except Exception:
+            pass
+
+def download_huggingface_dataset(repo_name, target_subdir, max_images, image_key='image', prefix='img', test_run=False):
+    print("\n" + "="*50)
+    print(f"  STREAMING DATASET FROM HF: {repo_name}")
+    print("="*50)
+    
+    target_dir = os.path.join(DATASET_BASE, target_subdir)
+    os.makedirs(target_dir, exist_ok=True)
+    
+    limit = 5 if test_run else max_images
+    
+    existing = len([f for f in os.listdir(target_dir) if f.endswith(".png")])
+    if existing >= limit:
+        print(f"[OK] {target_subdir} already contains {existing} images. Skipping.")
+        return
+        
+    print(f"Streaming from HF hub. Target count: {limit} images...")
+    try:
+        ds = load_dataset(repo_name, split="train", streaming=True)
+        saved = 0
+        for i, item in enumerate(ds):
+            if saved >= limit:
+                break
+                
+            dest = os.path.join(target_dir, f"{prefix}_{saved:06d}.png")
+            if os.path.exists(dest):
+                saved += 1
+                continue
+                
+            try:
+                # Retrieve the PIL image
+                pil_img = item[image_key]
+                if not isinstance(pil_img, Image.Image):
+                    # In case it's a dict with bytes
+                    import io
+                    pil_img = Image.open(io.BytesIO(pil_img['bytes']))
+                
+                # Keep high-res, save as PNG
+                pil_img.convert("RGB").save(dest, "PNG")
+                saved += 1
+                if saved % 100 == 0 or test_run:
+                    print(f"  -> Saved {saved}/{limit} images")
+            except Exception as e:
+                print(f"  [Error] Failed to save image index {i}: {e}")
+                
+        print(f"[OK] Successfully streamed and saved {saved} images to {target_dir}")
+    except Exception as e:
+        print(f"[FAIL] Error streaming dataset {repo_name}: {e}")
+
+def main():
+    parser = argparse.ArgumentParser(description="Download and integrate high-resolution datasets for super-resolution training.")
+    parser.add_argument("--test-run", action="store_true", help="Only download 5 images from each dataset for verification.")
+    args = parser.parse_args()
+    
+    # 1. Download DIV2K (2K resolution)
+    download_div2k(test_run=args.test_run)
+    
+    # 2. Stream Flickr2K Subset (2K resolution, yangtao9009/Flickr2K)
+    download_huggingface_dataset(
+        repo_name="yangtao9009/Flickr2K",
+        target_subdir="flickr2k_subset",
+        max_images=1000,
+        image_key="image",
+        prefix="flickr2k",
+        test_run=args.test_run
+    )
+    
+    # 3. Stream FFHQ-1024x1024 Subset (1024px resolution, Iceclear/FFHQ-HQ1024)
+    download_huggingface_dataset(
+        repo_name="Iceclear/FFHQ-HQ1024",
+        target_subdir="ffhq1024_subset",
+        max_images=10000,
+        image_key="image",
+        prefix="ffhq1024",
+        test_run=args.test_run
+    )
+
+if __name__ == "__main__":
+    main()

tools/download_weights.py

@@ -31,11 +31,11 @@ def download_file(url, save_path):
 
 def main():
     weights_to_download = {
-        "weights/CodeFormer/codeformer.pth": "https://github.com/sczhou/CodeFormer/releases/download/v0.1.0/codeformer.pth",
-        "weights/facelib/detection_Resnet50_Final.pth": "https://github.com/sczhou/CodeFormer/releases/download/v0.1.0/detection_Resnet50_Final.pth",
-        "weights/facelib/parsing_parsenet.pth": "https://github.com/sczhou/CodeFormer/releases/download/v0.1.0/parsing_parsenet.pth",
-        "weights/facelib/yolov5l-face.pth": "https://github.com/sczhou/CodeFormer/releases/download/v0.1.0/yolov5l-face.pth",
-        "weights/realesrgan/RealESRGAN_x2plus.pth": "https://github.com/sczhou/CodeFormer/releases/download/v0.1.0/RealESRGAN_x2plus.pth"
+        "weights/CodeFormer/codeformer.pth": "https://huggingface.co/spaces/sczhou/CodeFormer/resolve/main/weights/CodeFormer/codeformer.pth",
+        "weights/facelib/detection_Resnet50_Final.pth": "https://huggingface.co/spaces/sczhou/CodeFormer/resolve/main/weights/facelib/detection_Resnet50_Final.pth",
+        "weights/facelib/parsing_parsenet.pth": "https://huggingface.co/spaces/sczhou/CodeFormer/resolve/main/weights/facelib/parsing_parsenet.pth",
+        "weights/facelib/yolov5l-face.pth": "https://huggingface.co/spaces/sczhou/CodeFormer/resolve/main/weights/facelib/yolov5l-face.pth",
+        "weights/realesrgan/RealESRGAN_x2plus.pth": "https://huggingface.co/spaces/sczhou/CodeFormer/resolve/main/weights/realesrgan/RealESRGAN_x2plus.pth"
     }
     
     for save_path, url in weights_to_download.items():

tools/export_onnx.py

@@ -104,7 +104,7 @@ def export_codeformer():
         print(f"[WARNING] ONNX validation failed: {e}")
         return True
 
-def export_realesrgan():
+def export_realesrgan(num_block=None):
     print("\n--- Exporting Real-ESRGAN to ONNX ---")
     device = torch.device("cpu")
     
@@ -115,21 +115,26 @@ def export_realesrgan():
     if checkpoint_path:
         print(f"Found custom Real-ESRGAN checkpoint: {checkpoint_path}")
         scale = 4
+        default_blocks = 6
     else:
         checkpoint_path = os.path.join(project_dir, "weights", "realesrgan", "RealESRGAN_x2plus.pth")
         print(f"No custom checkpoint found. Using pretrained weights: {checkpoint_path}")
         scale = 2
+        default_blocks = 23
         
     if not os.path.exists(checkpoint_path):
         print(f"[ERROR] Real-ESRGAN weights not found at: {checkpoint_path}")
         return False
         
+    blocks = num_block if num_block is not None else default_blocks
+    print(f"Instantiating RRDBNet with scale={scale}, num_block={blocks}")
+    
     # 2. Instantiate RRDBNet architecture with correct scale
     net = RRDBNet(
         num_in_ch=3,
         num_out_ch=3,
         num_feat=64,
-        num_block=23,
+        num_block=blocks,
         num_grow_ch=32,
         scale=scale
     )
@@ -175,8 +180,13 @@ def export_realesrgan():
         return True
 
 def main():
+    import argparse
+    parser = argparse.ArgumentParser(description="Export CodeFormer and Real-ESRGAN models to ONNX")
+    parser.add_argument('--num-block', type=int, default=None, help="Number of RRDB blocks for Real-ESRGAN (defaults to 6 for custom checkpoint, 23 for pretrained)")
+    args = parser.parse_args()
+    
     success_cf = export_codeformer()
-    success_re = export_realesrgan()
+    success_re = export_realesrgan(num_block=args.num_block)
     if success_cf and success_re:
         print("\n=== ALL MODELS EXPORTED SUCCESSFULLY ===")
     else:

tools/test_pipeline.py

@@ -36,7 +36,7 @@ def main():
         enhanced_img = pipeline.process_image(
             img, 
             w=0.5, 
-            detection_model='retinaface_resnet50', 
+            detection_model='retinaface_mobile0.25', 
             upscale=upscale, 
             blend_softness=0.5
         )

train_custom.py

@@ -7,7 +7,13 @@ 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")
     
@@ -84,24 +90,31 @@ def main():
         except ValueError:
             continue
     
-    if latest_state and latest_state_iter < config.get("train", {}).get("total_iter", 50):
+    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
-        # When resuming, we don't need pretrain_network_g as it will be loaded from state
         config["path"]["pretrain_network_g"] = None
-        print(f"    Resuming training from iter {latest_state_iter} to {config['train']['total_iter']}...")
-    elif latest_state and latest_state_iter >= config.get("train", {}).get("total_iter", 50):
-        print(f"\n>>> Training already completed ({latest_state_iter} >= {config['train']['total_iter']} total_iter)")
-        print("    To train more iterations, increase total_iter in the config.")
-        sys.exit(0)
+        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)}...")
     
     # 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}")
+    print(f"Updated configuration file: device={device.upper()}, num_gpu={num_gpus}, total_iter={config['train']['total_iter']}")
     
     # 5. Run the training process
     train_script = os.path.join("basicsr", "train.py")

train_realesrgan.py

@@ -24,6 +24,9 @@ ANIME_DIR = os.path.join(DATASET_BASE, "anime")
 FACE_DIR = os.path.join(DATASET_BASE, "faces")
 STARRAIL_DIR = os.path.join(DATASET_BASE, "starrail")
 BLUEARCHIVE_DIR = os.path.join(DATASET_BASE, "bluearchive")
+DIV2K_DIR = os.path.join(DATASET_BASE, "div2k")
+FLICKR2K_DIR = os.path.join(DATASET_BASE, "flickr2k_subset")
+FFHQ1024_DIR = os.path.join(DATASET_BASE, "ffhq1024_subset")
 
 # Combined GT (high-quality) folder for Real-ESRGAN training
 REALESRGAN_GT_DIR = os.path.join(PROJECT_DIR, "datasets", "realesrgan_gt")
@@ -76,6 +79,9 @@ def prepare_gt_dataset():
         "face": FACE_DIR,
         "starrail": STARRAIL_DIR,
         "bluearchive": BLUEARCHIVE_DIR,
+        "div2k": DIV2K_DIR,
+        "flickr2k": FLICKR2K_DIR,
+        "ffhq1024": FFHQ1024_DIR,
     }
 
     for category, src_dir in sources.items():
@@ -300,6 +306,11 @@ def find_latest_state() -> str | None:
 
 
 def main():
+    import argparse
+    parser = argparse.ArgumentParser(description="Train Real-ESRGAN with custom parameters.")
+    parser.add_argument("--verify", action="store_true", help="Run 2 iterations for verification purposes.")
+    args = parser.parse_args()
+
     print("=" * 55)
     print("  Real-ESRGAN Custom Training Runner")
     print("  Target: Landscape + Anime + Face enhancement")
@@ -347,10 +358,27 @@ def main():
             cfg = yaml.safe_load(f)
         cfg["path"]["resume_state"] = latest_state
         cfg["path"]["pretrain_network_g"] = None
+        if args.verify:
+            cfg["train"]["total_iter"] = iter_num + 2
+            print(f"    Resuming training from iter {iter_num} to {cfg['train']['total_iter']} (verification mode)...")
+        else:
+            if iter_num >= cfg.get("train", {}).get("total_iter", 15000):
+                print(f"\n>>> Training already completed ({iter_num} >= {cfg.get('train', {}).get('total_iter', 15000)} total_iter)")
+                sys.exit(0)
+            print(f"    Resuming training from iter {iter_num} to {cfg['train']['total_iter']}...")
         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.")
+        with open(config_path, "r", encoding="utf-8") as f:
+            cfg = yaml.safe_load(f)
+        if args.verify:
+            cfg["train"]["total_iter"] = 2
+            print(f"    Training from scratch to {cfg['train']['total_iter']} (verification mode)...")
+        else:
+            print(f"    Training from scratch to {cfg.get('train', {}).get('total_iter', 15000)}...")
+        with open(config_path, "w", encoding="utf-8") as f:
+            yaml.dump(cfg, f, default_flow_style=False, sort_keys=False)
 
     # 6. Run training
     train_script = os.path.join("realesrgan", "train.py")
@@ -384,12 +412,8 @@ def main():
     # Real-ESRGAN's train.py imports from basicsr.
     # CodeFormer's basicsr lacks degradations module, so we use the full
     # BasicSR source cloned at D:\Temp\BasicSR_src (no install needed).
-    basicsr_src = r"D:\Temp\BasicSR_src"
-    codeformer_dir = os.path.join(PROJECT_DIR, "models", "CodeFormer")
     env["PYTHONPATH"] = os.path.pathsep.join([
         REALESRGAN_DIR,
-        basicsr_src,       # full BasicSR with degradations
-        codeformer_dir,    # fallback
         env.get("PYTHONPATH", ""),
     ])