76 lines
2.0 KiB
Python
76 lines
2.0 KiB
Python
|
import logging
|
||
|
|
import os
|
||
|
|
import shutil
|
||
|
|
import uuid
|
||
|
|
from typing import List, Tuple
|
||
|
|
|
||
|
|
from fastapi import UploadFile
|
||
|
|
|
||
|
|
from app.config import settings
|
||
|
|
|
||
|
|
logger = logging.getLogger(__name__)
|
||
|
|
|
||
|
|
|
||
|
|
def ensure_directories() -> None:
|
||
|
|
for path in (settings.upload_dir, settings.output_dir, settings.models_dir,
|
||
|
|
settings.temp_dir, settings.jobs_dir):
|
||
|
|
os.makedirs(path, exist_ok=True)
|
||
|
|
|
||
|
|
|
||
|
|
def create_request_dir() -> str:
|
||
|
|
request_id = str(uuid.uuid4())
|
||
|
|
request_dir = os.path.join(settings.upload_dir, request_id)
|
||
|
|
os.makedirs(request_dir, exist_ok=True)
|
||
|
|
return request_dir
|
||
|
|
|
||
|
|
|
||
|
|
async def save_upload(file: UploadFile, directory: str) -> str:
|
||
|
|
ext = os.path.splitext(file.filename or '')[1] or ''
|
||
|
|
filename = f'{uuid.uuid4()}{ext}'
|
||
|
|
filepath = os.path.join(directory, filename)
|
||
|
|
|
||
|
|
with open(filepath, 'wb') as f:
|
||
|
|
while chunk := await file.read(1024 * 1024):
|
||
|
|
f.write(chunk)
|
||
|
|
|
||
|
|
return filepath
|
||
|
|
|
||
|
|
|
||
|
|
async def save_uploads(files: List[UploadFile], directory: str) -> List[str]:
|
||
|
|
paths = []
|
||
|
|
for file in files:
|
||
|
|
path = await save_upload(file, directory)
|
||
|
|
paths.append(path)
|
||
|
|
return paths
|
||
|
|
|
||
|
|
|
||
|
|
def generate_output_path(target_path: str) -> str:
|
||
|
|
ext = os.path.splitext(target_path)[1] or '.png'
|
||
|
|
filename = f'{uuid.uuid4()}{ext}'
|
||
|
|
return os.path.join(settings.output_dir, filename)
|
||
|
|
|
||
|
|
|
||
|
|
def cleanup_directory(directory: str) -> None:
|
||
|
|
if os.path.isdir(directory):
|
||
|
|
shutil.rmtree(directory, ignore_errors=True)
|
||
|
|
|
||
|
|
|
||
|
|
def cleanup_file(filepath: str) -> None:
|
||
|
|
if os.path.isfile(filepath):
|
||
|
|
os.remove(filepath)
|
||
|
|
|
||
|
|
|
||
|
|
def list_model_files() -> List[Tuple[str, str, int]]:
|
||
|
|
"""Return list of (name, path, size_bytes) for all .onnx files in models dir."""
|
||
|
|
models = []
|
||
|
|
models_dir = settings.models_dir
|
||
|
|
if not os.path.isdir(models_dir):
|
||
|
|
return models
|
||
|
|
|
||
|
|
for name in sorted(os.listdir(models_dir)):
|
||
|
|
if name.endswith('.onnx'):
|
||
|
|
path = os.path.join(models_dir, name)
|
||
|
|
size = os.path.getsize(path)
|
||
|
|
models.append((name, path, size))
|
||
|
|
return models
|