Vixy
e4294b57e6
- SDXL image generation using RealVisXL_V4.0 - IP-Adapter FaceID integration for consistent face generation - Simplified API (removed client_id requirement) - New params: face_image, face_strength - 'vixy' shortcut for face-locked generation - Queue-based async job processing - FastAPI with proper error handling Co-authored-by: Alex <alex@k4zka.online>
133 lines
3.6 KiB
Python
Executable File
133 lines
3.6 KiB
Python
Executable File
"""
|
|
File Storage Manager
|
|
|
|
Handles saving and retrieving generated images.
|
|
"""
|
|
|
|
import logging
|
|
from pathlib import Path
|
|
from typing import Optional
|
|
from PIL import Image
|
|
import aiofiles
|
|
import os
|
|
|
|
import config
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
class FileManager:
|
|
"""Manages image file storage and retrieval."""
|
|
|
|
def __init__(self):
|
|
self.images_dir = config.IMAGES_DIR
|
|
self.image_format = config.IMAGE_FORMAT
|
|
|
|
# Ensure storage directory exists
|
|
self.images_dir.mkdir(parents=True, exist_ok=True)
|
|
logger.info(f"Image storage directory: {self.images_dir}")
|
|
|
|
async def save_image(self, job_id: str, image: Image.Image) -> str:
|
|
"""
|
|
Save generated image to disk.
|
|
|
|
Args:
|
|
job_id: Job identifier (used as filename)
|
|
image: PIL Image to save
|
|
|
|
Returns:
|
|
file_path: Absolute path to saved image
|
|
|
|
Raises:
|
|
IOError: If save fails
|
|
"""
|
|
filename = f"{job_id}.{self.image_format.lower()}"
|
|
file_path = self.images_dir / filename
|
|
|
|
try:
|
|
# Save in thread pool to avoid blocking
|
|
import asyncio
|
|
loop = asyncio.get_event_loop()
|
|
await loop.run_in_executor(
|
|
None,
|
|
lambda: image.save(
|
|
file_path,
|
|
format=self.image_format,
|
|
quality=config.IMAGE_QUALITY if self.image_format == "JPEG" else None
|
|
)
|
|
)
|
|
|
|
logger.info(f"Image saved: {file_path} ({os.path.getsize(file_path) / 1024:.1f} KB)")
|
|
return str(file_path)
|
|
|
|
except Exception as e:
|
|
logger.error(f"Failed to save image {job_id}: {e}")
|
|
raise IOError(f"Failed to save image: {e}")
|
|
|
|
def get_image_path(self, job_id: str) -> Optional[Path]:
|
|
"""
|
|
Get path to image file if it exists.
|
|
|
|
Args:
|
|
job_id: Job identifier
|
|
|
|
Returns:
|
|
Path to image file or None if not found
|
|
"""
|
|
filename = f"{job_id}.{self.image_format.lower()}"
|
|
file_path = self.images_dir / filename
|
|
|
|
if file_path.exists():
|
|
return file_path
|
|
return None
|
|
|
|
def image_exists(self, job_id: str) -> bool:
|
|
"""Check if image file exists."""
|
|
return self.get_image_path(job_id) is not None
|
|
|
|
async def delete_image(self, job_id: str) -> bool:
|
|
"""
|
|
Delete an image file.
|
|
|
|
Args:
|
|
job_id: Job identifier
|
|
|
|
Returns:
|
|
True if deleted, False if not found
|
|
"""
|
|
file_path = self.get_image_path(job_id)
|
|
|
|
if file_path:
|
|
try:
|
|
file_path.unlink()
|
|
logger.info(f"Deleted image: {file_path}")
|
|
return True
|
|
except Exception as e:
|
|
logger.error(f"Failed to delete image {job_id}: {e}")
|
|
return False
|
|
|
|
return False
|
|
|
|
def get_storage_stats(self) -> dict:
|
|
"""Get storage statistics."""
|
|
try:
|
|
files = list(self.images_dir.glob(f"*.{self.image_format.lower()}"))
|
|
total_size = sum(f.stat().st_size for f in files)
|
|
|
|
return {
|
|
"total_images": len(files),
|
|
"total_size_mb": total_size / (1024 * 1024),
|
|
"storage_path": str(self.images_dir)
|
|
}
|
|
except Exception as e:
|
|
logger.error(f"Failed to get storage stats: {e}")
|
|
return {
|
|
"total_images": 0,
|
|
"total_size_mb": 0,
|
|
"storage_path": str(self.images_dir)
|
|
}
|
|
|
|
|
|
# Global file manager instance
|
|
file_manager = FileManager()
|