feat: 上传文件

api.py

@@ -8,6 +8,7 @@ from io import BytesIO
 from datetime import datetime
 import hashlib
 import time
+from huggingface_hub import HfApi
 
 
 async def verify_internal_token(request: Request):
@@ -45,18 +46,41 @@ def download_file(url: str):
 def read_root():
     return {"Hello": "World!"}
 
-@router.post("/upload")
-async def upload_file(file: UploadFile = File(...)):
-    # 定义文件存储路径和名称
-    file_location = os.path.join(UPLOAD_DIRECTORY, file.filename)
-
-    # 将上传的文件写入本地文件系统
-    with open(file_location, "wb") as buffer:
-        buffer.write(await file.read())
-
-    return {"info": f"file '{file.filename}' saved at '{file_location}'"}
+@router.post("/upload/{modelId}/{userId}/{filename}")
+async def upload_file(modelId: str, userId: str, filename: str, file: UploadFile = File(...)):
+    file_location = os.path.join(UPLOAD_DIRECTORY,userId, filename)
+    try:
+        with open(file_location, "wb") as buffer:
+            buffer.write(await file.read())
+        callback(modelId,'UPLOADING')
+        pathInRepo = f"{userId}/{filename}"
+        huggingfaceApi = HfApi()
+        HfFolder.save_token(os.environ.get("HF_TOKEN"))
+        huggingfaceApi.upload_file(
+            path_or_fileobj=file_location,
+            path_in_repo=pathInRepo,
+            repo_id='wsj1995/aigc-user-uploaded-models',
+            repo_type="model"
+        )
+        callback(modelId,'UPLOADED')
+    except Exception as e:
+        print(e)
+        callback(modelId,'FAIL')
+    os.remove(file_location)
+    return {'success': True, 'id': modelId}
 
 
 @router.get("/download")
 def download():
-    return StreamingResponse(download_file("https://huggingface.co/wsj1995/stable-diffusion-models/resolve/main/3Guofeng3_v34.safetensors?download=true"), media_type="application/octet-stream")
+    return StreamingResponse(download_file("https://huggingface.co/wsj1995/stable-diffusion-models/resolve/main/3Guofeng3_v34.safetensors?download=true"), media_type="application/octet-stream")
+
+
+
+def callback(modelId,status):
+    res = requests.post(os.environ.get("CALLBACK_URL"),json={
+        'status': status,
+        'model_id': modelId
+    },headers={
+        os.environ.get("CALLBACK_SECRET_KEY"): os.environ.get("CALLBACK_SECRET")
+    })
+    print(f"回调结果 {res.status_code}")

requirements.txt

@@ -4,4 +4,5 @@ sentencepiece==0.1.*
 torch==1.11.*
 transformers==4.*
 uvicorn[standard]==0.17.*
-python-multipart
+python-multipart
+huggingface_hub