Update main.py

main.py

@@ -1,3 +1,4 @@
+import os
 from fastapi import FastAPI, File, UploadFile, Form
 from fastapi.responses import StreamingResponse
 import torch