Upload 2 files

app.py

@@ -21,9 +21,13 @@ from diffusers import FluxInpaintPipeline
 from transformers import BlipProcessor, BlipForConditionalGeneration
 from spandrel import ModelLoader, ImageModelDescriptor
 from huggingface_hub import hf_hub_download
-import onnxruntime as ort
+import torch
+from ddcolor import DDColor
+from huggingface_hub import PyTorchModelHubMixin
 from skimage.metrics import peak_signal_noise_ratio as psnr
 from skimage.metrics import structural_similarity as ssim
+import sys
+sys.path.insert(0, os.path.join(os.path.dirname(__file__), "DDColor"))
 
 # ================================================================
 # 5. LOAD ALL MODELS AT STARTUP (no lazy loading)
@@ -64,10 +68,14 @@ esrgan_model.to("cuda").eval()
 ESRGAN_SCALE = int(esrgan_model.scale) 
 print("Enhancement model loaded. Scale factor:", ESRGAN_SCALE)
 
-print("🚀 Loading DDColor ONNX...")
-model_path = hf_hub_download("edgetools/ddcolor", "ddcolor-tiny-fp16.onnx")
-cuda_options = {"cudnn_conv_algo_search": "HEURISTIC"}
-session = ort.InferenceSession(model_path, providers=['CPUExecutionProvider'])   
+class DDColorHF(DDColor, PyTorchModelHubMixin):
+    def __init__(self, config=None, **kwargs):
+        if isinstance(config, dict):
+            kwargs = {**config, **kwargs}
+        super().__init__(**kwargs)
+
+print("🚀 Loading DDColor (PyTorch tiny)...")
+ddcolor_model = DDColorHF.from_pretrained("piddnad/ddcolor_paper_tiny").eval()
 print("✅ DDColor loaded.")
 
 print("✅ All models loaded. Starting Gradio interface...")
@@ -320,32 +328,32 @@ def enhance_image(image, scale_factor, benchmark_mode=False):
 
     return output_image, info
 
-# ================================================================
-# 9. COLORIZATION (uses pre-loaded ddcolor_pipe)
-# ================================================================
 def colorize_ddcolor(image_rgb: np.ndarray) -> np.ndarray:
+    """Uses the PyTorch ddcolor_model (already on .cuda())."""
     h, w = image_rgb.shape[:2]
     img_float = image_rgb.astype(np.float32) / 255.0
     l_orig = cv2.cvtColor(img_float, cv2.COLOR_RGB2Lab)[:, :, :1]
 
     l_resized = cv2.resize(l_orig, (512, 512), interpolation=cv2.INTER_AREA)
-    l_resized = l_resized.reshape(512, 512, 1)  # ensure 3D
-
+    l_resized = l_resized.reshape(512, 512, 1)
     gray_lab = np.concatenate([l_resized, np.zeros_like(l_resized), np.zeros_like(l_resized)], axis=-1)
     gray_rgb = cv2.cvtColor(gray_lab, cv2.COLOR_LAB2RGB)
-    input_tensor = gray_rgb.transpose(2, 0, 1)[np.newaxis, ...]
 
-    output = session.run(None, {session.get_inputs()[0].name: input_tensor})[0]
-    print("ab range:", output.min(), output.max())
+    # PyTorch tensor (not numpy for ONNX)
+    input_tensor = torch.from_numpy(gray_rgb.transpose(2, 0, 1)).float().unsqueeze(0).to("cuda")
 
-    a = cv2.resize(output[0, 0], (w, h), interpolation=cv2.INTER_AREA)
-    b = cv2.resize(output[0, 1], (w, h), interpolation=cv2.INTER_AREA)
-    a = a.reshape(h, w, 1)  # same fix
-    b = b.reshape(h, w, 1)
+    with torch.no_grad():
+        output_ab = ddcolor_model(input_tensor)  # (1, 2, 512, 512) in [0,1]
+
+    output_ab = output_ab.cpu().numpy()[0].transpose(1, 2, 0)  # (512, 512, 2)
+
+    a = cv2.resize(output_ab[:, :, 0:1], (w, h), interpolation=cv2.INTER_AREA).reshape(h, w, 1)
+    b = cv2.resize(output_ab[:, :, 1:2], (w, h), interpolation=cv2.INTER_AREA).reshape(h, w, 1)
 
     lab_out = np.concatenate([l_orig, a, b], axis=-1)
     out_rgb = cv2.cvtColor(lab_out, cv2.COLOR_LAB2RGB)
-    return (out_rgb * 255).clip(0, 255).astype(np.uint8)       
+    return (out_rgb * 255).clip(0, 255).astype(np.uint8)
+
 
 @spaces.GPU(duration=58)
 def colour_image(image, strength):
@@ -356,11 +364,13 @@ def colour_image(image, strength):
     small = rgb.copy()
     small.thumbnail((768, 768), Image.LANCZOS)
 
-    # 1. Colorize
-    coloured = colorize_ddcolor(np.asarray(small))  # (h, w, 3) uint8 RGB
-    coloured_full = Image.fromarray(coloured).resize((W, H), Image.LANCZOS)
+    # GPU: inference
+    ddcolor_model.cuda()
+    coloured = colorize_ddcolor(np.asarray(small))
+    ddcolor_model.cpu()
 
-    # 2. LAB strength
+    # CPU: LAB strength (no GPU needed)
+    coloured_full = Image.fromarray(coloured).resize((W, H), Image.LANCZOS)
     lab_c = cv2.cvtColor(np.asarray(coloured_full), cv2.COLOR_RGB2LAB)
     l_orig = cv2.cvtColor(np.asarray(rgb), cv2.COLOR_RGB2LAB)[:, :, 0]
     a = lab_c[:, :, 1].astype(np.float32)
@@ -370,7 +380,7 @@ def colour_image(image, strength):
     a = cv2.bilateralFilter(a, 9, 40, 9)
     b = cv2.bilateralFilter(b, 9, 40, 9)
     out = cv2.cvtColor(cv2.merge([l_orig, a, b]), cv2.COLOR_LAB2RGB)
-    return Image.fromarray(out)   
+    return Image.fromarray(out)    
 
 # ================================================================
 # 10. TAB HELPERS

requirements.txt

@@ -1,5 +1,4 @@
 setuptools<82
-onnxruntime
 spandrel
 safetensors
 huggingface_hub
@@ -13,4 +12,6 @@ opencv-python-headless
 addict
 simplejson
 sortedcontainers
-scikit-image>=0.22   
+scikit-image>=0.22   
+timm
+einops