supervisely 6.73.268__py3-none-any.whl → 6.73.270__py3-none-any.whl

supervisely/nn/inference/inference.py

@@ -1,3 +1,5 @@
+import argparse
+import asyncio
 import inspect
 import json
 import os
@@ -11,12 +13,14 @@ from collections import OrderedDict, defaultdict
 from concurrent.futures import ThreadPoolExecutor
 from dataclasses import asdict
 from functools import partial, wraps
+from pathlib import Path
 from queue import Queue
 from typing import Any, Callable, Dict, List, Optional, Tuple, Union
 from urllib.request import urlopen
 
 import numpy as np
 import requests
+import uvicorn
 import yaml
 from fastapi import Form, HTTPException, Request, Response, UploadFile, status
 from fastapi.responses import JSONResponse
@@ -24,8 +28,7 @@ from requests.structures import CaseInsensitiveDict
 
 import supervisely.app.development as sly_app_development
 import supervisely.imaging.image as sly_image
-import supervisely.io.env as env
-import supervisely.io.fs as fs
+import supervisely.io.env as sly_env
 import supervisely.io.fs as sly_fs
 import supervisely.io.json as sly_json
 import supervisely.nn.inference.gui as GUI
@@ -33,6 +36,7 @@ from supervisely import DatasetInfo, ProjectInfo, VideoAnnotation, batched
 from supervisely._utils import (
     add_callback,
     get_filename_from_headers,
+    get_or_create_event_loop,
     is_debug_with_sly_net,
     is_production,
     rand_str,
@@ -43,6 +47,7 @@ from supervisely.annotation.obj_class import ObjClass
 from supervisely.annotation.tag_collection import TagCollection
 from supervisely.annotation.tag_meta import TagMeta, TagValueType
 from supervisely.api.api import Api
+from supervisely.api.app_api import WorkflowMeta, WorkflowSettings
 from supervisely.api.image_api import ImageInfo
 from supervisely.app.content import StateJson, get_data_dir
 from supervisely.app.exceptions import DialogWindowError
@@ -106,9 +111,19 @@ class Inference:
         multithread_inference: Optional[bool] = True,
         use_serving_gui_template: Optional[bool] = False,
     ):
+
+        self._args, self._is_local_deploy = self._parse_local_deploy_args()
+
         if model_dir is None:
-            model_dir = os.path.join(get_data_dir(), "models")
-            fs.mkdir(model_dir)
+            if self._is_local_deploy is True:
+                try:
+                    model_dir = get_data_dir()
+                except:
+                    model_dir = Path("~/.cache/supervisely/app_data").expanduser()
+            else:
+                model_dir = os.path.join(get_data_dir(), "models")
+        sly_fs.mkdir(model_dir)
+
         self.device: str = None
         self.runtime: str = None
         self.model_precision: str = None
@@ -128,7 +143,7 @@ class Inference:
         self._autostart_delay_time = 5 * 60  # 5 min
         self._tracker = None
         self._hardware: str = None
-        self.pretrained_models = self._load_models_json(self.MODELS) if self.MODELS else None
+        self.pretrained_models = self._load_models_json_file(self.MODELS) if self.MODELS else None
         if custom_inference_settings is None:
             if self.INFERENCE_SETTINGS is not None:
                 custom_inference_settings = self.INFERENCE_SETTINGS
@@ -136,7 +151,7 @@ class Inference:
                 logger.debug("Custom inference settings are not provided.")
                 custom_inference_settings = {}
         if isinstance(custom_inference_settings, str):
-            if fs.file_exists(custom_inference_settings):
+            if sly_fs.file_exists(custom_inference_settings):
                 with open(custom_inference_settings, "r") as f:
                     custom_inference_settings = f.read()
             else:
@@ -146,12 +161,24 @@ class Inference:
         self._use_gui = use_gui
         self._use_serving_gui_template = use_serving_gui_template
         self._gui = None
