server/model_loader.py

import os

_segmenter = None
_inpainter = None


def get_segmenter():
    global _segmenter
    if _segmenter is None:
        from server.segmentation import Sam2ModelSegmenter

        _segmenter = Sam2ModelSegmenter(
            hf_model_id=os.environ.get("SAM2_HF_MODEL_ID", "facebook/sam2.1-hiera-base-plus"),
            device=os.environ.get("INFERENCE_DEVICE", "cuda"),
        )
    return _segmenter


def get_inpainter():
    global _inpainter
    if _inpainter is None:
        device = os.environ.get("INFERENCE_DEVICE", "cuda")
        backend = os.environ.get("INPAINT_BACKEND", "powerpaint")  # "powerpaint", "sd", "lama", or "remote"

        if backend == "lama":
            from server.inpainting import LamaModelInpainter

            _inpainter = LamaModelInpainter(device=device)
        elif backend == "remote":
            from server.inpainting import RemoteInpainter

            _inpainter = RemoteInpainter(
                base_url=os.environ["REMOTE_INPAINT_URL"],
                timeout_seconds=float(os.environ.get("REMOTE_INPAINT_TIMEOUT_SECONDS", "300")),
            )
        elif backend == "powerpaint":
            from server.inpainting import PowerPaintInpainter

            _inpainter = PowerPaintInpainter(
                device=device,
                num_inference_steps=int(os.environ.get("SD_NUM_INFERENCE_STEPS", "45")),
                guidance_scale=float(os.environ.get("POWERPAINT_GUIDANCE_SCALE", "10.0")),
            )
        else:
            from server.inpainting import StableDiffusionInpainter

            working_size_str = os.environ.get("SD_WORKING_SIZE", "1024x1024")
            working_width, working_height = (int(v) for v in working_size_str.split("x"))
            _inpainter = StableDiffusionInpainter(
                device=device,
                model_id=os.environ.get(
                    "SD_INPAINT_MODEL_ID", "diffusers/stable-diffusion-xl-1.0-inpainting-0.1"
                ),
                working_size=(working_width, working_height),
                step_pause_seconds=float(os.environ.get("SD_STEP_PAUSE_SECONDS", "0")),
                num_inference_steps=int(os.environ.get("SD_NUM_INFERENCE_STEPS", "30")),
            )
    return _inpainter