Deploy from 3dvalley spaces/material-maps

app.py

@@ -79,13 +79,13 @@ with torch.no_grad():
     _tok = CLIPTokenizer.from_pretrained(LOTUS_REPO, subfolder="tokenizer")
     _enc = CLIPTextModel.from_pretrained(LOTUS_REPO, subfolder="text_encoder")
     _ids = _tok([""], padding="max_length", max_length=_tok.model_max_length, return_tensors="pt").input_ids
-    EMPTY_PROMPT = _enc(_ids)[0].half().to(DEVICE)
+    EMPTY_PROMPT = _enc(_ids)[0].half().to(DEVICE)  # computed on the CPU, like the CLIP classes below
     del _tok, _enc
 # 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, torch_dtype=torch.float16).eval().to(DEVICE)
+clip_model = CLIPModel.from_pretrained(CLIP_REPO).eval()
 clip_processor = CLIPProcessor.from_pretrained(CLIP_REPO)
 
 # (label, words for CLIP, roughness level, metal): "metal" is bare metal all
@@ -112,9 +112,12 @@ CLASSES = [
     ("polished metal", "a shiny polished metal, chrome, gold or copper texture", 0.15, "metal"),
     ("snow", "a snow or ice texture", 0.3, None),
 ]
+# The class words are embedded once, on the CPU: ZeroGPU only lends a GPU
+# inside a @spaces.GPU call, so nothing runs on "cuda" at start-up.
 with torch.no_grad():
-    _t = clip_processor(text=[c[1] for c in CLASSES], return_tensors="pt", padding=True).to(DEVICE)
-    CLASS_EMB = F.normalize(clip_model.get_text_features(**_t).float(), dim=-1)
+    _t = clip_processor(text=[c[1] for c in CLASSES], return_tensors="pt", padding=True)
+    CLASS_EMB = F.normalize(clip_model.get_text_features(**_t).float(), dim=-1).to(DEVICE)
+clip_model = clip_model.half().to(DEVICE)
 _log("models ready")