Merge pull request #1 from synapsec-ai/release/1.0.0

Dockerfile

@@ -0,0 +1,26 @@
+FROM python:3.10.14-bookworm
+
+ARG USER_UID=10002
+ARG USER_GID=$USER_UID
+ARG USERNAME=modelapi
+
+RUN groupadd --gid $USER_GID $USERNAME \
+&& useradd --uid $USER_UID --gid $USER_GID -m $USERNAME
+
+# Copy required files
+RUN mkdir -p /modelapi && mkdir -p /home/$USERNAME/.modelapi
+COPY app /modelapi/app
+COPY pyproject.toml /modelapi/pyproject.toml
+
+# Setup permissions
+RUN chown -R $USER_UID:$USER_GID /modelapi \
+&& chown -R $USER_UID:$USER_GID /home/$USERNAME/.modelapi \
+&& chown -R $USER_UID:$USER_GID /home/$USERNAME \
+&& chmod -R 755 /home/$USERNAME \
+&& chmod -R 755 /modelapi \
+&& chmod -R 777 /home/$USERNAME/.modelapi
+
+# Change to the user and do subnet installation
+USER $USERNAME
+
+RUN /bin/bash -c "python3 -m venv /modelapi/.venv && source /modelapi/.venv/bin/activate && pip3 install -e /modelapi/."

README.md

