Spaces:
Running on Zero
Running on Zero
| """ | |
| Material maps for the textures this Space makes, and for any picture of a | |
| surface (/material_maps): one picture in, the maps a PBR renderer needs out, | |
| all tiling when the picture tiles. `compute` runs on the GPU and is called | |
| from inside app.py's GPU functions, so a texture and its maps are one run. | |
| How it goes: | |
| - Two small ESRGAN nets trained on texture sets (Joey Ballentine's Material | |
| Map Generator, Apache-2.0): one gives a tangent-space normal map, the other | |
| displacement and roughness. They are what reads a painted brick as raised | |
| and its white mortar as sunk, which brightness alone gets backwards. | |
| - Lotus-G normal (Apache-2.0, Stable Diffusion 2 fine-tuned for surface | |
| normals, one step) gives the broad shape: the rounded top of a cobble, the | |
| bevel of a brick. Its low frequencies and the ESRGAN detail are merged in | |
| slope space with the slopes of the displacement map, so normal and height | |
| agree. | |
| - CLIP (MIT) names the material class, which sets the roughness level and | |
| whether anything is metal; the ESRGAN roughness adds the variation. | |
| Every convolution pads circularly (the ESRGAN input is wrapped, the Lotus UNet | |
| and VAE have their padding mode switched), and every filter wraps, so a | |
| seamless picture gives seamless maps. | |
| """ | |
| import os | |
| import tempfile | |
| import time | |
| os.environ["GRADIO_TEMP_DIR"] = os.path.join(tempfile.gettempdir(), "gradio") | |
| os.makedirs(os.environ["GRADIO_TEMP_DIR"], exist_ok=True) | |
| import gradio as gr | |
| import numpy as np | |
| import torch | |
| import torch.nn.functional as F | |
| from diffusers import AutoencoderKL, UNet2DConditionModel | |
| from huggingface_hub import hf_hub_download | |
| from PIL import Image | |
| from spandrel import ModelLoader | |
| from transformers import CLIPModel, CLIPProcessor | |
| import embeddings | |
| from embeddings import CLASSES, CLIP_REPO, LOTUS_REPO, features | |
| DEVICE = "cuda" | |
| MAX_SIDE = 2048 | |
| # The side the models read the picture at; the maps go out at the picture's own | |
| # size. Measured 2026-10-07 on six Poly Haven scans at 2K: read at 2048 the nets | |
| # gave normals twice too steep, 14 to 31 degrees off the scan against 8 to 25 | |
| # when read at 1024 and enlarged, with nothing above 1024 that matched the scan. | |
| MODEL_SIDE = 1024 | |
| # Slopes scaled to half: measured against Poly Haven scans (scripts/compare-normal-maps.mjs), | |
| # full strength was about twice too steep, and invented relief on painted, flat surfaces. | |
| NORMAL_STRENGTH = 0.5 | |
| LOTUS_SIDE = 768 # Lotus-G is Stable Diffusion 2 base: 512 to 768 is home. | |
| MAPS_REPO = "InvokeAI/pbr-material-maps" | |
| MAPS_REVISION = "b7ca9ebc6e14688a69d41872d2b9c80ea453e8f0" | |
| def _log(*parts): | |
| print("[maps]", *parts, flush=True) | |
| def _circular(module: torch.nn.Module) -> None: | |
| for m in module.modules(): | |
| if isinstance(m, torch.nn.Conv2d) and m.padding not in (0, (0, 0)): | |
| m.padding_mode = "circular" | |
| def _esrgan(name: str) -> torch.nn.Module: | |
| path = hf_hub_download(MAPS_REPO, name, revision=MAPS_REVISION) | |
| return ModelLoader().load_from_file(path).model.eval().half().to(DEVICE) | |
| normal_net = _esrgan("normal_map_generator.safetensors") | |
| franken_net = _esrgan("franken_map_generator.safetensors") | |
| lotus_unet = UNet2DConditionModel.from_pretrained(LOTUS_REPO, subfolder="unet", torch_dtype=torch.float16).to(DEVICE) | |
| lotus_vae = AutoencoderKL.from_pretrained(LOTUS_REPO, subfolder="vae", torch_dtype=torch.float16).to(DEVICE) | |
| _circular(lotus_unet) | |
| _circular(lotus_vae) | |
| # The task embedding that selects the normal head (see Lotus's infer.py). | |
| _task = torch.tensor([[1.0, 0.0]]) | |
| TASK_EMB = torch.cat([torch.sin(_task), torch.cos(_task)], dim=-1).half().to(DEVICE) | |
| clip_model = CLIPModel.from_pretrained(CLIP_REPO).eval().half().to(DEVICE) | |
| clip_processor = CLIPProcessor.from_pretrained(CLIP_REPO) | |
| # The empty prompt Lotus runs with and the class words for CLIP, embedded once | |
| # in a process of their own (embeddings.py says why). | |
| _embedded = embeddings.load() | |
| EMPTY_PROMPT = _embedded["empty_prompt"].half().to(DEVICE) | |
| CLASS_EMB = _embedded["class_emb"].to(DEVICE) | |
| _log("models ready") | |
| # --- plain-array helpers, all wrapping at the edges ------------------------- | |
| def _blur(a: torch.Tensor, sigma: float) -> torch.Tensor: | |
| """Separable Gaussian on an (H, W) tensor, wrapping around the edges.""" | |
| if sigma <= 0: | |
| return a | |
| radius = max(1, int(3 * sigma)) | |
| x = torch.arange(-radius, radius + 1, device=a.device, dtype=a.dtype) | |
| k = torch.exp(-(x**2) / (2 * sigma**2)) | |
| k = k / k.sum() | |
| out = F.pad(a[None, None], (radius, radius, 0, 0), mode="circular") | |
| out = F.conv2d(out, k.view(1, 1, 1, -1)) | |
| out = F.pad(out, (0, 0, radius, radius), mode="circular") | |
| return F.conv2d(out, k.view(1, 1, -1, 1))[0, 0] | |
| def _resize_wrap(x: torch.Tensor, size: tuple[int, int]) -> torch.Tensor: | |
| """Resize (N, C, H, W) so the result still tiles: pad by wrapping, scale, crop.""" | |
| h, w = x.shape[-2:] | |
| if (h, w) == size: | |
| return x | |
| pad = 4 | |
| big = F.pad(x, (pad, pad, pad, pad), mode="circular") | |
| sy, sx = size[0] / h, size[1] / w | |
| out = F.interpolate(big, size=(round((h + 2 * pad) * sy), round((w + 2 * pad) * sx)), mode="bicubic", align_corners=False) | |
| oy, ox = round(pad * sy), round(pad * sx) | |
| return out[..., oy : oy + size[0], ox : ox + size[1]] | |
| def _enlarge(a: torch.Tensor, size: tuple[int, int]) -> torch.Tensor: | |
| """An (H, W) map at the picture's size, still tiling.""" | |
| return _resize_wrap(a[None, None], size)[0, 0] | |
| def _slopes(n: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: | |
| """(3, H, W) normals, x right, y up → slopes dh/dx and dh/dy_up, tilt removed. | |
| A normal is (-dh/dx, -dh/dy, 1) normalised.""" | |
| nz = n[2].clamp(min=0.2) | |
| p, q = -n[0] / nz, -n[1] / nz | |
| return p - p.mean(), q - q.mean() | |
| def _grad(h: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: | |
| """Central differences with wrap: dh/dx and dh/dy_up (rows run down).""" | |
| gx = (torch.roll(h, -1, 1) - torch.roll(h, 1, 1)) / 2 | |
| gy = (torch.roll(h, 1, 0) - torch.roll(h, -1, 0)) / 2 | |
| return gx, gy | |
| def _stretch(a: torch.Tensor, lo: float = 0.005, hi: float = 0.995) -> torch.Tensor: | |
| flat = a.flatten() | |
| if flat.numel() > 1_000_000: | |
| flat = flat[:: flat.numel() // 1_000_000 + 1] | |
| a_lo, a_hi = torch.quantile(flat, lo), torch.quantile(flat, hi) | |
| return ((a - a_lo) / (a_hi - a_lo).clamp(min=1e-6)).clamp(0, 1) | |
| def _smoothstep(e0: float, e1: float, x: torch.Tensor) -> torch.Tensor: | |
| t = ((x - e0) / (e1 - e0)).clamp(0, 1) | |
| return t * t * (3 - 2 * t) | |
| # --- the models -------------------------------------------------------------- | |
| def _run_esrgan(net: torch.nn.Module, rgb: torch.Tensor) -> torch.Tensor: | |
| """(1, 3, H, W) in [0, 1] → (3, H, W) in [0, 1]; wrapped so the edges tile.""" | |
| pad = 32 | |
| x = F.pad(rgb, (pad, pad, pad, pad), mode="circular").half() | |
| return net(x)[0, :, pad:-pad, pad:-pad].float().clamp(0, 1) | |
| def _run_lotus(rgb: torch.Tensor) -> torch.Tensor: | |
| """(1, 3, H, W) in [0, 1] → (3, H, W) unit normals, x right, y up, z out.""" | |
| h, w = rgb.shape[-2:] | |
| scale = min(1.0, LOTUS_SIDE / max(h, w)) | |
| size = (max(64, round(h * scale / 64) * 64), max(64, round(w * scale / 64) * 64)) | |
| x = _resize_wrap(rgb, size) * 2 - 1 | |
| latents = lotus_vae.encode(x.half()).latent_dist.mode() * lotus_vae.config.scaling_factor | |
| noise = torch.randn(latents.shape, generator=torch.Generator(DEVICE).manual_seed(0), device=DEVICE, dtype=latents.dtype) | |
| x0 = lotus_unet( | |
| torch.cat([latents, noise], dim=1), | |
| torch.tensor([999], device=DEVICE), | |
| encoder_hidden_states=EMPTY_PROMPT, | |
| class_labels=TASK_EMB, | |
| return_dict=False, | |
| )[0] | |
| decoded = lotus_vae.decode(x0 / lotus_vae.config.scaling_factor, return_dict=False)[0].float().clamp(-1, 1) | |
| n = _resize_wrap(decoded, (h, w))[0] | |
| return n / n.norm(dim=0, keepdim=True).clamp(min=1e-6) | |
| def _classify(image: Image.Image) -> tuple[torch.Tensor, list[tuple[str, float]]]: | |
| inputs = clip_processor(images=image, return_tensors="pt").to(DEVICE) | |
| emb = F.normalize(features(clip_model.get_image_features(pixel_values=inputs.pixel_values.half())).float(), dim=-1) | |
| probs = (100 * emb @ CLASS_EMB.T).softmax(dim=-1)[0] | |
| order = probs.argsort(descending=True)[:3].tolist() | |
| return probs, [(CLASSES[i][0], round(float(probs[i]), 3)) for i in order] | |
| def compute(image: Image.Image): | |
| """On the GPU: the four maps as arrays in [0, 1] at the picture's size, and | |
| what the surface was taken for. The models read it at MODEL_SIDE at most.""" | |
| t0 = time.time() | |
| full, picture = (image.height, image.width), image | |
| if max(image.size) > MODEL_SIDE: | |
| if image.width == image.height and image.width % MODEL_SIDE == 0: | |
| image = image.reduce(image.width // MODEL_SIDE) # whole blocks: a tile keeps its wrap | |
| else: | |
| scale = MODEL_SIDE / max(image.size) | |
| image = image.resize((max(8, round(image.width * scale)), max(8, round(image.height * scale))), Image.LANCZOS) | |
| w, h = image.size | |
| rgb = torch.from_numpy(np.asarray(image, np.float32) / 255).permute(2, 0, 1)[None].to(DEVICE) | |
| s = max(w, h) / 768 # filter sizes were tuned at 768 px | |
| es_normal = _run_esrgan(normal_net, rgb) * 2 - 1 | |
| franken = _run_esrgan(franken_net, rgb) | |
| lotus = _run_lotus(rgb) | |
| probs, top = _classify(picture) # CLIP sizes it itself | |
| _log(f"models {time.time() - t0:.2f}s", top) | |
| # Height: the texture-trained displacement. Normal: the displacement's | |
| # slopes, plus Lotus's broad shape and the ESRGAN normal's fine detail. | |
| height = _stretch(franken[2]) | |
| pd, qd = _grad(height) | |
| pd, qd = pd * max(w, h) / 40, qd * max(w, h) / 40 | |
| pl, ql = _slopes(lotus) | |
| pe, qe = _slopes(es_normal) | |
| p = (_blur(pl, 2 * s) + pe - _blur(pe, 3 * s) + pd) / 2 * NORMAL_STRENGTH | |
| q = (_blur(ql, 2 * s) + qe - _blur(qe, 3 * s) + qd) / 2 * NORMAL_STRENGTH | |
| p, q = _enlarge(p, full), _enlarge(q, full) | |
| normal = torch.stack([-p, -q, torch.ones_like(p)]) | |
| normal = normal / normal.norm(dim=0, keepdim=True) | |
| # Roughness: the class sets the level, the ESRGAN map the variation. | |
| level = sum(float(probs[i]) * c[2] for i, c in enumerate(CLASSES)) | |
| rough = franken[1] | |
| roughness = (level + (rough - rough.median()) * 1.2).clamp(0.04, 1) | |
| # Metallic: bare metal is metal all over; rusted metal only where grey | |
| # steel shows (low saturation). Everything else is not metal. | |
| metal = sum(float(probs[i]) for i, c in enumerate(CLASSES) if c[3] == "metal") | |
| rust = sum(float(probs[i]) for i, c in enumerate(CLASSES) if c[3] == "rust") | |
| mx, mn = rgb[0].max(dim=0).values, rgb[0].min(dim=0).values | |
| saturation = (mx - mn) / mx.clamp(min=1e-3) | |
| bare = _smoothstep(0.35, 0.15, saturation) | |
| metallic = _smoothstep(0.35, 0.65, metal + rust * bare) | |
| roughness = roughness - metallic * 0.15 | |
| height, roughness, metallic = (_enlarge(m, full).clamp(0, 1) for m in (height, roughness, metallic)) | |
| info = { | |
| "material": top[0][0], | |
| "classes": [{"label": label, "p": p_} for label, p_ in top], | |
| "roughness_level": round(level, 3), | |
| "metal": round(metal, 3), | |
| "gpu_seconds": round(time.time() - t0, 2), | |
| } | |
| out = ( | |
| (normal.permute(1, 2, 0) * 0.5 + 0.5).clamp(0, 1).cpu().numpy(), | |
| height.cpu().numpy(), | |
| roughness.clamp(0, 1).cpu().numpy(), | |
| metallic.clamp(0, 1).cpu().numpy(), | |
| ) | |
| torch.cuda.empty_cache() | |
| return out, info | |
| def _save_png(array: np.ndarray, stem: str, bits: int = 8) -> str: | |
| path = os.path.join(os.environ["GRADIO_TEMP_DIR"], f"{stem}-{time.time_ns()}.png") | |
| if bits == 16: | |
| Image.fromarray((array * 65535).round().astype(np.uint16)).save(path) | |
| else: | |
| Image.fromarray((array * 255).round().astype(np.uint8)).save(path, optimize=False, compress_level=6) | |
| return path | |
| def prepare(image) -> Image.Image: | |
| """The picture as `compute` takes it: RGB, no larger than MAX_SIDE.""" | |
| if image is None: | |
| raise gr.Error("Upload a picture of a surface.") | |
| if not isinstance(image, Image.Image): | |
| image = Image.open(image) | |
| image = image.convert("RGB") | |
| if max(image.size) > MAX_SIDE: | |
| scale = MAX_SIDE / max(image.size) | |
| image = image.resize((max(8, round(image.width * scale)), max(8, round(image.height * scale))), Image.LANCZOS) | |
| return image | |
| def save(arrays, info: dict, directx: bool = False): | |
| """The maps as PNG files: normal (OpenGL unless `directx`), height | |
| (16-bit), roughness and metallic, and the JSON note.""" | |
| normal, height, roughness, metallic = arrays | |
| if directx: | |
| normal = normal.copy() | |
| normal[..., 1] = 1 - normal[..., 1] | |
| info["convention"] = "directx" if directx else "opengl" | |
| info["size"] = [normal.shape[1], normal.shape[0]] | |
| files = ( | |
| _save_png(normal, "normal"), | |
| _save_png(height, "height", bits=16), | |
| _save_png(roughness, "roughness"), | |
| _save_png(metallic, "metallic"), | |
| ) | |
| _log("done", info) | |
| return (*files, info) | |