Spaces:
Running on Zero
Running on Zero
| """ | |
| The two text embeddings the maps need, made once at start-up: the empty | |
| prompt Lotus runs with, and the class words CLIP compares a picture to. | |
| They are made in a process of their own (`python embeddings.py out.safetensors`, | |
| started by `load`). On ZeroGPU the main process must never run a model: a | |
| forward pass there, on the CPU, left every later GPU call failing with | |
| "No CUDA GPUs are available". Loading weights and `.to("cuda")` are fine. | |
| """ | |
| import os | |
| import subprocess | |
| import sys | |
| import tempfile | |
| LOTUS_REPO = "jingheya/lotus-normal-g-v1-1" | |
| CLIP_REPO = "openai/clip-vit-base-patch32" | |
| # (label, words for CLIP, roughness level, metal): "metal" is bare metal all | |
| # over, "rust" is metal only where grey steel shows through. | |
| CLASSES = [ | |
| ("brick", "a brick wall texture", 0.85, None), | |
| ("stone", "a cobblestone or stone paving texture", 0.8, None), | |
| ("rock", "a rough natural rock texture", 0.85, None), | |
| ("concrete", "a concrete or plaster wall texture", 0.9, None), | |
| ("asphalt", "an asphalt road texture", 0.9, None), | |
| ("wood", "a wooden planks texture", 0.7, None), | |
| ("varnished wood", "a varnished polished wood floor texture", 0.35, None), | |
| ("bark", "a tree bark texture", 0.9, None), | |
| ("ground", "a dirt, soil, mud or sand ground texture", 0.95, None), | |
| ("vegetation", "a grass, moss or leaves texture", 0.8, None), | |
| ("marble", "a polished marble texture", 0.2, None), | |
| ("tiles", "a glazed ceramic tiles texture", 0.25, None), | |
| ("fabric", "a fabric, cloth or carpet texture", 0.9, None), | |
| ("leather", "a leather texture", 0.6, None), | |
| ("plastic", "a plastic surface texture", 0.4, None), | |
| ("painted metal", "a painted metal surface texture", 0.5, None), | |
| ("rusted metal", "a rusty corroded metal texture", 0.8, "rust"), | |
| ("brushed metal", "a brushed steel or aluminium metal texture", 0.35, "metal"), | |
| ("polished metal", "a shiny polished metal, chrome, gold or copper texture", 0.15, "metal"), | |
| ("snow", "a snow or ice texture", 0.3, None), | |
| ] | |
| def features(out): | |
| """CLIP's projected embedding: transformers 5 wraps it in a model output.""" | |
| return out if hasattr(out, "shape") else out.pooler_output | |
| def load() -> dict: | |
| """Both embeddings, on the CPU, from a process that ends before any GPU is asked for.""" | |
| from safetensors.torch import load_file | |
| path = os.path.join(tempfile.gettempdir(), "maps-embeddings.safetensors") | |
| subprocess.run([sys.executable, os.path.abspath(__file__), path], check=True) | |
| return load_file(path) | |
| def _embed(path: str) -> None: | |
| import torch | |
| import torch.nn.functional as F | |
| from safetensors.torch import save_file | |
| from transformers import CLIPModel, CLIPProcessor, CLIPTextModel, CLIPTokenizer | |
| with torch.no_grad(): | |
| tokenizer = CLIPTokenizer.from_pretrained(LOTUS_REPO, subfolder="tokenizer") | |
| encoder = CLIPTextModel.from_pretrained(LOTUS_REPO, subfolder="text_encoder") | |
| ids = tokenizer([""], padding="max_length", max_length=tokenizer.model_max_length, return_tensors="pt").input_ids | |
| empty_prompt = encoder(ids)[0].float() | |
| clip = CLIPModel.from_pretrained(CLIP_REPO).eval() | |
| words = CLIPProcessor.from_pretrained(CLIP_REPO)(text=[c[1] for c in CLASSES], return_tensors="pt", padding=True) | |
| class_emb = F.normalize(features(clip.get_text_features(**words)).float(), dim=-1) | |
| save_file({"empty_prompt": empty_prompt.contiguous(), "class_emb": class_emb.contiguous()}, path) | |
| if __name__ == "__main__": | |
| _embed(sys.argv[1]) | |