File size: 3,731 Bytes
576ccef c302e80 576ccef c302e80 576ccef | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 | from __future__ import annotations
import os
import threading
from dataclasses import dataclass
import numpy as np
from PIL import Image
try:
import spaces
gpu_task = spaces.GPU(duration=180)
except ImportError:
def gpu_task(function):
return function
MODEL_ID = os.getenv("MODEL_ID", "ZhengPeng7/BiRefNet")
MODEL_REVISION = os.getenv(
"MODEL_REVISION", "e2bf8e4460fc8fa32bba5ea4d94b3233d367b0e4"
)
MODEL_INPUT_SIZE = int(os.getenv("MODEL_INPUT_SIZE", "1024"))
_runtime: "BiRefNetRuntime | None" = None
_load_lock = threading.Lock()
@dataclass
class BiRefNetRuntime:
model: object
device: object
dtype: object
@classmethod
def load(cls) -> "BiRefNetRuntime":
import torch
from transformers import AutoModelForImageSegmentation
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
dtype = torch.float16 if device.type == "cuda" else torch.float32
torch.set_float32_matmul_precision("high")
model = AutoModelForImageSegmentation.from_pretrained(
MODEL_ID,
revision=MODEL_REVISION,
trust_remote_code=True,
)
model.to(device=device, dtype=dtype)
model.eval()
return cls(model=model, device=device, dtype=dtype)
def predict(self, image: Image.Image) -> Image.Image:
import torch
from torchvision.transforms import functional as TF
boxed, content_box = _letterbox(image, MODEL_INPUT_SIZE)
tensor = TF.to_tensor(boxed)
tensor = TF.normalize(
tensor,
mean=(0.485, 0.456, 0.406),
std=(0.229, 0.224, 0.225),
).unsqueeze(0)
tensor = tensor.to(device=self.device, dtype=self.dtype)
with torch.inference_mode():
prediction = self.model(tensor)
logits = _last_tensor(prediction)
probability = logits.sigmoid()[0].squeeze().float().cpu().numpy()
probability = np.clip(probability * 255.0, 0, 255).astype(np.uint8)
square_mask = Image.fromarray(probability, mode="L")
content_mask = square_mask.crop(content_box)
return content_mask.resize(image.size, Image.Resampling.LANCZOS)
def _letterbox(image: Image.Image, size: int) -> tuple[Image.Image, tuple[int, int, int, int]]:
scale = min(size / image.width, size / image.height)
resized_size = (max(1, round(image.width * scale)), max(1, round(image.height * scale)))
resized = image.resize(resized_size, Image.Resampling.LANCZOS)
# ImageNet mean becomes approximately zero after normalization.
canvas = Image.new("RGB", (size, size), (124, 116, 104))
left = (size - resized.width) // 2
top = (size - resized.height) // 2
canvas.paste(resized, (left, top))
return canvas, (left, top, left + resized.width, top + resized.height)
def _last_tensor(value):
import torch
if torch.is_tensor(value):
return value
if hasattr(value, "logits"):
return _last_tensor(value.logits)
if isinstance(value, (tuple, list)) and value:
for item in reversed(value):
try:
return _last_tensor(item)
except (TypeError, ValueError):
continue
raise TypeError("BiRefNet beklenmeyen bir çıktı biçimi döndürdü.")
def get_runtime() -> BiRefNetRuntime:
global _runtime
if _runtime is None:
with _load_lock:
if _runtime is None:
_runtime = BiRefNetRuntime.load()
return _runtime
@gpu_task
def predict_mask(image: Image.Image) -> Image.Image:
"""Run segmentation on ZeroGPU/classical GPU, with transparent CPU fallback."""
return get_runtime().predict(image)
|