+        self._uvicorn_server = None
 
         self.load_on_device = LOAD_ON_DEVICE_DECORATOR(self.load_on_device)
         self.load_on_device = add_callback(self.load_on_device, self._set_served_callback)
 
         self.load_model = LOAD_MODEL_DECORATOR(self.load_model)
 
+        if self._is_local_deploy:
+            # self._args = self._parse_local_deploy_args()
+            self._use_gui = False
+            deploy_params, need_download = self._get_deploy_params_from_args()
+            if need_download:
+                local_model_files = self._download_model_files(
+                    deploy_params["model_source"], deploy_params["model_files"], False
+                )
+                deploy_params["model_files"] = local_model_files
+            self._load_model_headless(**deploy_params)
+
         if self._use_gui:
             initialize_custom_gui_method = getattr(self, "initialize_custom_gui", None)
             original_initialize_custom_gui_method = getattr(
@@ -224,10 +251,10 @@ class Inference:
         self.get_info = self._check_serve_before_call(self.get_info)
 
         self.cache = InferenceImageCache(
-            maxsize=env.smart_cache_size(),
-            ttl=env.smart_cache_ttl(),
+            maxsize=sly_env.smart_cache_size(),
+            ttl=sly_env.smart_cache_ttl(),
             is_persistent=True,
-            base_folder=env.smart_cache_container_dir(),
+            base_folder=sly_env.smart_cache_container_dir(),
             log_progress=True,
         )
 
@@ -248,9 +275,18 @@ class Inference:
                 )
                 device = "cpu"
 
-    def _load_models_json(self, models: str) -> List[Dict[str, Any]]:
+    def _load_json_file(self, file_path: str) -> dict:
+        if isinstance(file_path, str):
+            if sly_fs.file_exists(file_path) and sly_fs.get_file_ext(file_path) == ".json":
+                return sly_json.load_json_file(file_path)
+            else:
+                raise ValueError("File not found or invalid file format.")
+        else:
+            raise ValueError("Invalid file. Please provide a valid '.json' file.")
+
+    def _load_models_json_file(self, models: str) -> List[Dict[str, Any]]:
         """
-        Loads models from the provided file or list of model configurations.
+        Loads dictionary from the provided file.
         """
         if isinstance(models, str):
             if sly_fs.file_exists(models) and sly_fs.get_file_ext(models) == ".json":
@@ -258,9 +294,7 @@ class Inference:
             else:
                 raise ValueError("File not found or invalid file format.")
         else:
-            raise ValueError(
-                "Invalid models file. Please provide a valid '.json' file with list of model configurations."
-            )
+            raise ValueError("Invalid file. Please provide a valid '.json' file.")
 
         if not isinstance(models, list):
             raise ValueError("models parameters must be a list of dicts")
@@ -394,13 +428,13 @@ class Inference:
         else:
             progress = None
 
