Package: tungstenkit
Tungstenkit: ML container made simpleIntroduction|Installation|Documentation|CommunityTungstenkitis ML containerization tool with a focus on developer productivity and versatility.Have you ever struggled to use models from github?
You may have repeated tedious steps like: cuda/dependency problems, file handling, and scripting for testing.Standing on the shoulder of Docker, this project aims to make using ML models less painful by adding functionalities for typical use cases - REST API server, GUI, CLI, and Python script.With Tungstenkit, sharing and consuming ML models can be quick and enjoyable.FeaturesRequires only a few lines of Python codeBuild once, use everywhere:REST API serverGUI applicationCLI applicationPython functionFramework-agnostic and lightweightPydantic input/output definitions with convenient file handlingSupports batched predictionSupports clustering with distributed machines (coming soon)Take the tourRequires only a few lines of python codeBuilding a Tungsten model is easy. All you have to do is write a simpletungsten_model.pylike:fromtypingimportListimporttorchfromtungstenkitimportBaseIO,Image,define_modelclassInput(BaseIO):prompt:strclassOutput(BaseIO):image:Image@define_model(input=Input,output=Output,gpu=True,python_packages=["torch","torchvision"],batch_size=4,gpu_mem_gb=16,)classTextToImageModel:defsetup(self):weights=torch.load("./weights.pth")self.model=load_torch_model(weights)defpredict(self,inputs:List[Input])->List[Output]:input_tensor=preprocess(inputs)output_tensor=self.model(input_tensor)outputs=postprocess(output_tensor)returnoutputsStart a build process:$tungstenbuild.-ntext-to-image✅ Successfully built tungsten model: 'text-to-image:e3a5de56'Check the built image:$ tungsten models

Repository        Tag       Create Time          Docker Image ID
----------------  --------  -------------------  ---------------
text-to-image     latest    2023-04-26 05:23:58  830eb82f0fcd
text-to-image     e3a5de56  2023-04-26 05:23:58  830eb82f0fcdBuild once, use everywhereREST API serverStart a server:$tungstenservetext-to-image-p3000INFO:     Uvicorn running on http://0.0.0.0:3000 (Press CTRL+C to quit)Send a prediction request with a JSON payload:$curl-X'POST''http://localhost:3000/predictions'\-H'Accept: application/json'\-H'Content-Type: application/json'\-d'[{"prompt": "a professional photograph of an astronaut riding a horse"}]'{"prediction_id": "39c9eb6b"}Get the result:$curl-X'GET''http://localhost:3000/predictions/39c9eb6b'\-H'Accept: application/json'{"outputs": [{"image": "data:image/png;base64,..."}],"status": "success"}GUI applicationIf you need a more user-friendly way to make predictions, start a GUI app with the following command:$tungstendemotext-to-image-p8080INFO:     Uvicorn running on http://localhost:8080 (Press CTRL+C to quit)CLI applicationRun a prediction in a terminal:$tungstenpredicttext-to-image\-iprompt="a professional photograph of an astronaut riding a horse"{"image": "./output.png"}Python functionIf you want to run a model in your Python application, use the Python API:>>>fromtungstenkitimportmodels>>>model=models.get("text-to-image")>>>model.predict({"prompt":"a professional photograph of an astronaut riding a horse"}){"image":PosixPath("./output.png")}Framework-agnostic and lightweightTungstenkit doesn't restrict you to use specific ML libraries. Just use any library you want, and declare dependencies:# The latest cpu-only build of Tensorflow will be included@define_model(gpu=False,python_packages=["tensorflow"])classTensorflowModel:defpredict(self,inputs):"""Run a batch prediction"""# ...ops using tensorflow...returnoutputsPydantic input/output definitions with convenient file handlingLet's look at the example below:fromtungstenkitimportBaseIO,Image,define_modelclassInput(BaseIO):image:ImageclassOutput(BaseIO):image:Image@define_model(input=Input,output=Output)classStyleTransferModel:...As you see, input/output types are defined as subclasses of theBaseIOclass. TheBaseIOclass is a simple wrapper of theBaseModelclass ofPydantic, and Tungstenkit validates JSON requests utilizing functionalities Pydantic provides.Also, you can see that theImageclass is used. Tungstenkit provides four file classes for easing file handling -Image,Audio,Video, andBinary. They have useful methods for writing a model'spredictmethod:classStyleTransferModel:defpredict(self,inputs:List[Input])->List[Output]:# Preprocessinginput_pil_images=[inp.image.to_pil_image()forinpininputs]# Inferenceoutput_pil_images=do_inference(input_pil_images)# Postprocessingoutput_images=[Image.from_pil_image(pil_image)forpil_imageinoutput_pil_images]outputs=[Output(image=image)forimageinoutput_images]returnoutputsSupports batched predictionTungstenkit supports both server-side and client-side batching.Server-side batchingA server groups inputs across multiple requests and processes them together.
You can configure the max batch size:@define_model(input=Input,output=Output,gpu=True,batch_size=32)The max batch size can be changed when running a server:$tungstenservemymodel-p3000--batch-size16Client-side batchingAlso, you can reduce traffic volume by putting multiple inputs in a single prediction request:$curl-X'POST''http://localhost:3000/predictions'\-H'accept: application/json'\-H'Content-Type: application/json'\-d'[{"field": "input1"}, {"field": "input2"}, {"field": "input3"}]'PrerequisitesPython 3.7+DockerInstallationpipinstalltungstenkitDocumentationGetting StartedBuilding Your ModelRunning ModelsPushing and Pulling ModelsCLI ReferenceREST API ReferenceExamplesJoin our communityIf you have questions about anything related to Tungstenkit, you're always welcome to ask our community onDiscord.