fix(hf): integrate Static INT8 ONNX build into Dockerfile and fix thread state guard in app.py

Dockerfile

@@ -33,8 +33,10 @@ COPY . /app/
 # Create weights directory and make sure it has permissions
 RUN mkdir -p /app/weights && chmod -R 777 /app/weights
 
-# Download model weights during build phase to prevent runtime startup hang
+# Download model weights, export ONNX models, and generate static INT8 quantization during build phase
 RUN python tools/download_weights.py
+RUN python tools/export_onnx.py
+RUN python tools/quantize_onnx_static.py
 
 # Expose port 7860 for Hugging Face Spaces
 EXPOSE 7860

models/CodeFormer/basicsr/models/codeformer_joint_model.py

@@ -105,11 +105,12 @@ class CodeFormerJointModel(SRModel):
         self.setup_schedulers()
 
     def calculate_adaptive_weight(self, recon_loss, g_loss, last_layer, disc_weight_max):
-        recon_grads = torch.autograd.grad(recon_loss, last_layer, retain_graph=True)[0]
-        g_grads = torch.autograd.grad(g_loss, last_layer, retain_graph=True)[0]
+        recon_grads = torch.autograd.grad(recon_loss, last_layer, retain_graph=True)[0].detach()
+        g_grads = torch.autograd.grad(g_loss, last_layer, retain_graph=True)[0].detach()
 
         d_weight = torch.norm(recon_grads) / (torch.norm(g_grads) + 1e-4)
         d_weight = torch.clamp(d_weight, 0.0, disc_weight_max).detach()
+        del recon_grads, g_grads
         return d_weight
 
     def setup_optimizers(self):

models/CodeFormer/basicsr/train.py

@@ -203,13 +203,6 @@ def train_pipeline(root_path):
             if current_iter % opt['logger']['save_checkpoint_freq'] == 0:
                 logger.info('Saving models and training states.')
                 model.save(epoch, current_iter)
-                import gc
-                gc.collect()
-
-            # Garbage collect periodically on CPU to prevent RAM memory leaks
-            if current_iter % 50 == 0:
-                import gc
-                gc.collect()
 
             # validation
             if opt.get('val') is not None and opt['datasets'].get('val') is not None \

pipeline.py

@@ -404,10 +404,10 @@ class LocalAIEnhancerPipeline:
             
             self._report_progress("restoration", 0.8, "Face restoration complete")
         else:
-        # B2 FIX: Only use ThreadPoolExecutor when there are multiple faces.
-        # For single-face images (the common case), spawning a thread pool
-        # adds ~20 ms of overhead with zero parallelism benefit.
-        if parallel and len(face_helper.cropped_faces) > 1:
+            # B2 FIX: Only use ThreadPoolExecutor when there are multiple faces.
+            # For single-face images (the common case), spawning a thread pool
+            # adds ~20 ms of overhead with zero parallelism benefit.
+            if parallel and len(face_helper.cropped_faces) > 1:
                 def _process_face(idx, cropped_face):
                     if self.use_onnx:
                         try:

tools/benchmark_quant.py

@@ -0,0 +1,103 @@
+import os
+import sys
+import time
+import cv2
+import numpy as np
+import onnxruntime as ort
+
+project_dir = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
+weights_dir = os.path.join(project_dir, "weights", "CodeFormer")
+calib_dir = os.path.join(project_dir, "models", "CodeFormer", "datasets", "ffhq", "ffhq_512")
+
+def calculate_psnr(img1, img2):
+    mse = np.mean((img1.astype(np.float64) - img2.astype(np.float64)) ** 2)
+    if mse == 0:
+        return float('inf')
+    return 20 * np.log10(255.0 / np.sqrt(mse))
+
+import gc
+
+def benchmark_model(model_name, model_path, test_images, ort_opts):
+    if not os.path.exists(model_path):
+        print(f"[{model_name}] Model file not found: {model_path}")
+        return None
+        
+    session = ort.InferenceSession(model_path, sess_options=ort_opts, providers=['CPUExecutionProvider'])
+    input_names = [inp.name for inp in session.get_inputs()]
+    
+    times = []
+    outputs = []
+    
+    # Warmup
+    dummy_input = np.zeros((1, 3, 512, 512), dtype=np.float32)
+    w = np.array([0.5], dtype=np.float32)
+    session.run(None, {input_names[0]: dummy_input, input_names[1]: w})
+    
+    for img_path in test_images:
+        img = cv2.imread(img_path)
+        if img is None:
+            continue
+        img = cv2.resize(img, (512, 512), interpolation=cv2.INTER_LANCZOS4)
+        img_rgb = cv2.cvtColor(img, cv2.COLOR_BGR2RGB).astype(np.float32) / 255.0
+        img_norm = (img_rgb - 0.5) / 0.5
+        tensor = np.transpose(img_norm, (2, 0, 1))[np.newaxis, ...].astype(np.float32)
+        
+        t0 = time.time()
+        out = session.run(None, {input_names[0]: tensor, input_names[1]: w})[0]
+        t1 = time.time()
+        times.append(t1 - t0)
+        outputs.append(out)
+        
+    del session
+    gc.collect()
+    
+    avg_time = np.mean(times) if times else 0
+    return avg_time, outputs
+
+def main():
+    print("=== ONNX Quantization Benchmark (P5.4) ===")
+    opts = ort.SessionOptions()
+    opts.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL
+    
+    test_files = [os.path.join(calib_dir, f) for f in os.listdir(calib_dir)[:10] if f.endswith('.png')]
+    print(f"Benchmarking on {len(test_files)} sample face crops...")
+    
+    models = {
+        "FP32 Baseline": os.path.join(weights_dir, "codeformer.onnx"),
+        "Dynamic INT8": os.path.join(weights_dir, "codeformer_int8.onnx"),
+        "Static INT8 (v2)": os.path.join(weights_dir, "codeformer_int8_v2.onnx")
+    }
+    
+    results = {}
+    baseline_outputs = None
+    
+    for name, path in models.items():
+        res = benchmark_model(name, path, test_files, opts)
+        if res is not None:
+            avg_time, outputs = res
+            results[name] = {"time": avg_time, "outputs": outputs}
+            if name == "FP32 Baseline":
+                baseline_outputs = outputs
+                
+    print("\n-----------------------------------------------------------")
+    print(f"{'Model Variant':<20} | {'Avg Latency (ms)':<18} | {'PSNR vs FP32 (dB)':<18}")
+    print("-----------------------------------------------------------")
+    
+    for name, data in results.items():
+        avg_ms = data["time"] * 1000
+        if name == "FP32 Baseline":
+            psnr_str = "N/A (Baseline)"
+        else:
+            psnrs = []
+            for b_out, q_out in zip(baseline_outputs, data["outputs"]):
+                b_img = ((np.squeeze(b_out).transpose(1, 2, 0) + 1.0) * 127.5).clip(0, 255).astype(np.uint8)
+                q_img = ((np.squeeze(q_out).transpose(1, 2, 0) + 1.0) * 127.5).clip(0, 255).astype(np.uint8)
+                psnrs.append(calculate_psnr(b_img, q_img))
+            mean_psnr = np.mean(psnrs)
+            psnr_str = f"{mean_psnr:.2f} dB"
+            
+        print(f"{name:<20} | {avg_ms:<18.2f} | {psnr_str:<18}")
+    print("-----------------------------------------------------------\n")
+
+if __name__ == "__main__":
+    main()

tools/download_weights.py

@@ -39,6 +39,7 @@ def main():
         "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/facelib/yolov5n-face.pth": "https://github.com/sczhou/CodeFormer/releases/download/v0.1.0/yolov5n-face.pth",
+        "weights/facelib/arcface_resnet18.pth": "https://github.com/sczhou/CodeFormer/releases/download/v0.1.0/arcface_resnet18.pth",
         "weights/realesrgan/RealESRGAN_x2plus.pth": "https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.1/RealESRGAN_x2plus.pth"
     }
     

tools/quantize_onnx_static.py

@@ -9,7 +9,7 @@ from onnxruntime.quantization import quantize_static, CalibrationDataReader, Qua
 
 class CodeFormerCalibrationDataReader(CalibrationDataReader):
     """Calibration Data Reader for ONNX Runtime static quantization of CodeFormer."""
-    def __init__(self, calibration_folder: str, w_val: float = 0.5, max_samples: int = 100):
+    def __init__(self, calibration_folder: str, w_val: float = 0.5, max_samples: int = 30):
         super().__init__()
         valid_exts = ('.png', '.jpg', '.jpeg', '.webp')
         self.image_paths = [
@@ -41,6 +41,16 @@ class CodeFormerCalibrationDataReader(CalibrationDataReader):
             })
             count += 1
 
+        if len(self.enum_data_dicts) == 0:
+            print("[Calib] No calibration images found. Generating 10 synthetic calibration samples...")
+            np.random.seed(42)
+            for _ in range(10):
+                dummy_tensor = np.random.uniform(-1.0, 1.0, (1, 3, 512, 512)).astype(np.float32)
+                self.enum_data_dicts.append({
+                    "input": dummy_tensor,
+                    "w": self.w_val
+                })
+
         print(f"[Calib] Prepared {len(self.enum_data_dicts)} calibration samples.")
         self.enum_data = iter(self.enum_data_dicts)
 
@@ -59,14 +69,26 @@ def quantize_static_model(model_path: str, output_path: str, calibration_folder:
         print(f"[ERROR] Calibration folder is empty or invalid: {calibration_folder}")
         sys.exit(1)
 
-    print(f"[Static Quant] Quantizing {model_path} -> {output_path}")
+    # Pre-compute shape inference to file to avoid loading entire 370MB external data into RAM
+    inferred_model_path = model_path.replace(".onnx", "_inferred.onnx")
+    print(f"[Static Quant] Inferring shapes to {inferred_model_path}...")
+    try:
+        onnx.shape_inference.infer_shapes_path(model_path, inferred_model_path)
+        input_for_quant = inferred_model_path
+    except Exception as e:
+        print(f"[Static Quant] Warning during shape inference: {e}. Falling back to input model.")
+        input_for_quant = model_path
+
+    print(f"[Static Quant] Quantizing {input_for_quant} -> {output_path}")
     quantize_static(
-        model_input=model_path,
+        model_input=input_for_quant,
         model_output=output_path,
         calibration_data_reader=dr,
         quant_format=QuantType.QInt8,
         activation_type=QuantType.QUInt8,
-        weight_type=QuantType.QInt8
+        weight_type=QuantType.QInt8,
+        use_external_data_format=True,
+        extra_options={"DisableShapeInference": True}
     )
     print(f"[Static Quant] Completed successfully: {output_path}")