-        if fs.dir_exists(src_path) or fs.file_exists(
+        if sly_fs.dir_exists(src_path) or sly_fs.file_exists(
             src_path
         ):  # only during debug, has no effect in production
             dst_path = os.path.abspath(src_path)
             logger.info(f"File {dst_path} found.")
         elif src_path.startswith("/"):  # folder from Team Files
-            team_id = env.team_id()
+            team_id = sly_env.team_id()
 
             if src_path.endswith("/") and self.api.file.dir_exists(team_id, src_path):
 
@@ -436,7 +470,7 @@ class Inference:
                 def download_file(team_id, src_path, dst_path, progress_cb=None):
                     self.api.file.download(team_id, src_path, dst_path, progress_cb=progress_cb)
 
-                file_info = self.api.file.get_info_by_path(env.team_id(), src_path)
+                file_info = self.api.file.get_info_by_path(sly_env.team_id(), src_path)
                 if progress is None:
                     download_file(team_id, src_path, dst_path)
                 else:
@@ -451,8 +485,8 @@ class Inference:
                 logger.info(f"📥 File {basename} has been successfully downloaded from Team Files")
                 logger.info(f"File {basename} path: {dst_path}")
         else:  # external url
-            if not fs.dir_exists(os.path.dirname(dst_path)):
-                fs.mkdir(os.path.dirname(dst_path))
+            if not sly_fs.dir_exists(os.path.dirname(dst_path)):
+                sly_fs.mkdir(os.path.dirname(dst_path))
 
             def download_external_file(url, save_path, progress=None):
                 def download_content(save_path, progress_cb=None):
@@ -531,13 +565,15 @@ class Inference:
     def _checkpoints_cache_dir(self):
         return os.path.join(os.path.expanduser("~"), ".cache", "supervisely", "checkpoints")
 
-    def _download_model_files(self, model_source: str, model_files: List[str]) -> dict:
+    def _download_model_files(
+        self, model_source: str, model_files: List[str], log_progress: bool = True
+    ) -> dict:
         if model_source == ModelSource.PRETRAINED:
-            return self._download_pretrained_model(model_files)
+            return self._download_pretrained_model(model_files, log_progress)
         elif model_source == ModelSource.CUSTOM:
-            return self._download_custom_model(model_files)
+            return self._download_custom_model(model_files, log_progress)
 
-    def _download_pretrained_model(self, model_files: dict):
+    def _download_pretrained_model(self, model_files: dict, log_progress: bool = True):
         """
         Downloads the pretrained model data.
         """
@@ -564,54 +600,61 @@ class Inference:
                         logger.debug(f"Model: '{file_name}' was found in model dir")
                         continue
 
-                    with self.gui.download_progress(
-                        message=f"Downloading: '{file_name}'",
-                        total=file_size,
-                        unit="bytes",
-                        unit_scale=True,
-                    ) as download_pbar:
-                        self.gui.download_progress.show()
-                        sly_fs.download(
-                            url=file_url,
-                            save_path=file_path,
-                            progress=download_pbar.update,
-                        )
+                    if log_progress:
+                        with self.gui.download_progress(
+                            message=f"Downloading: '{file_name}'",
+                            total=file_size,
+                            unit="bytes",
+                            unit_scale=True,
+                        ) as download_pbar:
+                            self.gui.download_progress.show()
+                            sly_fs.download(
+                                url=file_url, save_path=file_path, progress=download_pbar.update
+                            )
+                    else:
+                        sly_fs.download(url=file_url, save_path=file_path)
                     local_model_files[file] = file_path
             else:
                 local_model_files[file] = file_url
-        self.gui.download_progress.hide()
+
+        if log_progress:
+            self.gui.download_progress.hide()
         return local_model_files
 
-    def _download_custom_model(self, model_files: dict):
+    def _download_custom_model(self, model_files: dict, log_progress: bool = True):
         """
         Downloads the custom model data.
         """
-
-        team_id = env.team_id()
+        team_id = sly_env.team_id()
         local_model_files = {}
-
         for file in model_files:
             file_url = model_files[file]
             file_info = self.api.file.get_info_by_path(team_id, file_url)
             if file_info is None:
-                raise FileNotFoundError(
-                    f"File '{file_url}' not found in Team Files. Make sure the file exists."
-                )
+                if sly_fs.file_exists(file_url):
+                    local_model_files[file] = file_url
+                    continue
+                else:
+                    raise FileNotFoundError(f"File '{file_url}' not found in Team Files")
             file_size = file_info.sizeb
             file_name = os.path.basename(file_url)
             file_path = os.path.join(self.model_dir, file_name)
-            with self.gui.download_progress(
-                message=f"Downloading: '{file_name}'",
-                total=file_size,
-                unit="bytes",
-                unit_scale=True,
-            ) as download_pbar:
-                self.gui.download_progress.show()
-                self.api.file.download(
-                    team_id, file_url, file_path, progress_cb=download_pbar.update
-                )
+            if log_progress:
+                with self.gui.download_progress(
+                    message=f"Downloading: '{file_name}'",
+                    total=file_size,
+                    unit="bytes",
+                    unit_scale=True,
+                ) as download_pbar:
+                    self.gui.download_progress.show()
+                    self.api.file.download(
+                        team_id, file_url, file_path, progress_cb=download_pbar.update
+                    )
+            else:
+                self.api.file.download(team_id, file_url, file_path)
             local_model_files[file] = file_path
-        self.gui.download_progress.hide()
+        if log_progress:
+            self.gui.download_progress.hide()
         return local_model_files
 
     def _load_model(self, deploy_params: dict):
@@ -647,6 +690,13 @@ class Inference:
         if model_source == ModelSource.CUSTOM:
             self._set_model_meta_custom_model(model_info)
             self._set_checkpoint_info_custom_model(deploy_params)
+
+        try:
+            if is_production():
+                self._add_workflow_input(model_source, model_files, model_info)
+        except Exception as e:
+            logger.warning(f"Failed to add input to the workflow: {repr(e)}")
+
         self._load_model(deploy_params)
         if self._model_meta is None:
             self._set_model_meta_from_classes()
@@ -679,7 +729,7 @@ class Inference:
                 model_info.get("artifacts_dir"), "checkpoints", checkpoint_name
             )
             checkpoint_file_info = self.api.file.get_info_by_path(
-                env.team_id(), checkpoint_file_path
+                sly_env.team_id(), checkpoint_file_path
             )
             if checkpoint_file_info is None:
                 checkpoint_url = None