@@ -1,2 +1,36 @@
-# SoundsRightModelContainer
-Model container template for the SoundsRight subnet.
+# Container Template for SoundsRight Subnet Miners
+
+Miners in [Bittensor's](https://bittensor.com/) [SoundsRight Subnet](https://github.com/synapsec-ai/SoundsRightSubnet) must containerize their models before uploading to HuggingFace. This repo serves as a template.
+
+The branches `DENOISING_16000HZ` and `DEREVERBERATION_16000HZ` contain this template fitted with [SGMSE+](https://huggingface.co/sp-uhh/speech-enhancement-sgmse) and are also helpful resources for how to incorporate your model.
+
+The `main` branch contains a template for a container that will spin up an API to communicate with the validator. The following entrypoints cannot be altered: 
+
+1. `/status/` : Communicates API status
+2. `/prepare/` : Makes necessary preparations (downloading checkpoints, etc.) and initializes model
+3. `/upload-audio/` : Upload audio files, save to noisy audio directory
+4. `/enhance/` : Initialize model, enhance audio files, save to enhanced audio directory
+5. `/download-enhanced/` : Download enhanced audio files
+
+To add your own model to this template, there are a few things that a miner must do:
+
+1. Add the model files under the `model` directory.
+2. Modify the `modelapi.prepare` method in `app/app.py` with necessary preparations to initialize your model.
+3. Modify the `modelapi.enhance` method in `app/app.py` with the logic your model uses to enhance audio.
+4. Update `dependencies` in `pyproject.toml` with the dependencies used by your model.
+5. If you have directories other than `app` in your repository, be sure to modify the `Dockerfile` accordingly (reference line 12 in the `Dockerfile` for how to do this).
+6. Cite your sources (if applicable).
+
+For your model to be processed by validators, there are a few formatting requirements. Note that the template already has been formatted to fit these guidelines. 
+
+1. API endpoints must as outlined above.
+2. Port must be 6500. 
+3. There should only be one service in `docker-compose.yml`.
+4. Container must be configured to run as non-root user.
+5. Container names cannot be any of the following:
+
+- common-validator 
+- soundsright-validator-debug-mode
+- soundsright-validator-debug-mode-dev
+- soundsright-validator 
+- soundsright-validator-dev 

app/__init__.py

@@ -0,0 +1 @@
+from .app import ModelAPI

app/app.py

@@ -0,0 +1,158 @@
+import fastapi 
+import shutil 
+import os 
+import zipfile  
+import io 
+import uvicorn 
+import glob
+from typing import List 
+
+class ModelAPI:
+    
+    def __init__(self, host, port):
+        
+        self.host = host 
+        self.port = port
+        
+        self.base_path = os.path.join(os.path.expanduser("~"), ".modelapi")
+        self.noisy_audio_path = os.path.join(self.base_path, "noisy_audio")
+        self.enhanced_audio_path = os.path.join(self.base_path, "enhanced_audio")
+
+        # Create directories if they do not exist
+        for audio_path in [self.noisy_audio_path, self.enhanced_audio_path]:
+            if not os.path.exists(audio_path):
+                os.makedirs(audio_path)
+                
+            # Loop through all the files and subdirectories in the directory
+            for filename in os.listdir(audio_path):
+                file_path = os.path.join(audio_path, filename)
+                
+                # Check if it's a file or directory and remove accordingly
+                try:
+                    if os.path.isfile(file_path) or os.path.islink(file_path):
+                        os.unlink(file_path)  # Remove the file or link
+                    elif os.path.isdir(file_path):
+                        shutil.rmtree(file_path)  # Remove the directory and its contents
+                except Exception as e:
+                    raise e
+                
+        self.app = fastapi.FastAPI()
+        self._setup_routes()
+        
+    def _prepare(self):
+        """Miners should modify this function to fit their fine-tuned models.
+        
+        This function will make any preparations necessary to initialize the
+        speech enhancement model (i.e. downloading checkpoint files, etc.)
+        """
+        # Continue from here
+        pass
+
+    def _enhance(self):
+        """
+        Miners should modify this function to fit their fine-tuned models.
+        
+        This function will:
+        1. Open each noisy .wav file
+        2. Enhance the audio with the model 
+        3. Save the enhanced audio in .wav format to MinerAPI.enhanced_audio_path
+        """
+        
+        # Define file paths for all noisy files to be enhanced
+        noisy_files = sorted(glob.glob(os.path.join(self.noisy_audio_path, '*.wav')))
+        for noisy_file in noisy_files: 
+            # Continue from here
+            pass
+        
+    def _setup_routes(self):
+        """
+        Setup API routes:
+        
+        /status/ : Communicates API status
+        /upload-audio/ : Upload audio files, save to noisy audio directory
+        /enhance/ : Enhance audio files, save to enhanced audio directory
+        /download-enhanced/ : Download enhanced audio files
+        """
+        self.app.get("/status/")(self.get_status)
+        self.app.post("/prepare/")(self.prepare)
+        self.app.post("/upload-audio/")(self.upload_audio)
+        self.app.post("/enhance/")(self.enhance_audio)
+        self.app.get("/download-enhanced/")(self.download_enhanced)
+        
+    def get_status(self):
+        try:
+            return {"container_running": True}
+        except:
+            raise fastapi.HTTPException(status_code=500, detail="An error occurred while fetching API status.")
+        
+    def prepare(self):
+        try:
+            self._prepare()
+            return {'preparations': True}
+        except:
+            return fastapi.HTTPException(status_code=500, detail="An error occurred while fetching API status.")
+        
+    def upload_audio(self, files: List[fastapi.UploadFile] = fastapi.File(...)):
+        
+        uploaded_files = []
+        
+        for file in files:
+            try: 
+                # Define the path to save the file
+                file_path = os.path.join(self.noisy_audio_path, file.filename)
+                
+                # Save the uploaded file
+                with open(file_path, "wb") as f:
+                    while contents := file.file.read(1024*1024):
+                        f.write(contents)
+                
+                # Append the file name to the list of uploaded files
+                uploaded_files.append(file.filename)                
+                
+            except: 
+                raise fastapi.HTTPException(status_code=500, detail="An error occurred while uploading the noisy files.")
+            finally:
+                file.file.close()
+            
+        return {"uploaded_files": uploaded_files, "status": True}
+
+    def enhance_audio(self):
+        try:
+            # Enhance audio
+            self._enhance()
+            # Obtain list of file paths for enhanced audio
+            wav_files = glob.glob(os.path.join(self.enhanced_audio_path, '*.wav'))
+            # Extract just the file names
+            enhanced_files = [os.path.basename(file) for file in wav_files]
+            return {"status": True}
+    
+        except Exception as e:
+            raise fastapi.HTTPException(status_code=500, detail="An error occurred while enhancing the noisy files.")
+        
+    def download_enhanced(self):
+        try:
+            # Create an in-memory zip file to hold all the enhanced audio files
+            zip_buffer = io.BytesIO()
+
+            with zipfile.ZipFile(zip_buffer, "w") as zip_file:
+                # Add each .wav file in the enhanced_audio_path directory to the zip file
+                for wav_file in glob.glob(os.path.join(self.enhanced_audio_path, '*.wav')):
+                    zip_file.write(wav_file, arcname=os.path.basename(wav_file))
+
+            # Make sure to seek back to the start of the BytesIO object before sending it
+            zip_buffer.seek(0)
+
+            # Send the zip file to the client as a downloadable file
+            return fastapi.responses.StreamingResponse(
+                iter([zip_buffer.getvalue()]),  # Stream the in-memory content
+                media_type="application/zip",
+                headers={"Content-Disposition": "attachment; filename=enhanced_audio_files.zip"}
+            )
+
+        except Exception as e:
+            # Log the error if needed, and raise an HTTPException to inform the client
+            raise fastapi.HTTPException(status_code=500, detail=f"An error occurred while creating the download file: {str(e)}")
+
+    def run(self):
+        
+        uvicorn.run(self.app, host=self.host, port=self.port)

app/run.py

@@ -0,0 +1,8 @@
+from app import ModelAPI
+
+api = ModelAPI(
+    host = "0.0.0.0",
+    port = 6500
+)
+
+api.run()

docker-compose.yml

@@ -0,0 +1,12 @@
+services:
+  modelapi:
+    build:
+      context: .
+    container_name: modelapi
+    ports:
+      - "6500:6500"
+    environment:
+      USER_UID: 10002
+      USER_GID: 10002
+      USERNAME: modelapi
+    command: /bin/bash -c "source /modelapi/.venv/bin/activate && python3 /modelapi/app/run.py"

pyproject.toml

@@ -0,0 +1,32 @@
+[build-system]
+requires = ["setuptools", "wheel"]
+build-backend = "setuptools.build_meta"
+
+[project]
+name = "modelapi"
+version = "1.0.0"
+description = "This project implements a container for a fine-tuned audio enhancement model."
+readme = { file = "README.md", content-type = "text/markdown" }
+license = { file = "LICENSE" }
+classifiers = [
+    "Development Status :: 3 - Beta",
+    "Intended Audience :: Developers",
+    "Topic :: Software Development :: Build Tools",
+    "License :: OSI Approved :: MIT License",
+    "Programming Language :: Python :: 3 :: Only",
+    "Programming Language :: Python :: 3.10",
+    "Topic :: Scientific/Engineering",
+    "Topic :: Scientific/Engineering :: Mathematics",
+    "Topic :: Scientific/Engineering :: Artificial Intelligence",
+    "Topic :: Software Development",
+    "Topic :: Software Development :: Libraries",
+    "Topic :: Software Development :: Libraries :: Python Modules"
+]
+requires-python = ">=3.10,<3.12"
+
+dependencies = [
+    "fastapi==0.115.5", "uvicorn==0.32.0",  "python-multipart==0.0.17"
+]
+
+[tool.setuptools.packages.find]
+include = ["app","model"]