CelestialWandererOfTheVoid commited on
Commit
657d16d
·
verified ·
1 Parent(s): 6e2b0f5

Upload app files

Browse files
Files changed (2) hide show
  1. app/app.py +159 -0
  2. app/requirements.txt +3 -0
app/app.py ADDED
@@ -0,0 +1,159 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import fastapi
2
+ import shutil
3
+ import os
4
+ import zipfile
5
+ import io
6
+ import uvicorn
7
+ import glob
8
+ from typing import List
9
+
10
+ class ModelAPI:
11
+
12
+ def __init__(self, host, port):
13
+
14
+ self.host = host
15
+ self.port = port
16
+
17
+ self.base_path = os.path.join(os.path.expanduser("~"), ".modelapi")
18
+ self.noisy_audio_path = os.path.join(self.base_path, "noisy_audio")
19
+ self.enhanced_audio_path = os.path.join(self.base_path, "enhanced_audio")
20
+
21
+ # Create directories if they do not exist
22
+ for audio_path in [self.noisy_audio_path, self.enhanced_audio_path]:
23
+ if not os.path.exists(audio_path):
24
+ os.makedirs(audio_path)
25
+
26
+ # Loop through all the files and subdirectories in the directory
27
+ for filename in os.listdir(audio_path):
28
+ file_path = os.path.join(audio_path, filename)
29
+
30
+ # Check if it's a file or directory and remove accordingly
31
+ try:
32
+ if os.path.isfile(file_path) or os.path.islink(file_path):
33
+ os.unlink(file_path) # Remove the file or link
34
+ elif os.path.isdir(file_path):
35
+ shutil.rmtree(file_path) # Remove the directory and its contents
36
+ except Exception as e:
37
+ raise e
38
+
39
+ self.app = fastapi.FastAPI()
40
+ self._setup_routes()
41
+
42
+ def _prepare(self):
43
+ """Miners should modify this function to fit their fine-tuned models.
44
+
45
+ This function will make any preparations necessary to initialize the
46
+ speech enhancement model (i.e. downloading checkpoint files, etc.)
47
+ """
48
+ # Continue from here
49
+ pass
50
+
51
+ def _enhance(self):
52
+ """
53
+ Miners should modify this function to fit their fine-tuned models.
54
+
55
+ This function will:
56
+ 1. Open each noisy .wav file
57
+ 2. Enhance the audio with the model
58
+ 3. Save the enhanced audio in .wav format to MinerAPI.enhanced_audio_path
59
+ """
60
+
61
+ # Define file paths for all noisy files to be enhanced
62
+ noisy_files = sorted(glob.glob(os.path.join(self.noisy_audio_path, '*.wav')))
63
+ for noisy_file in noisy_files:
64
+ # Continue from here
65
+ pass
66
+
67
+ def _setup_routes(self):
68
+ """
69
+ Setup API routes:
70
+
71
+ /status/ : Communicates API status
72
+ /prepare/ : Makes necessary preparations (downloading checkpoints, etc.) and initializes model
73
+ /upload-audio/ : Upload audio files, save to noisy audio directory
74
+ /enhance/ : Initialize model, enhance audio files, save to enhanced audio directory
75
+ /download-enhanced/ : Download enhanced audio files
76
+ """
77
+ self.app.get("/status/")(self.get_status)
78
+ self.app.post("/prepare/")(self.prepare)
79
+ self.app.post("/upload-audio/")(self.upload_audio)
80
+ self.app.post("/enhance/")(self.enhance_audio)
81
+ self.app.get("/download-enhanced/")(self.download_enhanced)
82
+
83
+ def get_status(self):
84
+ try:
85
+ return {"container_running": True}
86
+ except:
87
+ raise fastapi.HTTPException(status_code=500, detail="An error occurred while fetching API status.")
88
+
89
+ def prepare(self):
90
+ try:
91
+ self._prepare()
92
+ return {'preparations': True}
93
+ except:
94
+ return fastapi.HTTPException(status_code=500, detail="An error occurred while fetching API status.")
95
+
96
+ def upload_audio(self, files: List[fastapi.UploadFile] = fastapi.File(...)):
97
+
98
+ uploaded_files = []
99
+
100
+ for file in files:
101
+ try:
102
+ # Define the path to save the file
103
+ file_path = os.path.join(self.noisy_audio_path, file.filename)
104
+
105
+ # Save the uploaded file
106
+ with open(file_path, "wb") as f:
107
+ while contents := file.file.read(1024*1024):
108
+ f.write(contents)
109
+
110
+ # Append the file name to the list of uploaded files
111
+ uploaded_files.append(file.filename)
112
+
113
+ except:
114
+ raise fastapi.HTTPException(status_code=500, detail="An error occurred while uploading the noisy files.")
115
+ finally:
116
+ file.file.close()
117
+
118
+ return {"uploaded_files": uploaded_files, "status": True}
119
+
120
+ def enhance_audio(self):
121
+ try:
122
+ # Enhance audio
123
+ self._enhance()
124
+ # Obtain list of file paths for enhanced audio
125
+ wav_files = glob.glob(os.path.join(self.enhanced_audio_path, '*.wav'))
126
+ # Extract just the file names
127
+ enhanced_files = [os.path.basename(file) for file in wav_files]
128
+ return {"status": True}
129
+
130
+ except Exception as e:
131
+ raise fastapi.HTTPException(status_code=500, detail="An error occurred while enhancing the noisy files.")
132
+
133
+ def download_enhanced(self):
134
+ try:
135
+ # Create an in-memory zip file to hold all the enhanced audio files
136
+ zip_buffer = io.BytesIO()
137
+
138
+ with zipfile.ZipFile(zip_buffer, "w") as zip_file:
139
+ # Add each .wav file in the enhanced_audio_path directory to the zip file
140
+ for wav_file in glob.glob(os.path.join(self.enhanced_audio_path, '*.wav')):
141
+ zip_file.write(wav_file, arcname=os.path.basename(wav_file))
142
+
143
+ # Make sure to seek back to the start of the BytesIO object before sending it
144
+ zip_buffer.seek(0)
145
+
146
+ # Send the zip file to the client as a downloadable file
147
+ return fastapi.responses.StreamingResponse(
148
+ iter([zip_buffer.getvalue()]), # Stream the in-memory content
149
+ media_type="application/zip",
150
+ headers={"Content-Disposition": "attachment; filename=enhanced_audio_files.zip"}
151
+ )
152
+
153
+ except Exception as e:
154
+ # Log the error if needed, and raise an HTTPException to inform the client
155
+ raise fastapi.HTTPException(status_code=500, detail=f"An error occurred while creating the download file: {str(e)}")
156
+
157
+ def run(self):
158
+
159
+ uvicorn.run(self.app, host=self.host, port=self.port)
app/requirements.txt ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ fastapi==0.95.1
2
+ uvicorn==0.22.0
3
+ python-multipart==0.0.6