@@ -1253,18 +1303,18 @@ class Inference:
         logger.debug("Inferring image_url...", extra={"state": state})
         settings = self._get_inference_settings(state)
         image_url = state["image_url"]
-        ext = fs.get_file_ext(image_url)
+        ext = sly_fs.get_file_ext(image_url)
         if ext == "":
             ext = ".jpg"
         image_path = os.path.join(get_data_dir(), rand_str(15) + ext)
-        fs.download(image_url, image_path)
+        sly_fs.download(image_url, image_path)
         logger.debug("Inference settings:", extra=settings)
         logger.debug(f"Downloaded path: {image_path}")
         anns, slides_data = self._inference_auto(
             [image_path],
             settings=settings,
         )
-        fs.silent_remove(image_path)
+        sly_fs.silent_remove(image_path)
         return self._format_output(anns, slides_data)[0]
 
     def _inference_video_id(self, api: Api, state: dict, async_inference_request_uuid: str = None):
@@ -2208,28 +2258,46 @@ class Inference:
 
         if is_debug_with_sly_net():
             # advanced debug for Supervisely Team
-            logger.warn(
+            logger.warning(
                 "Serving is running in advanced development mode with Supervisely VPN Network"
             )
-            team_id = env.team_id()
+            team_id = sly_env.team_id()
             # sly_app_development.supervisely_vpn_network(action="down") # for debug
             sly_app_development.supervisely_vpn_network(action="up")
             task = sly_app_development.create_debug_task(team_id, port="8000")
             self._task_id = task["id"]
             os.environ["TASK_ID"] = str(self._task_id)
         else:
-            self._task_id = env.task_id() if is_production() else None
+            if not self._is_local_deploy:
+                self._task_id = sly_env.task_id() if is_production() else None
 
-        if isinstance(self.gui, GUI.InferenceGUI):
+        if isinstance(self.gui, GUI.ServingGUITemplate):
             self._app = Application(layout=self.get_ui())
         elif isinstance(self.gui, GUI.ServingGUI):
             self._app = Application(layout=self._app_layout)
+        # elif isinstance(self.gui, GUI.InferenceGUI):
+        #     self._app = Application(layout=self.get_ui())
         else:
             self._app = Application(layout=self.get_ui())
 
         server = self._app.get_server()
         self._app.set_ready_check_function(self.is_model_deployed)
 
+        if self._is_local_deploy:
+            # Predict and shutdown
+            if any(
+                [
+                    self._args.predict_project,
+                    self._args.predict_dataset,
+                    self._args.predict_dir,
+                    self._args.predict_image,
+                ]
+            ):
+                self._inference_by_local_deploy_args()
+                # Gracefully shut down the server
+                self._app.shutdown()
+        # else: run server after endpoints
+
         @call_on_autostart()
         def autostart_func():
             gpu_count = get_gpu_count()
@@ -2725,6 +2793,285 @@ class Inference:
         def _get_deploy_info():
             return asdict(self._get_deploy_info())
 
+        # Local deploy without predict args
+        if self._is_local_deploy:
+            self._run_server()
+
+    def _parse_local_deploy_args(self):
+        parser = argparse.ArgumentParser(description="Run Inference Serving")
+
+        # Deploy args
+        parser.add_argument(
+            "--model",
+            type=str,
+            help="Name of the pretrained model or path to custom checkpoint file",
+        )
+        parser.add_argument(
+            "--device",
+            type=str,
+            choices=["cpu", "cuda", "cuda:0", "cuda:1", "cuda:2", "cuda:3"],
+            default="cuda:0",
+            help="Device to use for inference (default: 'cuda:0')",
+        )
+        parser.add_argument(
+            "--runtime",
+            type=str,
+            choices=[RuntimeType.PYTORCH, RuntimeType.ONNXRUNTIME, RuntimeType.TENSORRT],
+            default=RuntimeType.PYTORCH,
+            help="Runtime type for inference (default: PYTORCH)",
+        )
+        # -------------------------- #
+
+        # Predict args
+        parser.add_argument("--predict-project", type=int, required=False, help="ID of the project")
+        parser.add_argument(
+            "--predict-dataset",
+            type=lambda x: [int(i) for i in x.split(",")] if "," in x else int(x),
+            required=False,
+            help="ID of the dataset or a comma-separated list of dataset IDs",
+        )
+        parser.add_argument(
+            "--predict-dir",
+            type=str,
+            required=False,
+            help="Not implemented yet. Path to the local directory with images",
+        )
+        parser.add_argument(
+            "--predict-image",
+            type=str,
+            required=False,
+            help="Image ID on Supervisely instance or path to local image",
+        )
+        # -------------------------- #
+
+        # Output args
+        parser.add_argument("--output", type=str, required=False, help="Not implemented yet")
+        parser.add_argument("--output-dir", type=str, required=False, help="Not implemented yet")
+        # -------------------------- #
+
+        # Parse arguments
+        args, _ = parser.parse_known_args()
+        if args.model is None:
+            # raise ValueError("Argument '--model' is required for local deployment")
+            return None, False
+        if isinstance(args.predict_dataset, int):
+            args.predict_dataset = [args.predict_dataset]
+        if args.predict_image is not None:
+            if args.predict_image.isdigit():
+                args.predict_image = int(args.predict_image)
+        return args, True
+
+    def _get_pretrained_model_params_from_args(self):
+        model_files = None
+        model_source = None
+        model_info = None
+        need_download = True
+
+        model = self._args.model
+        for m in self.pretrained_models:
+            meta = m.get("meta", None)
+            if meta is None:
+                continue
+            model_name = meta.get("model_name", None)
+            if model_name is None:
+                continue
+            m_files = meta.get("model_files", None)
+            if m_files is None:
+                continue
+            checkpoint = m_files.get("checkpoint", None)
+            if checkpoint is None:
+                continue
+            if model == m["meta"]["model_name"]:
+                model_info = m
+                model_source = ModelSource.PRETRAINED
+                model_files = {"checkpoint": checkpoint}
+                config = m_files.get("config", None)
+                if config is not None:
+                    model_files["config"] = config
+                break
+
+        return model_files, model_source, model_info, need_download
+
+    def _get_custom_model_params_from_args(self):
+        def _load_experiment_info(artifacts_dir):
+            experiment_path = os.path.join(artifacts_dir, "experiment_info.json")
+            model_info = self._load_json_file(experiment_path)
+            original_model_files = model_info.get("model_files")
+            if not original_model_files:
+                raise ValueError("Invalid 'experiment_info.json'. Missing 'model_files' key.")
+            return model_info, original_model_files
+
+        def _prepare_local_model_files(artifacts_dir, checkpoint_path, original_model_files):
+            return {k: os.path.join(artifacts_dir, v) for k, v in original_model_files.items()} | {
+                "checkpoint": checkpoint_path
+            }
+
+        def _download_remote_files(team_id, artifacts_dir, local_artifacts_dir):
+            sly_fs.mkdir(local_artifacts_dir, True)
+            file_infos = self.api.file.list(team_id, artifacts_dir, False, "fileinfo")
+            remote_paths = [f.path for f in file_infos if not f.is_dir]
+            local_paths = [
+                os.path.join(local_artifacts_dir, f.name) for f in file_infos if not f.is_dir
+            ]
+
+            coro = self.api.file.download_bulk_async(team_id, remote_paths, local_paths)
+            loop = get_or_create_event_loop()
+            if loop.is_running():
+                future = asyncio.run_coroutine_threadsafe(coro, loop)
+                future.result()
+            else:
+                loop.run_until_complete(coro)
+
+        model_source = ModelSource.CUSTOM
+        need_download = False
+        checkpoint_path = self._args.model
+
+        if not os.path.isfile(checkpoint_path):
+            team_id = sly_env.team_id(raise_not_found=False)
+            if not team_id:
+                raise ValueError(
+                    "Team ID not found in env. Required for remote custom checkpoints."
+                )
+            file_info = self.api.file.get_info_by_path(team_id, checkpoint_path)
+            if not file_info:
+                raise ValueError(
+                    f"Couldn't find: '{checkpoint_path}' locally or remotely in Team ID."
+                )
+            need_download = True
+
+        artifacts_dir = os.path.dirname(os.path.dirname(checkpoint_path))
+        if not need_download:
+            model_info, original_model_files = _load_experiment_info(artifacts_dir)
+            model_files = _prepare_local_model_files(
+                artifacts_dir, checkpoint_path, original_model_files
+            )
+        else:
+            local_artifacts_dir = os.path.join(
+                self.model_dir, "local_deploy", os.path.basename(artifacts_dir)
+            )
+            _download_remote_files(team_id, artifacts_dir, local_artifacts_dir)
+
+            model_info, original_model_files = _load_experiment_info(local_artifacts_dir)
+            model_files = _prepare_local_model_files(
+                local_artifacts_dir, checkpoint_path, original_model_files
+            )
+        return model_files, model_source, model_info, need_download
+
+    def _get_deploy_params_from_args(self):
+        # Ensure model directory exists
+        device = self._args.device if self._args.device else "cuda:0"
+        runtime = self._args.runtime if self._args.runtime else RuntimeType.PYTORCH
+
+        model_files, model_source, model_info, need_download = (
+            self._get_pretrained_model_params_from_args()
+        )
+        if model_source is None:
+            model_files, model_source, model_info, need_download = (
+                self._get_custom_model_params_from_args()
+            )
+
+        if model_source is None:
+            raise ValueError("Couldn't create 'model_source' from args")
+        if model_files is None:
+            raise ValueError("Couldn't create 'model_files' from args")
+        if model_info is None:
+            raise ValueError("Couldn't create 'model_info' from args")
+
+        deploy_params = {
+            "model_files": model_files,
+            "model_source": model_source,
+            "model_info": model_info,
+            "device": device,
+            "runtime": runtime,
+        }
+
+        logger.info(f"Deploy parameters: {deploy_params}")
+        return deploy_params, need_download
+
+    def _run_server(self):
+        config = uvicorn.Config(app=self._app, host="0.0.0.0", port=8000, ws="websockets")
+        self._uvicorn_server = uvicorn.Server(config)
+        self._uvicorn_server.run()
+
+    def _inference_by_local_deploy_args(self):
+        def predict_project_by_args(api: Api, project_id: int, dataset_ids: List[int] = None):
+            source_project = api.project.get_info_by_id(project_id)
+            workspace_id = source_project.workspace_id
+            output_project = api.project.create(
+                workspace_id, f"{source_project.name} predicted", change_name_if_conflict=True
+            )
+            results = self._inference_project_id(
+                api=self.api,
+                state={
+                    "projectId": project_id,
+                    "dataset_ids": dataset_ids,
+                    "output_project_id": output_project.id,
+                },
+            )
+
+        def predict_datasets_by_args(api: Api, dataset_ids: List[int]):
+            dataset_infos = [api.dataset.get_info_by_id(dataset_id) for dataset_id in dataset_ids]
+            project_ids = list(set([dataset_info.project_id for dataset_info in dataset_infos]))
+            if len(project_ids) > 1:
+                raise ValueError("All datasets should belong to the same project")
+            predict_project_by_args(api, project_ids[0], dataset_ids)
+
+        def predict_image_by_args(api: Api, image: Union[str, int]):
+            def predict_image_np(image_np):
+                settings = self._get_inference_settings({})
+                anns, _ = self._inference_auto([image_np], settings)
+                if len(anns) == 0:
+                    return Annotation(img_size=image_np.shape[:2])
+                ann = anns[0]
+                return ann
+
+            if isinstance(image, int):
+                image_np = api.image.download_np(image)
+                ann = predict_image_np(image_np)
+                api.annotation.upload_ann(image, ann)
+            elif isinstance(image, str):
+                if sly_fs.file_exists(self._args.predict):
+                    image_np = sly_image.read(self._args.predict)
+                    ann = predict_image_np(image_np)
+                    pred_ann_path = image + ".json"
+                    sly_json.dump_json_file(ann.to_json(), pred_ann_path)
+                    # Save image for debug
+                    # ann.draw_pretty(image_np)
+                    # pred_path = os.path.join(os.path.dirname(self._args.predict), "pred_" + os.path.basename(self._args.predict))
+                    # sly_image.write(pred_path, image_np)
+
+        if self._args.predict_project is not None:
+            predict_project_by_args(self.api, self._args.predict_project)
+        elif self._args.predict_dataset is not None:
+            predict_datasets_by_args(self.api, self._args.predict_dataset)
+        elif self._args.predict_dir is not None:
+            raise NotImplementedError("Predict from directory is not implemented yet")
+        elif self._args.predict_image is not None:
+            predict_image_by_args(self.api, self._args.predict_image)
+
+    def _add_workflow_input(self, model_source: str, model_files: dict, model_info: dict):
+        if model_source == ModelSource.PRETRAINED:
+            checkpoint_url = model_info["meta"]["model_files"]["checkpoint"]
+            checkpoint_name = model_info["meta"]["model_name"]
+        else:
+            checkpoint_name = sly_fs.get_file_name_with_ext(model_files["checkpoint"])
+            checkpoint_url = os.path.join(
+                model_info["artifacts_dir"], "checkpoints", checkpoint_name
+            )
+
+        app_name = sly_env.app_name()
+        meta = WorkflowMeta(node_settings=WorkflowSettings(title=f"Serve {app_name}"))
+
+        logger.debug(
+            f"Workflow Input: Checkpoint URL - {checkpoint_url}, Checkpoint Name - {checkpoint_name}"
+        )
+        if checkpoint_url and self.api.file.exists(sly_env.team_id(), checkpoint_url):
+            self.api.app.workflow.add_input_file(checkpoint_url, model_weight=True, meta=meta)
+        else:
+            logger.debug(
+                f"Checkpoint {checkpoint_url} not found in Team Files. Cannot set workflow input"
+            )
+
 
 def _get_log_extra_for_inference_request(inference_request_uuid, inference_request: dict):
     log_extra = {
@@ -2998,7 +3345,7 @@ class TempImageWriter:
     def __init__(self, format: str = "png"):
         self.format = format
         self.temp_dir = os.path.join(get_data_dir(), rand_str(10))
-        fs.mkdir(self.temp_dir)
+        sly_fs.mkdir(self.temp_dir)
 
     def write(self, image: np.ndarray):
         image_path = os.path.join(self.temp_dir, f"{rand_str(10)}.{self.format}")
@@ -3006,7 +3353,7 @@ class TempImageWriter:
         return image_path
 
     def clean(self):
-        fs.remove_dir(self.temp_dir)
+        sly_fs.remove_dir(self.temp_dir)
 
 
 def get_hardware_info(device: str) -> str: