diff --git a/bioimageio_colab/register_sam_service.py b/bioimageio_colab/register_sam_service.py index d598cde..3ff65eb 100644 --- a/bioimageio_colab/register_sam_service.py +++ b/bioimageio_colab/register_sam_service.py @@ -153,29 +153,32 @@ async def compute_image_embedding( """ Compute the embeddings of an image using the specified model. """ - async with semaphore: - try: - user = context["user"] - if require_login and user["is_anonymous"]: - raise PermissionError("You must be logged in to use this service.") - - user_id = user["id"] + try: + user = context["user"] + if require_login and user["is_anonymous"]: + raise PermissionError("You must be logged in to use this service.") + user_id = user["id"] + + # Put image immediately into the object store to avoid memory issues + logger.info(f"User '{user_id}' - Putting image into the object store...") + obj_ref = ray.put(image) + del image + + # Compute the embedding, but limit the number of concurrent requests + async with semaphore: logger.info( f"User '{user_id}' - Computing embedding (model: '{model_id}')..." ) - - # Compute the embedding - # Returns: {"features": embedding, "input_size": input_size} + + # Format: {"features": embedding, "input_size": input_size} result = await app_handle.options(multiplexed_model_id=model_id).remote( - image + obj_ref ) - logger.info(f"User '{user_id}' - Embedding computed successfully.") - return result - except Exception as e: - logger.error(f"User '{user_id}' - Error computing embedding: {e}") - raise e + except Exception as e: + logger.error(f"User '{user_id}' - Error computing embedding: {e}") + raise e # def compute_mask( diff --git a/pyproject.toml b/pyproject.toml index 4bac6b4..21e40ec 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -3,7 +3,7 @@ requires = ["setuptools", "wheel"] [project] name = "bioimageio-colab" -version = "0.2.6" +version = "0.2.7" readme = "README.md" description = "Collaborative image annotation and model training with human in the loop." dependencies = [