updates

Dockerfile

@@ -10,7 +10,6 @@ RUN groupadd --gid $USER_GID $USERNAME \
 # Copy required files
 RUN mkdir -p /modelapi && mkdir -p /home/$USERNAME/.modelapi
 COPY app /modelapi/app
-COPY model /modelapi/model
 COPY pyproject.toml /modelapi/pyproject.toml
 
 # Setup permissions

README.md

@@ -4,10 +4,10 @@ Miners in [Bittensor's](https://bittensor.com/) [SoundsRight Subnet](https://git
 
 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.
 
-This repo contains a template for a container that will spin up an API to communicate with the validator. The following entrypoints cannot be altered: 
+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 to initialize model (downloading checkpoints, etc.)
+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
@@ -18,7 +18,8 @@ To add your own model to this template, there are a few things that a miner must
 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. Cite your sources (if applicable).
+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. 
 

app/app.py

@@ -35,25 +35,21 @@ class ModelAPI:
                     elif os.path.isdir(file_path):
                         shutil.rmtree(file_path)  # Remove the directory and its contents
                 except Exception as e:
-                    Utils.subnet_logger(
-                        severity="ERROR",
-                        message=f"Failed to delete {file_path}. Reason: {e}",
-                        log_level=self.log_level
-                    )
+                    raise e
                 
         self.app = fastapi.FastAPI()
         self._setup_routes()
         
-    def prepare(self):
+    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 
+        pass
 
-    def enhance(self):
+    def _enhance(self):
         """
         Miners should modify this function to fit their fine-tuned models.
         
@@ -61,16 +57,11 @@ class ModelAPI:
         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
-
-        Args:
-            noisy_dir (str): The path to the directory containing the audio files to be enhanced.
-            enhanced_dir (str): The path to the directory where the enhanced audio files will be saved (the filenames must be identical to those in the noisy_dir).
         """
         
         # 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
         
@@ -97,42 +88,46 @@ class ModelAPI:
         
     def prepare(self):
         try:
-            self.prepare()
+            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(...)):
-        try:
-            uploaded_files = []
-            
-            for file in files:
+        
+        uploaded_files = []
+        
+        for file in files:
+            try: 
                 # Define the path to save the file
-                file_path = os.path.join(noisy_audio_path, file.filename)
+                file_path = os.path.join(self.noisy_audio_path, file.filename)
                 
                 # Save the uploaded file
-                with open(file_path, "wb") as buffer:
-                    shutil.copyfileobj(file.file, buffer)
+                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)
+                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}
-        
-        except: 
-            raise fastapi.HTTPException(status_code=500, detail="An error occurred while uploading the noisy files.")
+        return {"uploaded_files": uploaded_files, "status": True}
 
     def enhance_audio(self):
         try:
             # Enhance audio
-            self.enhance()
+            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:
+        except Exception as e:
             raise fastapi.HTTPException(status_code=500, detail="An error occurred while enhancing the noisy files.")
         
     def download_enhanced(self):
@@ -149,15 +144,15 @@ class ModelAPI:
             zip_buffer.seek(0)
 
             # Send the zip file to the client as a downloadable file
-            return fastapi.responses.FileResponse(
-                zip_buffer,
+            return fastapi.responses.StreamingResponse(
+                iter([zip_buffer.getvalue()]),  # Stream the in-memory content
                 media_type="application/zip",
-                filename="enhanced_audio_files.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="An error occurred while creating the download file.")
+            raise fastapi.HTTPException(status_code=500, detail=f"An error occurred while creating the download file: {str(e)}")
 
     def run(self):