Spaces:
Runtime error
Runtime error
| # ============================================================================ | |
| # DeOldify WebApp — Model definitions & image processing pipeline | |
| # ---------------------------------------------------------------------------- | |
| # Author : Tariq Jamil | |
| # Version: 1.0.0 | |
| # Date : 2025-06-14 | |
| # License: MIT | |
| # ============================================================================ | |
| import os | |
| import warnings | |
| from pathlib import Path | |
| import cv2 | |
| import numpy as np | |
| import torch | |
| import torch.nn as nn | |
| import requests | |
| from PIL import Image, ImageEnhance, ImageOps, ImageFilter as PILFilter | |
| torch.set_num_threads(1) | |
| warnings.filterwarnings("ignore", ".*torch.distributed.*") | |
| _HF_DATA = Path("/data") | |
| MODEL_DIR = _HF_DATA / "models" if _HF_DATA.is_dir() else Path(__file__).parent / "models" | |
| MODEL_DIR.mkdir(parents=True, exist_ok=True) | |
| # --------------------------------------------------------------------------- | |
| # Remote model weight URLs (mirrors — HF Hub first, GitHub fallback) | |
| # --------------------------------------------------------------------------- | |
| MODEL_MIRRORS = { | |
| "ColorizeArtistic_gen.pth": [ | |
| "https://huggingface.co/databuzzword/deoldify-artistic/resolve/main/ColorizeArtistic_gen.pth", | |
| "https://github.com/databuzzword/deoldify-artistic/releases/download/v1.0/ColorizeArtistic_gen.pth", | |
| ], | |
| "RealESRGAN_x4plus.pth": [ | |
| "https://huggingface.co/lucataco/RealESRGAN-x4plus/resolve/main/RealESRGAN_x4plus.pth", | |
| "https://github.com/xinntao/Real-ESRGAN/releases/download/v0.1.0/RealESRGAN_x4plus.pth", | |
| ], | |
| "GFPGANv1.4.pth": [ | |
| "https://huggingface.co/Xintao/GFPGAN/resolve/main/GFPGANv1.4.pth", | |
| "https://github.com/TencentARC/GFPGAN/releases/download/v1.3.4/GFPGANv1.4.pth", | |
| ], | |
| } | |
| MODEL_SIZES = { | |
| "ColorizeArtistic_gen.pth": 243 * 1024 * 1024, | |
| "RealESRGAN_x4plus.pth": 64 * 1024 * 1024, | |
| "GFPGANv1.4.pth": 333 * 1024 * 1024, | |
| } | |
| # --------------------------------------------------------------------------- | |
| # Parallel model downloader | |
| # --------------------------------------------------------------------------- | |
| def _download_one(name: str, path: Path): | |
| if path.exists() and path.stat().st_size > 1024: | |
| return | |
| mirrors = MODEL_MIRRORS.get(name, []) | |
| if not mirrors: | |
| raise FileNotFoundError(f"No download URL for {name}") | |
| size_mb = MODEL_SIZES.get(name, 0) // 1024 // 1024 | |
| for url in mirrors: | |
| try: | |
| print(f"Downloading {name} ({size_mb} MB)...") | |
| r = requests.get(url, stream=True, timeout=300) | |
| r.raise_for_status() | |
| tmp = path.with_suffix(".tmp") | |
| with open(tmp, "wb") as f: | |
| for chunk in r.iter_content(8192): | |
| f.write(chunk) | |
| tmp.rename(path) | |
| print(f" {name} done") | |
| return | |
| except Exception as e: | |
| print(f" mirror failed: {e}") | |
| raise RuntimeError(f"Failed to download {name} from all mirrors") | |
| def _ensure_model(name: str): | |
| path = MODEL_DIR / name | |
| _download_one(name, path) | |
| return str(path) | |
| def _ensure_all_models(): | |
| from concurrent.futures import ThreadPoolExecutor, as_completed | |
| needed = [(name, MODEL_DIR / name) for name in MODEL_MIRRORS | |
| if not (MODEL_DIR / name).exists() or (MODEL_DIR / name).stat().st_size < 1024] | |
| if not needed: | |
| return | |
| print(f"Downloading {len(needed)} models in parallel...") | |
| with ThreadPoolExecutor(max_workers=len(needed)) as ex: | |
| fs = {ex.submit(_download_one, name, path): name for name, path in needed} | |
| for f in as_completed(fs): | |
| f.result() | |
| print("All models ready") | |
| # ============================================================================ | |
| # Colorization Model — UNet with ResNet34 encoder | |
| # ============================================================================ | |
| class ColorizationModel(nn.Module): | |
| """UNet with ResNet34 encoder for image colorization (RGB output).""" | |
| def __init__(self): | |
| super().__init__() | |
| from torchvision import models as tv_models | |
| resnet = tv_models.resnet34(weights=tv_models.ResNet34_Weights.IMAGENET1K_V1) | |
| self.enc_conv1 = nn.Sequential(resnet.conv1, resnet.bn1, resnet.relu) # 64 | |
| self.enc_maxpool = resnet.maxpool | |
| self.enc_layer1 = resnet.layer1 # 64 | |
| self.enc_layer2 = resnet.layer2 # 128 | |
| self.enc_layer3 = resnet.layer3 # 256 | |
| self.enc_layer4 = resnet.layer4 # 512 | |
| # Up-path: upsample → cat → conv | |
| def _up(in_ch, out_ch): | |
| return nn.Sequential( | |
| nn.Conv2d(in_ch, out_ch, 3, padding=1), | |
| nn.BatchNorm2d(out_ch), nn.ReLU(True), | |
| ) | |
| self.up4_t = nn.ConvTranspose2d(512, 256, 4, 2, 1) | |
| self.up4_c = _up(512, 256) | |
| self.up3_t = nn.ConvTranspose2d(256, 128, 4, 2, 1) | |
| self.up3_c = _up(256, 128) | |
| self.up2_t = nn.ConvTranspose2d(128, 64, 4, 2, 1) | |
| self.up2_c = _up(128, 64) | |
| self.up1_t = nn.ConvTranspose2d(64, 32, 4, 2, 1) | |
| self.up1_c = _up(96, 32) | |
| self.out = nn.Sequential( | |
| nn.Conv2d(32, 16, 3, padding=1), nn.ReLU(True), | |
| nn.Conv2d(16, 3, 1), nn.Sigmoid(), | |
| ) | |
| def forward(self, x): | |
| f0 = self.enc_conv1(x) # 64, 56×56 | |
| p = self.enc_maxpool(f0) # 64, 28×28 | |
| f1 = self.enc_layer1(p) # 64, 28×28 | |
| f2 = self.enc_layer2(f1) # 128, 14×14 | |
| f3 = self.enc_layer3(f2) # 256, 7×7 | |
| f4 = self.enc_layer4(f3) # 512, 7×7 | |
| d = self.up4_t(f4) # 256, 14×14 | |
| d = self.up4_c(torch.cat([d, f3], 1)) # 256, 14×14 | |
| d = self.up3_t(d) # 128, 28×28 | |
| d = self.up3_c(torch.cat([d, f2], 1)) # 128, 28×28 | |
| d = self.up2_t(d) # 64, 56×56 | |
| d = self.up2_c(torch.cat([d, f1], 1)) # 64, 56×56 | |
| d = self.up1_t(d) # 32, 112×112 | |
| d = self.up1_c(torch.cat([d, f0], 1)) # 32, 112×112 | |
| return self.out(d) # 3, 112×112 | |
| def load_encoder_from_deoldify(self, weights_path): | |
| """Load encoder weights from DeOldify state dict (best-effort).""" | |
| sd = torch.load(weights_path, map_location='cpu', weights_only=False) | |
| sd = sd['model'] if 'model' in sd else sd | |
| mapping = { | |
| 'enc_conv1.0.': 'layers.0.0.', # conv1 | |
| 'enc_conv1.1.': 'layers.0.1.', # bn1 | |
| 'enc_layer1.': 'layers.0.4.', # layer1 | |
| 'enc_layer2.': 'layers.0.5.', # layer2 | |
| 'enc_layer3.': 'layers.0.6.', # layer3 | |
| 'enc_layer4.': 'layers.0.7.', # layer4 | |
| } | |
| n = 0 | |
| for own_key in self.state_dict(): | |
| for prefix, deoldify_prefix in mapping.items(): | |
| if own_key.startswith(prefix): | |
| dk = own_key.replace(prefix, deoldify_prefix) | |
| if dk in sd: | |
| self.state_dict()[own_key].copy_(sd[dk]) | |
| n += 1 | |
| break | |
| print(f"Loaded {n}/{len(list(self.state_dict().keys()))} encoder keys from DeOldify") | |
| return n | |
| # ============================================================================ | |
| # Colorizer — High-level wrapper for colorization inference | |
| # ============================================================================ | |
| class Colorizer: | |
| def __init__(self): | |
| self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu") | |
| self.model = None | |
| self.ready = False | |
| def init(self): | |
| if self.ready: | |
| return | |
| self.model = ColorizationModel() | |
| wpath = _ensure_model("ColorizeArtistic_gen.pth") | |
| self.model.load_encoder_from_deoldify(wpath) | |
| self.model.to(self.device).eval() | |
| self.ready = True | |
| print(f"Colorizer ready on {self.device}") | |
| def colorize(self, img_rgb: np.ndarray) -> Image.Image: | |
| self.init() | |
| img = Image.fromarray(img_rgb).convert("RGB") | |
| w, h = img.size | |
| scale = min(224 / w, 224 / h) | |
| nw, nh = int(w * scale), int(h * scale) | |
| img_small = img.resize((nw, nh), Image.LANCZOS) | |
| pad_w, pad_h = 224 - nw, 224 - nh | |
| img_pad = Image.new("RGB", (224, 224), (0, 0, 0)) | |
| img_pad.paste(img_small, (pad_w // 2, pad_h // 2)) | |
| t = torch.from_numpy(np.array(img_pad).transpose(2, 0, 1)).float().div(255.).unsqueeze(0).to(self.device) | |
| if self.device.type == "cpu": | |
| t = t.half() | |
| with torch.no_grad(): | |
| m = self.model.half() | |
| out = m(t).clamp(0, 1).squeeze(0).cpu().float().numpy().transpose(1, 2, 0) | |
| self.model.float() | |
| else: | |
| with torch.no_grad(): | |
| out = self.model(t).clamp(0, 1).squeeze(0).cpu().numpy().transpose(1, 2, 0) | |
| out = out[pad_h // 2:pad_h // 2 + nh, pad_w // 2:pad_w // 2 + nw] | |
| return Image.fromarray((out * 255).astype(np.uint8)).resize((w, h), Image.LANCZOS) | |
| # ============================================================================ | |
| # Super-Resolution — Real-ESRGAN 4x upscaling | |
| # ============================================================================ | |
| class SuperResolution: | |
| def __init__(self): | |
| self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu") | |
| self.upsampler = None | |
| self.ready = False | |
| def init(self): | |
| if self.ready: | |
| return | |
| from basicsr.archs.rrdbnet_arch import RRDBNet | |
| from realesrgan import RealESRGANer | |
| model = RRDBNet(num_in_ch=3, num_out_ch=3, num_feat=64, num_block=23, num_grow_ch=32, scale=4) | |
| path = _ensure_model("RealESRGAN_x4plus.pth") | |
| self.upsampler = RealESRGANer( | |
| scale=4, model_path=path, model=model, | |
| tile=0, tile_pad=10, pre_pad=0, device=self.device, | |
| fp16=self.device.type == "cpu", | |
| ) | |
| self.ready = True | |
| print(f"SuperResolution ready on {self.device}") | |
| def upscale(self, img_rgb: np.ndarray, scale: int = 4) -> np.ndarray: | |
| self.init() | |
| img_bgr = cv2.cvtColor(img_rgb, cv2.COLOR_RGB2BGR) | |
| out, _ = self.upsampler.enhance(img_bgr, outscale=scale) | |
| return cv2.cvtColor(out, cv2.COLOR_BGR2RGB) | |
| # ============================================================================ | |
| # Face Restoration — GFPGAN face enhancement | |
| # ============================================================================ | |
| class FaceRestoration: | |
| def __init__(self): | |
| self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu") | |
| self.restorer = None | |
| self.ready = False | |
| def init(self): | |
| if self.ready: | |
| return | |
| from gfpgan import GFPGANer | |
| path = _ensure_model("GFPGANv1.4.pth") | |
| self.restorer = GFPGANer( | |
| model_path=path, upscale=1, arch='clean', | |
| channel_multiplier=2, device=self.device, | |
| fp16=self.device.type == "cpu", | |
| ) | |
| self.ready = True | |
| print(f"FaceRestoration ready on {self.device}") | |
| def restore(self, img_rgb: np.ndarray, weight: float = 0.5) -> np.ndarray: | |
| self.init() | |
| img_bgr = cv2.cvtColor(img_rgb, cv2.COLOR_RGB2BGR) | |
| _, _, restored_img = self.restorer.enhance(img_bgr, has_aligned=False, paste_back=True, weight=weight) | |
| if restored_img is None: | |
| return img_rgb | |
| return cv2.cvtColor(restored_img, cv2.COLOR_BGR2RGB) | |
| # ============================================================================ | |
| # Photo Filters — 20+ artistic & classic image effects | |
| # ============================================================================ | |
| class PhotoFilters: | |
| def mono(im): return ImageOps.grayscale(im).convert("RGB") | |
| def sepia(im): | |
| a = np.array(im.convert("RGB"), np.float32) | |
| k = np.array([[.393,.769,.189],[.349,.686,.168],[.272,.534,.131]]) | |
| return Image.fromarray(np.clip(a @ k.T, 0, 255).astype(np.uint8)) | |
| def film(im): | |
| a = np.array(im.convert("RGB"), np.float32) | |
| a = np.clip(a * [1., .93, .85] + np.random.randn(*a.shape).astype(np.float32) * 8, 0, 255) | |
| h, w = a.shape[:2] | |
| X, Y = np.meshgrid(np.linspace(-1, 1, w), np.linspace(-1, 1, h)) | |
| v = np.clip(1 - np.sqrt(X**2 + Y**2) * .4, .25, 1.) | |
| r = ImageEnhance.Contrast(Image.fromarray((a * v[..., None]).astype(np.uint8))).enhance(1.15) | |
| return ImageEnhance.Brightness(r).enhance(.95) | |
| def canny(im, t1=50, t2=150): | |
| e = cv2.Canny(cv2.cvtColor(np.array(im.convert("RGB")), cv2.COLOR_RGB2GRAY), t1, t2) | |
| return Image.fromarray(cv2.cvtColor(e, cv2.COLOR_GRAY2RGB)) | |
| def pencil(im): | |
| g = cv2.cvtColor(np.array(im.convert("RGB")), cv2.COLOR_RGB2GRAY) | |
| s = cv2.divide(g, 255 - cv2.GaussianBlur(255 - g, (21, 21), 0), scale=256) | |
| return Image.fromarray(cv2.cvtColor(s, cv2.COLOR_GRAY2RGB)) | |
| def cartoon(im): | |
| a = np.array(im.convert("RGB")) | |
| g = cv2.medianBlur(cv2.cvtColor(a, cv2.COLOR_RGB2GRAY), 5) | |
| e = cv2.adaptiveThreshold(g, 255, cv2.ADAPTIVE_THRESH_MEAN_C, cv2.THRESH_BINARY, 9, 10) | |
| return Image.fromarray(cv2.bitwise_and(cv2.bilateralFilter(a, 9, 300, 300), a, mask=e)) | |
| def emboss(im): return im.convert("RGB").filter(PILFilter.EMBOSS) | |
| def sharpen(im): return ImageEnhance.Sharpness(im.convert("RGB")).enhance(3.) | |
| def oil(im): | |
| return Image.fromarray(cv2.medianBlur(cv2.bilateralFilter(np.array(im.convert("RGB")), 9, 150, 150), 3)) | |
| def sobel(im): | |
| g = cv2.cvtColor(np.array(im.convert("RGB")), cv2.COLOR_RGB2GRAY).astype(np.float32) | |
| m = np.sqrt(cv2.Sobel(g, cv2.CV_64F, 1, 0, ksize=3)**2 + cv2.Sobel(g, cv2.CV_64F, 0, 1, ksize=3)**2) | |
| mx = m.max() | |
| return Image.fromarray(cv2.cvtColor(np.clip(m / mx * 255 if mx > 0 else 0, 0, 255).astype(np.uint8), cv2.COLOR_GRAY2RGB)) | |
| def pixelate(im, size=8): | |
| a = np.array(im.convert("RGB")) | |
| h, w = a.shape[:2] | |
| sm = cv2.resize(a, (w // size, h // size), interpolation=cv2.INTER_LINEAR) | |
| return Image.fromarray(cv2.resize(sm, (w, h), interpolation=cv2.INTER_NEAREST)) | |
| def invert(im): return ImageOps.invert(im.convert("RGB")) | |
| def posterize(im, bits=4): return ImageOps.posterize(im.convert("RGB"), bits) | |
| def solarize(im, thresh=128): return ImageOps.solarize(im.convert("RGB"), thresh) | |
| def edge_enhance(im): return im.convert("RGB").filter(PILFilter.EDGE_ENHANCE_MORE) | |
| def blur(im): return im.convert("RGB").filter(PILFilter.GaussianBlur(radius=5)) | |
| def detail(im): return im.convert("RGB").filter(PILFilter.DETAIL) | |
| def contour(im): | |
| a = cv2.Canny(cv2.cvtColor(np.array(im.convert("RGB")), cv2.COLOR_RGB2GRAY), 30, 100) | |
| return Image.fromarray(cv2.cvtColor(255 - a, cv2.COLOR_GRAY2RGB)) | |
| def thermal(im): | |
| a = cv2.cvtColor(np.array(im.convert("RGB")), cv2.COLOR_RGB2GRAY) | |
| return Image.fromarray(cv2.applyColorMap(a, cv2.COLORMAP_INFERNO)) | |
| FILTER_CHOICES = [ | |
| "Colorize (DeOldify)", "Monochrome", "Sepia", "Film Photo", | |
| "Canny Edges", "Sobel Edges", "Contour", "Pencil Sketch", | |
| "Cartoonify", "Oil Painting", "Pixelate", "Blur", | |
| "Emboss", "Edge Enhance", "Sharpen", "Detail", | |
| "Invert", "Posterize", "Solarize", "Thermal", | |
| "Super Resolution (4x)", "Face Restoration", | |
| ] | |
| # ============================================================================ | |
| # Processor — Orchestrates filter dispatch & resolution management | |
| # ============================================================================ | |
| class PhotoProcessor: | |
| def __init__(self): | |
| self.colorizer = Colorizer() | |
| self.sr = SuperResolution() | |
| self.fr = FaceRestoration() | |
| self.ready = False | |
| def init(self): | |
| if not self.ready: | |
| _ensure_all_models() | |
| self.colorizer.init() | |
| self.ready = True | |
| def _resize(im: Image.Image, r: int) -> Image.Image: | |
| w, h = im.size | |
| mx = max(w, h) | |
| if mx <= r: | |
| return im | |
| return im.resize((r, int(h * r / w)) if w >= h else (int(w * r / h), r), Image.LANCZOS) | |
| def process(self, im: Image.Image | None, flt: str = "Colorize (DeOldify)", res: int = 1024, fr_weight: float = 0.5) -> Image.Image | None: | |
| if im is None: | |
| return None | |
| self.init() | |
| d = { | |
| "Colorize (DeOldify)": lambda i: self.colorizer.colorize(np.array(i.convert("RGB"))), | |
| "Monochrome": PhotoFilters.mono, "Sepia": PhotoFilters.sepia, | |
| "Film Photo": PhotoFilters.film, | |
| "Canny Edges": lambda i: PhotoFilters.canny(i, 50, 150), | |
| "Sobel Edges": PhotoFilters.sobel, "Contour": PhotoFilters.contour, | |
| "Pencil Sketch": PhotoFilters.pencil, "Cartoonify": PhotoFilters.cartoon, | |
| "Oil Painting": PhotoFilters.oil, "Pixelate": PhotoFilters.pixelate, | |
| "Blur": PhotoFilters.blur, "Emboss": PhotoFilters.emboss, | |
| "Edge Enhance": PhotoFilters.edge_enhance, "Sharpen": PhotoFilters.sharpen, | |
| "Detail": PhotoFilters.detail, "Invert": PhotoFilters.invert, | |
| "Posterize": PhotoFilters.posterize, "Solarize": PhotoFilters.solarize, | |
| "Thermal": PhotoFilters.thermal, | |
| "Super Resolution (4x)": lambda i: Image.fromarray(self.sr.upscale(np.array(i.convert("RGB")))), | |
| "Face Restoration": lambda i: Image.fromarray(self.fr.restore(np.array(i.convert("RGB")), weight=fr_weight)), | |
| } | |
| r = d.get(flt, d["Colorize (DeOldify)"])(im.convert("RGB")) | |
| return self._resize(Image.fromarray(r) if isinstance(r, np.ndarray) else r, res) | |