| import numpy as np |
| import torch |
| from pathlib import Path |
| from huggingface_hub.constants import HF_HUB_CACHE |
| from PIL import Image |
| from spandrel import ImageModelDescriptor, ModelLoader |
|
|
| CLEAN_BUCKET = "HyperHail/CB" |
| CLEAN_FILE = "upscale_models/3_1x-Archivist_Soft.pth" |
| TILE_SIZE = 256 |
| TILE_PAD = 32 |
| PIXEL_MAX = 255 |
|
|
|
|
| def download_clean(): |
| path = Path(HF_HUB_CACHE) / f"buckets--{CLEAN_BUCKET.replace('/', '--')}" / CLEAN_FILE |
| if not path.is_file(): |
| path.parent.mkdir(parents=True, exist_ok=True) |
| url = f"https://huggingface.co/buckets/{CLEAN_BUCKET}/resolve/{CLEAN_FILE}" |
| torch.hub.download_url_to_file(url, str(path), progress=False) |
| return str(path) |
|
|
|
|
| def load_clean(device): |
| model = ModelLoader().load_from_file(download_clean()) |
| if not isinstance(model, ImageModelDescriptor) or model.scale != 1: |
| raise ValueError("Clean after requires a 1x image model") |
| if model.input_channels != 3 or model.output_channels != 3: |
| raise ValueError("Clean after requires an RGB model") |
| return model.to(device).eval() |
|
|
|
|
| @torch.inference_mode() |
| def clean_image(model, input_path, output_path): |
| device = next(model.model.parameters()).device |
| with Image.open(input_path) as source: |
| image = source.convert("RGB") |
| alpha = source.convert("RGBA").getchannel("A") if "A" in source.getbands() or "transparency" in source.info else None |
| width, height = image.size |
| result = Image.new("RGB", image.size) |
| for y in range(0, height, TILE_SIZE): |
| for x in range(0, width, TILE_SIZE): |
| right, bottom = min(x + TILE_SIZE, width), min(y + TILE_SIZE, height) |
| left, top = max(0, x - TILE_PAD), max(0, y - TILE_PAD) |
| box = (left, top, min(width, right + TILE_PAD), min(height, bottom + TILE_PAD)) |
| pixels = np.array(image.crop(box), dtype=np.float32) / PIXEL_MAX |
| tile = torch.from_numpy(pixels).permute(2, 0, 1).unsqueeze(0).to(device) |
| output = model(tile)[0].clamp(0, 1).mul(PIXEL_MAX).round().byte() |
| patch = Image.fromarray(output.permute(1, 2, 0).cpu().numpy()) |
| result.paste(patch.crop((x - left, y - top, right - left, bottom - top)), (x, y)) |
| if alpha is not None: |
| result.putalpha(alpha) |
| result.save(output_path, format="PNG") |
|
|
