feat: Implement image thumbnail generation, storage, and API endpoints for assets, including a regeneration utility.
This commit is contained in:
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
@@ -21,17 +21,31 @@ router = APIRouter(prefix="/api/assets", tags=["Assets"])
|
||||
|
||||
|
||||
@router.get("/{asset_id}")
|
||||
async def get_asset(asset_id: str, request: Request,dao: DAO = Depends(get_dao),) -> Response:
|
||||
logger.debug(f"get_asset called for ID: {asset_id}")
|
||||
async def get_asset(
|
||||
asset_id: str,
|
||||
request: Request,
|
||||
thumbnail: bool = False,
|
||||
dao: DAO = Depends(get_dao)
|
||||
) -> Response:
|
||||
logger.debug(f"get_asset called for ID: {asset_id}, thumbnail={thumbnail}")
|
||||
asset = await dao.assets.get_asset(asset_id)
|
||||
# 2. Проверка на существование
|
||||
if not asset:
|
||||
raise HTTPException(status_code=404, detail="Asset not found")
|
||||
|
||||
headers = {
|
||||
# Кэшировать на 1 год (31536000 сек)
|
||||
"Cache-Control": "public, max-age=31536000, immutable"
|
||||
}
|
||||
return Response(content=asset.data, media_type="image/png", headers=headers)
|
||||
|
||||
content = asset.data
|
||||
media_type = "image/png" # Default, or detect
|
||||
|
||||
if thumbnail and asset.thumbnail:
|
||||
content = asset.thumbnail
|
||||
media_type = "image/jpeg"
|
||||
|
||||
return Response(content=content, media_type=media_type, headers=headers)
|
||||
|
||||
|
||||
@router.get("")
|
||||
@@ -41,7 +55,13 @@ async def get_assets(request: Request, dao: DAO = Depends(get_dao), limit: int =
|
||||
# assets = await dao.assets.get_assets() # This line seemed redundant/conflicting in original code
|
||||
total_count = await dao.assets.get_asset_count()
|
||||
|
||||
return AssetsResponse(assets=assets, total_count=total_count)
|
||||
# Manually map to DTO to trigger computed fields validation if necessary,
|
||||
# but primarily to ensure valid Pydantic models for the response list.
|
||||
# Asset.model_dump() generally includes computed fields (url) if configured.
|
||||
# Let's ensure strict conversion.
|
||||
asset_responses = [AssetResponse.model_validate(a.model_dump()) for a in assets]
|
||||
|
||||
return AssetsResponse(assets=asset_responses, total_count=total_count)
|
||||
|
||||
|
||||
|
||||
@@ -62,11 +82,16 @@ async def upload_asset(
|
||||
if not data:
|
||||
raise HTTPException(status_code=400, detail="Empty file")
|
||||
|
||||
# Generate thumbnail
|
||||
from utils.image_utils import create_thumbnail
|
||||
thumbnail_bytes = await asyncio.to_thread(create_thumbnail, data)
|
||||
|
||||
asset = Asset(
|
||||
name=file.filename or "upload",
|
||||
type=AssetType.IMAGE,
|
||||
linked_char_id=linked_char_id,
|
||||
data=data,
|
||||
thumbnail=thumbnail_bytes
|
||||
)
|
||||
|
||||
asset_id = await dao.assets.create_asset(asset)
|
||||
@@ -79,4 +104,39 @@ async def upload_asset(
|
||||
type=asset.type.value if hasattr(asset.type, "value") else asset.type,
|
||||
linked_char_id=asset.linked_char_id,
|
||||
created_at=asset.created_at,
|
||||
)
|
||||
url=asset.url
|
||||
)
|
||||
|
||||
|
||||
@router.post("/regenerate_thumbnails")
|
||||
async def regenerate_thumbnails(dao: DAO = Depends(get_dao)):
|
||||
"""
|
||||
Regenerates thumbnails for all existing image assets that don't have one.
|
||||
"""
|
||||
logger.info("Starting thumbnail regeneration task")
|
||||
from utils.image_utils import create_thumbnail
|
||||
import asyncio
|
||||
|
||||
# Get all assets (pagination loop might be needed for huge datasets, but simple list for now)
|
||||
# We'll rely on DAO providing a method or just fetch large chunk.
|
||||
# Assuming get_assets might have limit, let's fetch in chunks or just all if possible within limit.
|
||||
# Ideally should use a specific repo method for iteration.
|
||||
# For now, let's fetch first 1000 or similar.
|
||||
assets = await dao.assets.get_assets(limit=1000, offset=0, with_data=True)
|
||||
logger.info(f"Found {len(assets)} assets")
|
||||
count = 0
|
||||
updated = 0
|
||||
|
||||
for asset in assets:
|
||||
if asset.type == AssetType.IMAGE and asset.data :
|
||||
try:
|
||||
thumb = await asyncio.to_thread(create_thumbnail, asset.data)
|
||||
if thumb:
|
||||
asset.thumbnail = thumb
|
||||
await dao.assets.update_asset(asset.id, asset)
|
||||
updated += 1
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to regenerate thumbnail for asset {asset.id}: {e}")
|
||||
count += 1
|
||||
|
||||
return {"status": "completed", "processed": count, "updated": updated}
|
||||
@@ -5,7 +5,7 @@ from pydantic import BaseModel
|
||||
from starlette.exceptions import HTTPException
|
||||
from starlette.requests import Request
|
||||
|
||||
from api.models.AssetDTO import AssetsResponse
|
||||
from api.models.AssetDTO import AssetsResponse, AssetResponse
|
||||
from api.models.GenerationRequest import GenerationRequest, GenerationResponse
|
||||
from models.Asset import Asset
|
||||
from models.Character import Character
|
||||
@@ -35,7 +35,9 @@ async def get_character_assets(character_id: str, dao: DAO = Depends(get_dao), l
|
||||
raise HTTPException(status_code=404, detail="Character not found")
|
||||
assets = await dao.assets.get_assets_by_char_id(character_id, limit, offset)
|
||||
total_count = await dao.assets.get_asset_count(character_id)
|
||||
return AssetsResponse(assets=assets, total_count=total_count)
|
||||
|
||||
asset_responses = [AssetResponse.model_validate(a.model_dump()) for a in assets]
|
||||
return AssetsResponse(assets=asset_responses, total_count=total_count)
|
||||
|
||||
|
||||
@router.get("/{character_id}", response_model=Character)
|
||||
|
||||
@@ -6,14 +6,14 @@ from pydantic import BaseModel
|
||||
from models.Asset import Asset
|
||||
|
||||
|
||||
class AssetsResponse(BaseModel):
|
||||
assets: List[Asset]
|
||||
total_count: int
|
||||
|
||||
|
||||
class AssetResponse(BaseModel):
|
||||
id: str
|
||||
name: str
|
||||
type: str
|
||||
linked_char_id: Optional[str] = None
|
||||
created_at: datetime
|
||||
created_at: datetime
|
||||
url: Optional[str] = None
|
||||
|
||||
class AssetsResponse(BaseModel):
|
||||
assets: List[AssetResponse]
|
||||
total_count: int
|
||||
Binary file not shown.
Binary file not shown.
Binary file not shown.
@@ -66,7 +66,7 @@ class GenerationService:
|
||||
async def ask_prompt_assistant(self, prompt: str, assets: List[str] = None) -> str:
|
||||
future_prompt = """You are an prompt-assistant. You improving user-entered prompts for image generation. User may upload reference image too.
|
||||
I will provide sources prompt entered by user. Understand user needs and generate best variation of prompt.
|
||||
ANSWER ONLY PROMPT STRING!!! USER_ENTERED_PROMPT:"""
|
||||
ANSWER ONLY PROMPT STRING!!! USER_ENTERED_PROMPT: """
|
||||
future_prompt += prompt
|
||||
assets_data = []
|
||||
if assets is not None:
|
||||
@@ -88,7 +88,7 @@ class GenerationService:
|
||||
|
||||
async def get_generations(self, character_id: Optional[str] = None, limit: int = 10, offset: int = 0) -> List[
|
||||
Generation]:
|
||||
return await self.dao.generations.get_generations(limit=limit, offset=offset)
|
||||
return await self.dao.generations.get_generations(character_id = character_id,limit=limit, offset=offset)
|
||||
|
||||
async def get_generation(self, generation_id: str) -> Optional[GenerationResponse]:
|
||||
gen = await self.dao.generations.get_generation(generation_id)
|
||||
@@ -162,7 +162,7 @@ class GenerationService:
|
||||
for asset in reference_assets
|
||||
if asset.data is not None and asset.type == AssetType.IMAGE
|
||||
)
|
||||
generation_prompt+=f"PROMPT: {generation.prompt}"
|
||||
generation_prompt+=f" PROMPT: {generation.prompt}"
|
||||
logger.info(f"Final generation prompt assembled. Length: {len(generation_prompt)}. Media count: {len(media_group_bytes)}")
|
||||
|
||||
# 3. Запускаем процесс генерации и симуляцию прогресса
|
||||
@@ -209,11 +209,19 @@ class GenerationService:
|
||||
created_assets: List[Asset] = []
|
||||
|
||||
for idx, img_bytes in enumerate(generated_bytes_list):
|
||||
# Generate thumbnail
|
||||
thumbnail_bytes = None
|
||||
# Assuming AssetType.IMAGE since we are in generated_bytes_list which are images usually
|
||||
# Or use explicit check if we have distinct types in list (not currently)
|
||||
from utils.image_utils import create_thumbnail
|
||||
thumbnail_bytes = await asyncio.to_thread(create_thumbnail, img_bytes)
|
||||
|
||||
new_asset = Asset(
|
||||
name=f"Generated_{generation.linked_character_id}_{random.randint(1000, 9999)}",
|
||||
type=AssetType.IMAGE,
|
||||
linked_char_id=generation.linked_character_id, # Если генерация привязана к персонажу
|
||||
data=img_bytes,
|
||||
thumbnail=thumbnail_bytes
|
||||
# Остальные поля заполнятся дефолтными значениями (created_at)
|
||||
)
|
||||
|
||||
@@ -236,6 +244,8 @@ class GenerationService:
|
||||
end_time = datetime.now()
|
||||
generation.execution_time_seconds = (end_time - start_time).total_seconds()
|
||||
|
||||
logger.info(f"DEBUG: Saving generation {generation.id}. Metrics: api_exec={generation.api_execution_time_seconds}, tokens={generation.token_usage}, exec={generation.execution_time_seconds}")
|
||||
|
||||
await self.dao.generations.update_generation(generation)
|
||||
logger.info(f"Generation {generation.id} completed successfully. {len(created_assets)} assets created. Total Time: {generation.execution_time_seconds:.2f}s")
|
||||
|
||||
|
||||
Reference in New Issue
Block a user