Spaces:
Sleeping
Sleeping
File size: 2,202 Bytes
ca4beaa 479a2b5 ca4beaa 838d140 479a2b5 e9c3e3f 479a2b5 838d140 ca4beaa 838d140 | 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 | from transformers import AutoProcessor
from ultralytics import SAM
import openvino as ov
import config as config
from optimum.intel import OVStableDiffusionXLInpaintPipeline
from simple_lama_inpainting import SimpleLama
def get_detector_model(MODELS: dict):
if MODELS["detector"] is None:
core = ov.Core()
model = core.read_model(config.DETECTOR_MODEL_PATH)
MODELS["detector"] = core.compile_model(model, "CPU")
return MODELS["detector"]
def get_detector_processor(MODELS: dict):
if MODELS["detector_processor"] is None:
MODELS["detector_processor"] = AutoProcessor.from_pretrained(config.DETECTOR_MODEL_NAME, use_fast=True)
return MODELS["detector_processor"]
def get_segmentor_model(MODELS: dict):
if MODELS["segmentor"] is None:
MODELS["segmentor"] = SAM(config.SEGMENTOR_MODEL_NAME)
return MODELS["segmentor"]
def get_inpaintor_model(MODELS: dict):
if MODELS["inpaintor"] is None:
model_source = config.INPAINTOR_MODEL_PATH if config.INPAINTOR_MODEL_PATH.exists() else config.INPAINTOR_MODEL_NAME
kwargs = {"device": config.INPAINTOR_DEVICE}
if model_source == config.INPAINTOR_MODEL_NAME:
kwargs["export"] = True
pipe = OVStableDiffusionXLInpaintPipeline.from_pretrained(
str(model_source),
**kwargs,
)
# ❌ REMOVE OR COMMENT OUT THIS ENTIRE BLOCK ❌
# pipe.reshape(
# batch_size=1,
# height=config.INPAINTOR_IMAGE_SIZE,
# width=config.INPAINTOR_IMAGE_SIZE,
# num_images_per_prompt=1,
# )
# Keep compile!
pipe.compile()
MODELS["inpaintor"] = pipe
return MODELS["inpaintor"]
def get_removing_model(MODELS: dict):
if MODELS["remover"] is None:
import torch
# Intercept the load function to force CPU mapping
original_load = torch.jit.load
torch.jit.load = lambda *a, **kw: original_load(*a, **{**kw, "map_location": "cpu"})
MODELS["remover"] = SimpleLama()
# Restore the original PyTorch load function right after
torch.jit.load = original_load
return MODELS["remover"]
|