jasonod888
Add normals; keep intrinsics path unchanged; pin diffusers 0.33.1 and peft>=0.17.0
109fd13 | import base64 | |
| import io | |
| from typing import Any, Dict, Optional | |
| import numpy as np | |
| import torch | |
| from PIL import Image | |
| from diffusers import MarigoldIntrinsicsPipeline, utils as diffusers_utils | |
| # We will import MarigoldNormalsPipeline lazily to avoid loading it unless asked | |
| NormalsPipe = None | |
| def _clean_b64(s: str) -> bytes: | |
| s = s.strip() | |
| if s.startswith("data:image"): | |
| s = s.split(",", 1)[1] | |
| s = "".join(s.split()) # strip whitespace/newlines | |
| pad = (-len(s)) % 4 | |
| if pad: | |
| s += "=" * pad | |
| try: | |
| return base64.b64decode(s, validate=False) | |
| except Exception: | |
| return base64.urlsafe_b64decode(s) | |
| def _image_from_inputs(obj: Any) -> Image.Image: | |
| if isinstance(obj, dict): | |
| if "image_url" in obj: | |
| return diffusers_utils.load_image(obj["image_url"]).convert("RGB") | |
| if "image_b64" in obj: | |
| raw = _clean_b64(obj["image_b64"]) | |
| return Image.open(io.BytesIO(raw)).convert("RGB") | |
| raise ValueError("inputs dict must contain 'image_url' or 'image_b64'.") | |
| if isinstance(obj, (str, bytes, bytearray)): | |
| if isinstance(obj, (bytes, bytearray)): | |
| obj = obj.decode("utf-8", errors="ignore") | |
| s = obj.strip() | |
| if s.startswith("http://") or s.startswith("https://"): | |
| return diffusers_utils.load_image(s).convert("RGB") | |
| raw = _clean_b64(s) | |
| return Image.open(io.BytesIO(raw)).convert("RGB") | |
| raise ValueError("Unsupported 'inputs' type. Provide URL string, base64 string, or {'image_url'|'image_b64'}.") | |
| def _to_uint8_png_b64(img_like: Any) -> Optional[str]: | |
| """ | |
| Best-effort: accept PIL.Image or np.ndarray / torch.Tensor (H,W,3) or (3,H,W), | |
| map to [0,255] and return PNG b64 string. | |
| """ | |
| try: | |
| if isinstance(img_like, Image.Image): | |
| buf = io.BytesIO() | |
| img_like.save(buf, format="PNG") | |
| return base64.b64encode(buf.getvalue()).decode("utf-8") | |
| if isinstance(img_like, torch.Tensor): | |
| arr = img_like.detach().cpu().float().numpy() | |
| elif isinstance(img_like, np.ndarray): | |
| arr = img_like | |
| else: | |
| return None | |
| # Move channels-last if needed | |
| if arr.ndim == 3 and arr.shape[0] in (1, 3) and arr.shape[-1] not in (1, 3): | |
| arr = np.transpose(arr, (1, 2, 0)) | |
| if arr.ndim != 3 or arr.shape[-1] not in (1, 3): | |
| # If it's something else (e.g., batch), try take first item | |
| if arr.ndim == 4: # (B,H,W,C) | |
| arr = arr[0] | |
| else: | |
| return None | |
| # Heuristic normalization | |
| arr = arr.astype(np.float32) | |
| a_min, a_max = float(np.nanmin(arr)), float(np.nanmax(arr)) | |
| if a_max <= 1.0 and a_min >= 0.0: | |
| # already [0,1] | |
| pass | |
| elif a_max <= 1.0 and a_min >= -1.0: | |
| # [-1,1] normals style | |
| arr = (arr + 1.0) * 0.5 | |
| else: | |
| # generic min-max | |
| rng = (a_max - a_min) if (a_max - a_min) > 1e-6 else 1.0 | |
| arr = (arr - a_min) / rng | |
| arr = np.clip(arr * 255.0, 0, 255).astype(np.uint8) | |
| img = Image.fromarray(arr) | |
| buf = io.BytesIO() | |
| img.save(buf, format="PNG") | |
| return base64.b64encode(buf.getvalue()).decode("utf-8") | |
| except Exception: | |
| return None | |
| class EndpointHandler: | |
| """HF Inference Toolkit-compatible class handler with Intrinsics (as-is) + optional Normals.""" | |
| def __init__(self, path: Optional[str] = None): | |
| device = "cuda" if torch.cuda.is_available() else "cpu" | |
| dtype = torch.float16 if device == "cuda" else torch.float32 | |
| # ---- Intrinsics pipeline (UNCHANGED from your working code) ---- | |
| self.pipe = MarigoldIntrinsicsPipeline.from_pretrained( | |
| "prs-eth/marigold-iid-lighting-v1-1", | |
| torch_dtype=dtype, | |
| ).to(device) | |
| try: | |
| self.pipe.set_progress_bar_config(disable=True) | |
| except Exception: | |
| pass | |
| # Expected targets for lighting | |
| self.targets = ["albedo", "shading", "residual"] | |
| # ---- Normals (lazy) ---- | |
| self._normals_pipe = None | |
| self._device = device | |
| self._dtype = dtype | |
| # Defaults | |
| self.default_intr_steps = 4 # as before | |
| self.fixed_processing_res = 768 # per your latest preference | |
| # ---------------- Normals helpers ---------------- | |
| def _ensure_normals_loaded(self): | |
| global NormalsPipe | |
| if self._normals_pipe is not None: | |
| return | |
| if NormalsPipe is None: | |
| # Lazy import to avoid loading unless requested | |
| from diffusers.pipelines.marigold.pipeline_marigold_normals import MarigoldNormalsPipeline as _NP | |
| NormalsPipe = _NP | |
| self._normals_pipe = NormalsPipe.from_pretrained( | |
| "prs-eth/marigold-normals-v1-1", | |
| torch_dtype=self._dtype, | |
| ).to(self._device) | |
| try: | |
| self._normals_pipe.set_progress_bar_config(disable=True) | |
| except Exception: | |
| pass | |
| def _run_normals(self, pil_img: Image.Image) -> Dict[str, Any]: | |
| try: | |
| self._ensure_normals_loaded() | |
| # Do NOT pass processing_res to avoid the earlier error. | |
| # Most versions work with just the image (and optionally steps if supported). | |
| out = self._normals_pipe(pil_img) | |
| # Try common fields in order | |
| pred = None | |
| for key in ("prediction", "pred_normals"): | |
| if hasattr(out, key): | |
| pred = getattr(out, key) | |
| break | |
| if pred is None and hasattr(out, "images") and out.images: | |
| # Some pipelines return PIL in images[0] | |
| pred = out.images[0] | |
| # Convert to info + small preview b64 | |
| if isinstance(pred, torch.Tensor): | |
| arr = pred.detach().cpu().numpy() | |
| shape = list(arr.shape) | |
| dtype_str = str(arr.dtype) | |
| png_b64 = _to_uint8_png_b64(arr) | |
| elif isinstance(pred, np.ndarray): | |
| shape = list(pred.shape) | |
| dtype_str = str(pred.dtype) | |
| png_b64 = _to_uint8_png_b64(pred) | |
| elif isinstance(pred, Image.Image): | |
| shape = [pred.height, pred.width, 3] | |
| dtype_str = "uint8" | |
| png_b64 = _to_uint8_png_b64(pred) | |
| else: | |
| return {"error": f"Unexpected prediction type: {type(pred)}"} | |
| if png_b64 is None: | |
| return {"error": "Too many dimensions: 4 > 3."} | |
| return { | |
| "tensors_info": {"shape": shape, "dtype": dtype_str}, | |
| "preview_png_b64": png_b64, | |
| } | |
| except torch.cuda.OutOfMemoryError: | |
| return {"error": "GPU out of memory. Try a smaller image."} | |
| except Exception as e: | |
| return {"error": f"Normals prediction failed: {str(e)}"} | |
| # ---------------- Intrinsics (unchanged behavior) ---------------- | |
| def _run_intrinsics(self, pil_img: Image.Image, num_steps: int) -> Dict[str, Any]: | |
| # Inference (as before) | |
| with torch.inference_mode(): | |
| out = self.pipe(pil_img, num_inference_steps=num_steps) | |
| pred = out.prediction | |
| # Support torch.Tensor OR numpy.ndarray | |
| if isinstance(pred, torch.Tensor): | |
| pred_np = pred.detach().cpu().numpy() | |
| elif isinstance(pred, np.ndarray): | |
| pred_np = pred | |
| else: | |
| return {"error": f"Unexpected prediction type: {type(pred)}"} | |
| # pred_np expected shape: (T, H, W, 3) or (H, W, 3) depending on pipeline | |
| shape = list(pred_np.shape) | |
| dtype_str = str(pred_np.dtype) | |
| # Visualizations (best-effort) — unchanged | |
| vis = {} | |
| vis_err = None | |
| try: | |
| # visualize_intrinsics expects numpy prediction and the pipeline's target_properties | |
| vis_list = self.pipe.image_processor.visualize_intrinsics(pred_np, self.pipe.target_properties) | |
| # vis_list is a list over batch; we run single image, so index 0 | |
| if vis_list and isinstance(vis_list[0], dict): | |
| for name, im in vis_list[0].items(): | |
| if isinstance(im, Image.Image): | |
| buf = io.BytesIO() | |
| im.save(buf, format="PNG") | |
| vis[name] = base64.b64encode(buf.getvalue()).decode("utf-8") | |
| except Exception as e: | |
| vis_err = f"Failed to visualize: {e}" | |
| return { | |
| "targets": self.targets, | |
| "tensors_info": {"shape": shape, "dtype": dtype_str}, | |
| "uncertainty_available": False, | |
| "visualizations_png_b64": vis or None, | |
| "visualizations_error": vis_err, | |
| } | |
| # ---------------- Entry point ---------------- | |
| def __call__(self, data: Dict[str, Any]) -> Dict[str, Any]: | |
| try: | |
| if "inputs" not in data: | |
| return {"error": "Body must include an 'inputs' key."} | |
| # Parameters can be at top-level or under 'parameters' | |
| params = {} | |
| params.update({k: v for k, v in data.items() if k != "inputs"}) | |
| if isinstance(data.get("parameters"), dict): | |
| params.update(data["parameters"]) | |
| # Which outputs to return? | |
| ret = params.get("return", ["intrinsics"]) | |
| if isinstance(ret, str): | |
| ret = [ret] | |
| ret = [r.lower() for r in ret] | |
| # Steps (intrinsics) | |
| try: | |
| num_steps = int(params.get("num_inference_steps", self.default_intr_steps)) | |
| except Exception: | |
| num_steps = self.default_intr_steps | |
| # Decode image | |
| try: | |
| pil_img = _image_from_inputs(data["inputs"]) | |
| except Exception as e: | |
| return {"error": f"Failed to decode 'inputs': {e}"} | |
| result: Dict[str, Any] = {"outputs": {}} | |
| if "normals" in ret: | |
| result["outputs"]["normals"] = self._run_normals(pil_img) | |
| if "intrinsics" in ret or not ret: | |
| intr_res = self._run_intrinsics(pil_img, num_steps=num_steps) | |
| # sanity check that vis contains expected keys when available | |
| if isinstance(intr_res, dict) and "targets" in intr_res: | |
| # Keep your original “intrinsics-only” structure nested | |
| result["outputs"]["intrinsics"] = intr_res | |
| else: | |
| result["outputs"]["intrinsics"] = { | |
| "error": "Intrinsics pipeline did not return expected outputs (albedo, diffuse_shading, residual)." | |
| } | |
| # If only one branch requested, return only that branch in 'outputs' | |
| return result | |
| except Exception as e: | |
| return {"error": f"{e}"} | |
| finally: | |
| if torch.cuda.is_available(): | |
| torch.cuda.empty_cache() | |