Merge pull request #438 from ollama-webui/fix

fix: gguf upload
This commit is contained in:
Timothy Jaeryang Baek 2024-01-09 13:34:04 -08:00 committed by GitHub
commit dcf16df166
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
2 changed files with 32 additions and 29 deletions

View file

@ -12,6 +12,8 @@ import json
from utils.misc import calculate_sha256 from utils.misc import calculate_sha256
from config import OLLAMA_API_BASE_URL from config import OLLAMA_API_BASE_URL
from constants import ERROR_MESSAGES
router = APIRouter() router = APIRouter()
@ -40,10 +42,7 @@ def parse_huggingface_url(hf_url):
return None return None
async def download_file_stream(url, async def download_file_stream(url, file_path, file_name, chunk_size=1024 * 1024):
file_path,
file_name,
chunk_size=1024 * 1024):
done = False done = False
if os.path.exists(file_path): if os.path.exists(file_path):
@ -57,8 +56,7 @@ async def download_file_stream(url,
async with aiohttp.ClientSession(timeout=timeout) as session: async with aiohttp.ClientSession(timeout=timeout) as session:
async with session.get(url, headers=headers) as response: async with session.get(url, headers=headers) as response:
total_size = int(response.headers.get("content-length", total_size = int(response.headers.get("content-length", 0)) + current_size
0)) + current_size
with open(file_path, "ab+") as file: with open(file_path, "ab+") as file:
async for data in response.content.iter_chunked(chunk_size): async for data in response.content.iter_chunked(chunk_size):
@ -91,7 +89,9 @@ async def download_file_stream(url,
@router.get("/download") @router.get("/download")
async def download(url: str, ): async def download(
url: str,
):
# url = "https://huggingface.co/TheBloke/stablelm-zephyr-3b-GGUF/resolve/main/stablelm-zephyr-3b.Q2_K.gguf" # url = "https://huggingface.co/TheBloke/stablelm-zephyr-3b-GGUF/resolve/main/stablelm-zephyr-3b.Q2_K.gguf"
file_name = parse_huggingface_url(url) file_name = parse_huggingface_url(url)
@ -108,25 +108,30 @@ async def download(url: str, ):
@router.post("/upload") @router.post("/upload")
async def upload(file: UploadFile = File(...)): def upload(file: UploadFile = File(...)):
os.makedirs("./uploads", exist_ok=True) os.makedirs("./data/uploads", exist_ok=True)
file_path = os.path.join("./uploads", file.filename) file_path = os.path.join("./data/uploads", file.filename)
async def file_write_stream(): # Save file in chunks
total = 0 with open(file_path, "wb+") as f:
total_size = file.size for chunk in file.file:
f.write(chunk)
def file_process_stream():
total_size = os.path.getsize(file_path)
chunk_size = 1024 * 1024 chunk_size = 1024 * 1024
done = False
try: try:
with open(file_path, "wb+") as f: with open(file_path, "rb") as f:
while True: total = 0
chunk = file.file.read(chunk_size) done = False
while not done:
chunk = f.read(chunk_size)
if not chunk: if not chunk:
break done = True
f.write(chunk) continue
total += len(chunk) total += len(chunk)
done = total_size == total
progress = round((total / total_size) * 100, 2) progress = round((total / total_size) * 100, 2)
res = { res = {
@ -134,7 +139,6 @@ async def upload(file: UploadFile = File(...)):
"total": total_size, "total": total_size,
"completed": total, "completed": total,
} }
yield f"data: {json.dumps(res)}\n\n" yield f"data: {json.dumps(res)}\n\n"
if done: if done:
@ -152,14 +156,14 @@ async def upload(file: UploadFile = File(...)):
"name": file.filename, "name": file.filename,
} }
os.remove(file_path) os.remove(file_path)
yield f"data: {json.dumps(res)}\n\n" yield f"data: {json.dumps(res)}\n\n"
else: else:
raise "Ollama: Could not create blob, Please try again." raise Exception(
"Ollama: Could not create blob, Please try again."
)
except Exception as e: except Exception as e:
res = {"error": str(e)} res = {"error": str(e)}
yield f"data: {json.dumps(res)}\n\n" yield f"data: {json.dumps(res)}\n\n"
return StreamingResponse(file_write_stream(), return StreamingResponse(file_process_stream(), media_type="text/event-stream")
media_type="text/event-stream")

View file

@ -363,7 +363,6 @@
fileResponse = await fetch(`${WEBUI_API_BASE_URL}/utils/upload`, { fileResponse = await fetch(`${WEBUI_API_BASE_URL}/utils/upload`, {
method: 'POST', method: 'POST',
headers: { headers: {
...($settings.authHeader && { Authorization: $settings.authHeader }),
...($user && { Authorization: `Bearer ${localStorage.token}` }) ...($user && { Authorization: `Bearer ${localStorage.token}` })
}, },
body: formData body: formData
@ -375,7 +374,6 @@
fileResponse = await fetch(`${WEBUI_API_BASE_URL}/utils/download?url=${modelFileUrl}`, { fileResponse = await fetch(`${WEBUI_API_BASE_URL}/utils/download?url=${modelFileUrl}`, {
method: 'GET', method: 'GET',
headers: { headers: {
...($settings.authHeader && { Authorization: $settings.authHeader }),
...($user && { Authorization: `Bearer ${localStorage.token}` }) ...($user && { Authorization: `Bearer ${localStorage.token}` })
} }
}).catch((error) => { }).catch((error) => {
@ -1388,9 +1386,10 @@
{#if uploadProgress !== null} {#if uploadProgress !== null}
<div class="mt-2"> <div class="mt-2">
<div class=" mb-2 text-xs">Upload Progress</div> <div class=" mb-2 text-xs">Upload Progress</div>
<div class="w-full rounded-full dark:bg-gray-800"> <div class="w-full rounded-full dark:bg-gray-800">
<div <div
class="dark:bg-gray-600 text-xs font-medium text-blue-100 text-center p-0.5 leading-none rounded-full" class="dark:bg-gray-600 bg-gray-500 text-xs font-medium text-gray-100 text-center p-0.5 leading-none rounded-full"
style="width: {Math.max(15, uploadProgress ?? 0)}%" style="width: {Math.max(15, uploadProgress ?? 0)}%"
> >
{uploadProgress ?? 0}% {uploadProgress ?? 0}%